Compare commits

...
62 Commits
Author SHA1 Message Date
geohot aec4d65241 ds compiled 2025-12-31 15:43:44 -05:00
geohot f022a7d8a7 assembly/amd: move more instructions to pcode 2025-12-31 15:42:59 -05:00
George HotzandGitHub 2bb07d4824 assembly/amd: move Reg out of the psuedocode (#13934)
* assembly/amd: move Reg out of the psuedocode

* remove extra

* fix pcode tests

* simpler pcode

* simpler

* simpler

* cleaner

* fix mypy
2025-12-31 15:34:51 -05:00
chenyuandGitHub 52acadc160 consolidate IGNORE_OOB=0 tests (#13937)
add a new unit test file and add more cases
2025-12-31 15:24:20 -05:00
chenyuandGitHub c0c1c1c8c8 remove unused validate rule (#13936) 2025-12-31 15:02:49 -05:00
chenyuandGitHub b6d08f247d assert z3_xor input type (#13933) 2025-12-31 13:37:57 -05:00
George HotzandGitHub f14428090f assembly/amd: speed up emulator (#13932) 2025-12-31 13:32:25 -05:00
sirhcmandGitHub 13973e4dea refactor image pitch (#13928) 2025-12-31 13:22:38 -05:00
chenyuandGitHub 051fe6c8bc less toposort iteration in oob validate (#13929) 2025-12-31 13:16:34 -05:00
chenyuandGitHub a9a7b33404 IGNORE_OOB=0 in CI (#13903) 2025-12-31 12:56:59 -05:00
George HotzandGitHub 29402034a1 assembly/amd: cleanups to asm and emu (#13912)
* a bunch of cleanups

* ops are back

* bug fixes

* cleanups

* a lil simpler

* more refactors

* _disasm_vop1

* sops

* more

* continue

* more

* num_srcs

* simpler

* no _is16

* op cleanups

* isinstnace
2025-12-31 12:46:11 -05:00
chenyuandGitHub ba9aa5cd6f skip some PTX IGNORE_OOB validation (#13927) 2025-12-31 12:40:21 -05:00
chenyuandGitHub 4968060ad4 fix IGNORE_OOB=0 for WEBGPU (#13926) 2025-12-31 10:41:28 -05:00
chenyuandGitHub 35bd39e4ba update mypy and torch version in ci (#13925) 2025-12-31 10:29:28 -05:00
George HotzandGitHub b998a80b5d assembly/amd: split generated stuff into enum/ins (#13924) 2025-12-31 10:10:52 -05:00
chenyuandGitHub 404755bafd merge ci ruff tests and update ruff version (#13922) 2025-12-31 09:53:49 -05:00
nimlgenandGitHub 25440f0f72 all2all (#13902)
* all2all

* um

* fix

* x

* um

* simler

* mypy

* fix

* t

* cmnts
2025-12-31 16:38:32 +03:00
nimlgenandGitHub f7ee644950 amd: lazy sdma queue allocation (#13920)
* ams: lazy queue

* nv

* linter

* f
2025-12-31 15:17:13 +03:00
nimlgenandGitHub b063518ea7 am: several sdmas (#13919)
* am: several sdmas

* fix
2025-12-31 14:19:22 +03:00
qazalandGitHub b23f4517ab prep mi350x gemm for python dsl (#13918)
* start by pruning existing asm

* better branch names

* split to template and real instructions
2025-12-31 20:00:57 +09:00
qazalandGitHub 3f3786ded9 mmapeak: fix compiler import (#13915) 2025-12-31 16:52:23 +09:00
sirhcmandGitHub a14896fff2 refactor QCOM arg parsing (#13914)
* refactor QCOM arg parsing

* ruff

* mypy
2025-12-30 19:26:02 -05:00
sirhcmandGitHub c475c3a6d7 remove useless cast (#13911) 2025-12-30 19:24:29 -05:00
George HotzandGitHub 0221b96761 assembly/amd: fix all ops tests (#13910)
* assembly/amd: fix all ops tests

* test_ops with smaller sizes

* ds store/load 2addr
2025-12-30 18:01:34 -05:00
chenyuandGitHub dc27eb48ac remove PYTHONPATH="." from test.yml (#13909) 2025-12-30 17:00:16 -05:00
George HotzandGitHub efc99d0c55 assembly/amd: more refactors (#13907)
* assembly/amd: more refactors

* more refactors

* more refactors

* simpler emu

* generate.py

* regen all

* cleanups

* more

* work

* more readme

* lil
2025-12-30 16:13:24 -05:00
George HotzandGitHub 49d1bf93d6 assembly/amd: refactor asm.py to be simpler (#13900)
* assembly/amd: refactor asm.py

* assembly/amd: refactor asm.py to be simpler

* multiple fxns

* fast

* more tests pass

* regen

* stop decode
2025-12-30 13:51:40 -05:00
George HotzandGitHub 04c79505ec no subnormal bf16 (#13905) 2025-12-30 13:02:53 -05:00
chenyuandGitHub 39f99b207a update IGNORE_OOB error message (#13904)
IGNORE_OOB=1 to disable
2025-12-30 12:25:55 -05:00
George HotzandGitHub 7e14cdcb06 assembly/amd: clean up clt/ctz hack (#13901)
* assembly/amd: clean up clt/ctz hack

* add breaks
2025-12-30 11:59:28 -05:00
George HotzandGitHub 69cdc8066d assembly/amd: add dtype tests to AMD IDE CI (#13899)
* add dtype tests to AMD IDE CI

* more tests

* add trig preop

* regen done

* split to amd autogen

* simpler
2025-12-30 11:09:51 -05:00
George HotzandGitHub 9c89be5235 assembly/amd: fix v_perm_b32 + PC fixes (#13897)
* assembly/amd: fix v_perm_b32

* add pc support
2025-12-30 09:25:40 -05:00
George HotzandGitHub 2b838dc1d8 assembly/amd: fix AMD_LLVM=1 support in emulator (#13881)
* fix AMD_LLVM=1 support in emulator

* more llvm with dtype

* work

* more fixes

* fix dtype
2025-12-30 09:09:57 -05:00
nimlgenandGitHub a19d21ea9c am: mi3xx smu clocks (#13894)
* am: mi3xx smu clocks

* x
2025-12-30 16:44:17 +03:00
qazalandGitHub b557c46233 assembly gemm clean ups, instructions for cli (#13892) 2025-12-30 16:14:06 +09:00
qazalandGitHub d7e1f26e3d command line interface for sqtt viz (#13891)
* command line interface for sqtt viz

* cleanup

* api surface area

* this confuses the llms

* document
2025-12-30 12:33:21 +09:00
chenyuandGitHub ab58926b00 update sampling in test_float_cast_to_unsigned (#13889)
filter is slow for small dtypes
2025-12-29 21:35:46 -05:00
sirhcmandGitHub 0497387e45 NIR: new-style (fix beam) (#13887)
* NIR: fix beam

* new reduce

* Revert "Revert "NIR: new-style compilers (#13875)" (#13888)"

This reverts commit fc4faed0b2.

* oops
2025-12-29 18:41:29 -05:00
sirhcmandGitHub fc4faed0b2 Revert "NIR: new-style compilers (#13875)" (#13888)
This reverts commit 72236bbd3d.
2025-12-29 17:42:28 -05:00
George HotzandGitHub 94bca91f3e assembly/amd: have asm go through the dsl (#13886)
* assembly/amd: have asm go through the dsl

* lil
2025-12-29 17:39:11 -05:00
George HotzandGitHub 7322d9ec4a assembly/amd: add new instruction support to pcode (#13885)
* assembly/amd: add new instruction support

* more

* regen all
2025-12-29 17:30:17 -05:00
George HotzandGitHub 0d326f5b9b fix missing instructions in psuedocode (#13884) 2025-12-29 16:11:22 -05:00
sirhcmandGitHub 9c6850fc01 remove try-catches on llvm import (#13883) 2025-12-29 15:56:17 -05:00
George HotzandGitHub 9d8397be11 add CDNA3+RDNA4 support (#13882)
* fix CI

* remove junk

* rename lib to dsl

* correct

* cleanups
2025-12-29 15:51:29 -05:00
sirhcmandGitHub 72236bbd3d NIR: new-style compilers (#13875)
* NIR: new-style compilers

* mypy

* simplify NIR compilers

* lvp compiler too

* mypy

* simplify

* mypy
2025-12-29 15:31:41 -05:00
George HotzandGitHub 81cf9ea0ab rename to extra.assembly.amd (#13879) 2025-12-29 14:10:55 -05:00
George HotzandGitHub 37f0fa11b6 rdna3 test cleanups (#13878)
* rdna3 test cleanups

* cleanups

* ugh DONT SKIP
2025-12-29 13:41:59 -05:00
George HotzandGitHub 35db73b231 add cdna4 support to parsers (#13877)
* add cdna4 support to parsers

* cdna4
2025-12-29 13:23:43 -05:00
Clément VerrierandGitHub d178235309 delete tree structure from CLAUDE.md (#13876)
Claude Code should be able to figure out the correct structure, and the
hardcoded tree structure might become outdated.
2025-12-29 13:23:20 -05:00
George HotzandGitHub ff856a74cb minor refactoring for rdna3 (#13873)
* minor refactoring for rdna3

* fix div scale stuff

* more bugfixes
2025-12-29 13:20:00 -05:00
C TandGitHub 39923203ba fix exception in cuda bindings code on windows (#13823)
* fix cuda on windows

* fix linter errors

* test github action install cuda-toolkit

* Revert "test github action install cuda-toolkit"

This reverts commit c18ad6f937.

* Revert "fix linter errors"

This reverts commit 00aa943e91.

* Revert "fix cuda on windows"

This reverts commit 7aea5256b1.

* fix windows sysconfig.get_config_var("MULTIARCH") is None
2025-12-29 12:58:22 -05:00
b1tgandGitHub 63a1bb8507 multi custom kernel: support input mixed with copy and shard (#13748) 2025-12-29 12:54:27 -05:00
chenyuandGitHub 0a98fd38b3 fix tests that failed locally on mac (#13872)
keccak output was silently broken without contiguous
2025-12-29 11:23:38 -05:00
0e409ff5ce fix indentation in UOp pretty_print for repeated references (#13857)
* fix correct indentation in UOp pretty_print for repeated references

When a UOp was referenced multiple times, the walrus operator notation
(e.g., x0:=) was correctly used for the first occurrence, but subsequent
references had misaligned indentation due to an extra space character.

Fix indentation misalignment in pretty_print() when UOps are referenced
multiple times.

* add simple unit tests for UOp repr

---------

Co-authored-by: chenyu <[email protected]>
2025-12-29 10:46:16 -05:00
George HotzandGitHub f1471a3b99 speed up rdna3 unit tests + add to CI (#13871)
* speed up rdna3 unit tests

* add test to CI

* faster and simpler

* speedups

* bugfixes

* use helper

* fix CI maybe

* test fixes

* llvm-21 on 24.04

* upd

* llvm-21

* fix test

* bring that back

* merge gen into lib

* test generators
2025-12-29 10:26:48 -05:00
h-vetinariandGitHub 37720fd6c0 also look for linux libraries in RHEL-themed paths (#13863) 2025-12-29 10:05:32 -05:00
George HotzandGitHub 25ef866e89 write python emulator from RDNA3 psuedocode in pdf (#13841)
* write python emulator from RDNA3 psuedocode in pdf

* emu2

* more emu

* working

* more psueod

* progress

* cleanups

* delete junk

* delete stale files

* just emu

* work

* emu compare

* bemu

* cleanups and more failures

* revert bench emu

* fix emu cmp

* four tests fail

* bugfixes

* dsl

* ext

* refactor

* dsl

* div scale fix

* test_emu

* fix emu tests

* pcode

* test pcode

* top imports

* fix test_emu to use run_asm

* emu tests on real hardware

* more tests

* more emu tests

* more

* work

* work

* bug fix

* bugfixes

* fix fp16 gemm

* all ops tests pass in emulator

* fix llvm tests

* fix a few more tests

* fix mockgpu timeout
2025-12-29 07:39:53 -05:00
nimlgenandGitHub 88eb230326 memory: correct pa allocator size (#13861) 2025-12-29 14:49:44 +03:00
qazalandGitHub f541540129 variable N for asm gemm (#13869)
* variable N for asm gemm

* cleanup spacing
2025-12-29 19:35:50 +09:00
nimlgenandGitHub c6769badc2 mockgpu: async support (#13868)
* mockgpu: async support

* cpu
2025-12-29 13:18:37 +03:00
qazalandGitHub fc5278746f mi350x assembly gemm cleanups (#13867) 2025-12-29 18:47:23 +09:00
George HotzandGitHub f07c39cfa4 hwtest fixes for rdna3 dsl (#13865) 2025-12-28 20:42:29 -05:00
104 changed files with 36322 additions and 19817 deletions
+56 -7
View File
@@ -5,6 +5,7 @@ env:
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
IGNORE_OOB: 0
on:
push:
@@ -36,6 +37,8 @@ jobs:
name: Docs
runs-on: ubuntu-22.04
timeout-minutes: 10
env:
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -102,15 +105,11 @@ jobs:
run: |
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Lint with ruff
run: |
pip3 install --upgrade --force-reinstall ruff==0.11.0
python3 -m ruff check extra/torch_backend/backend.py
- name: Test one op
run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
- name: Test ResNet-18
run: DEBUG=2 python3 extra/torch_backend/example.py
- name: My (custom) tests
- name: custom tests
run: python3 extra/torch_backend/test.py
- name: Test one op in torch tests
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
@@ -233,10 +232,11 @@ jobs:
run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string=' ' --recursive=y .
- name: Lint with ruff
run: |
pip3 install --upgrade --force-reinstall ruff==0.11.0
pip3 install --upgrade --force-reinstall ruff==0.14.10
python3 -m ruff check .
python3 -m ruff check examples/mlperf/ --ignore E501
python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
python3 -m ruff check extra/torch_backend/backend.py
- name: Run mypy
run: |
python -m mypy --strict-equality --lineprecision-report .
@@ -310,7 +310,7 @@ jobs:
deps: testing_unit
python-version: '3.14'
- name: Test SPEC=2
run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
fuzzing:
name: Fuzzing
@@ -473,6 +473,8 @@ jobs:
name: Test LLM
runs-on: ubuntu-24.04
timeout-minutes: 15
env:
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -654,6 +656,53 @@ jobs:
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdasm:
name: AMD ASM IDE
runs-on: ubuntu-24.04
timeout-minutes: 10
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: rdna3-emu
deps: testing_minimal
amd: 'true'
- name: Install LLVM 21
run: |
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
sudo apt-get update
sudo apt-get install llvm-21 llvm-21-tools cloc
- name: RDNA3 Line Count
run: cloc --by-file extra/assembly/amd/*.py
- name: Run RDNA3 emulator tests
run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 dtype tests
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
testamdautogen:
name: AMD autogen
runs-on: ubuntu-24.04
timeout-minutes: 10
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: rdna3-autogen
pydeps: "pdfplumber"
- name: Verify AMD autogen is up to date
run: |
python -m extra.assembly.amd.pdf --arch all
git diff --exit-code extra/assembly/amd/autogen/
testnvidia:
strategy:
fail-fast: false
+17 -28
View File
@@ -34,33 +34,6 @@ result = graph_rewrite(uop, pm)
### Schedule Cache
Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
## Directory Structure
```
tinygrad/
├── tensor.py # Tensor class, user API
├── device.py # Buffer, device management
├── dtype.py # Data types
├── helpers.py # Utilities, environment vars
├── uop/
│ ├── ops.py # UOp class, Ops enum, PatternMatcher
│ ├── spec.py # UOp type verification
│ └── symbolic.py # Symbolic math simplification
├── engine/
│ ├── schedule.py # Schedule creation, caching
│ ├── realize.py # Tensor realization
│ ├── jit.py # JIT compilation
│ └── memory.py # Memory planning
├── schedule/
│ ├── rangeify.py # Convert movements to ranges
│ └── indexing.py # Index calculations
├── codegen/
│ ├── kernel.py # Kernel optimization
│ └── uopgraph.py # UOp graph transformations
├── renderer/ # Code generation (CUDA, Metal, etc.)
└── runtime/ # Device backends
```
## Testing
```bash
@@ -79,7 +52,7 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
## Common Environment Variables
- `DEBUG=1-4` - Increasing verbosity
- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
- `VIZ=1` - Enable graph visualization
- `SPEC=1` - Enable UOp spec verification
- `NOOPT=1` - Disable optimizations
@@ -100,6 +73,16 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
- Run tests before proposing commits
- Test with `SPEC=2` when modifying UOp-related code
## Auto-generated Files (DO NOT EDIT)
The following files are auto-generated and should never be edited manually:
- `extra/assembly/amd/autogen/{arch}/__init__.py` - Generated by `python -m extra.assembly.amd.dsl --arch {arch}`
- `extra/assembly/amd/autogen/{arch}/gen_pcode.py` - Generated by `python -m extra.assembly.amd.pcode --arch {arch}`
Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
To add missing instruction implementations, add them to `extra/assembly/amd/emu.py` instead.
## Style Notes
- 2-space indentation, 150 char line limit
@@ -225,3 +208,9 @@ Key patterns to watch (from ResNet50 benchmark):
- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
## 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.
+31
View File
@@ -0,0 +1,31 @@
An integrated environment for AMD GPU assembly and emulation
Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`
* pdf.py -- extract assembly format + instruction psuedocode from AMD PDF
* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
* pcode.py -- psuedocode execution environment. psuedocode should be transformed as little as possible.
* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same as LLVM.
tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
The ops tests also pass, but they are very slow, so you should run them one at a time.
`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`
When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`. While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware, so if a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's likely 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.
Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines. Count lines with `cloc --by-file extra/assembly/amd/*.py`
+581
View File
@@ -0,0 +1,581 @@
# RDNA3 assembler and disassembler
from __future__ import annotations
import re
from extra.assembly.amd.dsl import Inst, RawImm, Reg, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory
from extra.assembly.amd.dsl import VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
from extra.assembly.amd.dsl import SPECIAL_GPRS, SPECIAL_PAIRS, FLOAT_DEC, FLOAT_ENC, decode_src
from extra.assembly.amd.autogen.rdna3 import ins
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP,
VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOPDOp, SOP1Op, SOPKOp, SOPPOp, SMEMOp, DSOp, MUBUFOp)
def _matches_encoding(word: int, cls: type[Inst]) -> bool:
"""Check if word matches the encoding pattern of an instruction class."""
if cls._encoding is None: return False
bf, val = cls._encoding
return ((word >> bf.lo) & bf.mask()) == val
# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, MUBUF, MTBUF, MIMG, SMEM, EXP]
_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2] # SOP2/VOP2 are catch-alls
def detect_format(data: bytes) -> type[Inst]:
"""Detect instruction format from machine code bytes."""
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
word = int.from_bytes(data[:4], 'little')
# Check 64-bit formats first (bits[31:30] == 0b11)
if (word >> 30) == 0b11:
for cls in _FORMATS_64:
if _matches_encoding(word, cls):
return VOP3SD if cls is VOP3 and ((word >> 16) & 0x3ff) in Inst._VOP3SD_OPS else cls
raise ValueError(f"unknown 64-bit format word={word:#010x}")
# 32-bit formats
for cls in _FORMATS_32:
if _matches_encoding(word, cls): return cls
raise ValueError(f"unknown 32-bit format word={word:#010x}")
# ═══════════════════════════════════════════════════════════════════════════════
# CONSTANTS
# ═══════════════════════════════════════════════════════════════════════════════
HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
HWREG_IDS = {v.lower(): k for k, v in HWREG.items()}
MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
# ═══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{b}" if n == 1 else f"{p}[{b}:{b+n-1}]"
def _sreg(b: int, n: int = 1) -> str: return _reg("s", b, n)
def _vreg(b: int, n: int = 1) -> str: return _reg("v", b, n)
def _ttmp(b: int, n: int = 1) -> str: return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _sreg_or_ttmp(b: int, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
def _fmt_sdst(v: int, n: int = 1) -> str:
if v == 124: return "null"
if t := _ttmp(v, n): return t
if n > 1: return SPECIAL_PAIRS.get(v) or _sreg(v, n)
return SPECIAL_GPRS.get(v, f"s{v}")
def _fmt_src(v: int, n: int = 1) -> str:
if n == 1: return decode_src(v)
if v >= 256: return _vreg(v - 256, n)
if v <= 105: return _sreg(v, n)
if n == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if t := _ttmp(v, n): return t
return decode_src(v)
def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
def _src16(inst, v: int) -> str: return _fmt_v16(v) if v >= 256 else inst.lit(v) # format 16-bit src: vgpr.h/l or literal
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool, any_hi: bool) -> str:
"""Format VOP3 source operand with modifiers."""
if n > 1: s = _fmt_src(v, n)
elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else (f"v{v - 256}.l" if any_hi else inst.lit(v))
else: s = inst.lit(v)
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
"""Format op_sel modifier string."""
if not need: return ""
if is16_d and (opsel & 8): return f" op_sel:[1,1,1{',1' if n == 3 else ''}]"
if n == 3: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
# ═══════════════════════════════════════════════════════════════════════════════
# DISASSEMBLER
# ═══════════════════════════════════════════════════════════════════════════════
def _disasm_vop1(inst: VOP1) -> str:
name = inst.op_name.lower()
if inst.op in (VOP1Op.V_NOP, VOP1Op.V_PIPEFLUSH): return name
if inst.op == VOP1Op.V_READFIRSTLANE_B32: return f"v_readfirstlane_b32 {decode_src(inst.vdst)}, v{inst.src0 - 256 if inst.src0 >= 256 else inst.src0}"
# 16-bit dst: uses .h/.l suffix (determined by name pattern, not dtype - e.g. sat_pk_u8_i16 outputs 8-bit but uses 16-bit encoding)
parts = name.split('_')
is_16d = any(p in ('f16','i16','u16','b16') for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in ('f16','i16','u16','b16') and 'cvt' not in name)
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
src = _fmt_src(inst.src0, inst.src_regs(0)) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if inst.is_src_16(0) and 'sat_pk' not in name else inst.lit(inst.src0)
return f"{name}_e32 {dst}, {src}"
def _disasm_vop2(inst: VOP2) -> str:
name = inst.op_name.lower()
suf = "" if inst.op == VOP2Op.V_DOT2ACC_F32_F16 else "_e32"
# fmaak: dst = src0 * vsrc1 + K, fmamk: dst = src0 * K + vsrc1
if inst.op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{inst._literal:x}"
if inst.op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{inst._literal:x}, v{inst.vsrc1}"
if inst.is_16bit(): return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (", vcc_lo" if inst.op == VOP2Op.V_CNDMASK_B32 else "")
def _disasm_vopc(inst: VOPC) -> str:
name = inst.op_name.lower()
s0 = _fmt_src(inst.src0, inst.src_regs(0)) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if inst.is_16bit() else inst.lit(inst.src0)
s1 = _vreg(inst.vsrc1, inst.src_regs(1)) if inst.src_regs(1) > 1 else _fmt_v16(inst.vsrc1, 0, 128) if inst.is_16bit() else f"v{inst.vsrc1}"
return f"{name}_e32 {s0}, {s1}" if inst.op.value >= 128 else f"{name}_e32 vcc_lo, {s0}, {s1}"
NO_ARG_SOPP = {SOPPOp.S_ENDPGM, SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE}
def _disasm_sopp(inst: SOPP) -> str:
name = inst.op_name.lower()
if inst.op in NO_ARG_SOPP: return name
if inst.op == SOPPOp.S_WAITCNT:
vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if inst.op == SOPPOp.S_DELAY_ALU:
deps, skips = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3'], ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
return f"{name} {inst.simm16}" if name.startswith(('s_cbranch', 's_branch')) else f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
name = inst.op_name.lower()
if inst.op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
off_s = f"{decode_src(inst.soffset)} offset:0x{inst.offset:x}" if inst.offset and inst.soffset != 124 else f"0x{inst.offset:x}" if inst.offset else decode_src(inst.soffset)
sbase_idx, sbase_count = inst.sbase * 2, 4 if (8 <= inst.op.value <= 12 or name == 's_atc_probe_buffer') else 2
sbase_str = _fmt_src(sbase_idx, sbase_count) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
return f"{name} {_fmt_sdst(inst.sdata, inst.dst_regs())}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (inst.dlc, " dlc"))
def _disasm_flat(inst: FLAT) -> str:
name = inst.op_name.lower()
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
w = inst.dst_regs() * (2 if 'cmpswap' in name else 1)
mods = f"{f' offset:{off_val}' if off_val else ''}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
# saddr
if seg == 'flat' or inst.saddr == 0x7F: saddr_s = ""
elif inst.saddr == 124: saddr_s = ", off"
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr)}"
elif inst.saddr in SPECIAL_PAIRS: saddr_s = f", {SPECIAL_PAIRS[inst.saddr]}"
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr)}"
# addtid: no addr
if 'addtid' in name: return f"{instr} v{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
# addr width
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, 1 if seg == 'scratch' or (inst.saddr not in (0x7F, 124)) else 2)
data_s, vdst_s = _vreg(inst.data, w), _vreg(inst.vdst, w // 2 if 'cmpswap' in name else w)
if 'atomic' in name:
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if inst.glc else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
return f"{instr} {_vreg(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
def _disasm_ds(inst: DS) -> str:
op, name = inst.op, inst.op_name.lower()
gds = " gds" if inst.gds else ""
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
off2 = f" offset0:{inst.offset0} offset1:{inst.offset1}" if inst.offset0 or inst.offset1 else ""
w = inst.dst_regs()
d0, d1, dst, addr = _vreg(inst.data0, w), _vreg(inst.data1, w), _vreg(inst.vdst, w), f"v{inst.addr}"
if op == DSOp.DS_NOP: return name
if op == DSOp.DS_BVH_STACK_RTN_B32: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}, {_vreg(inst.data1, 4)}{off}{gds}"
if 'gws_sema' in name and op != DSOp.DS_GWS_SEMA_BR: return f"{name}{off}{gds}"
if 'gws_' in name: return f"{name} {addr}{off}{gds}"
if op in (DSOp.DS_CONSUME, DSOp.DS_APPEND): return f"{name} v{inst.vdst}{off}{gds}"
if 'gs_reg' in name: return f"{name} {_vreg(inst.vdst, 2)}, v{inst.data0}{off}{gds}"
if '2addr' in name:
if 'load' in name: return f"{name} {_vreg(inst.vdst, w*2)}, {addr}{off2}{gds}"
if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
return f"{name} {_vreg(inst.vdst, w*2)}, {addr}, {d0}, {d1}{off2}{gds}"
if 'load' in name: return f"{name} v{inst.vdst}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
if 'store' in name and not _has(name, 'cmp', 'xchg'):
return f"{name} v{inst.data0}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
if 'swizzle' in name or op == DSOp.DS_ORDERED_COUNT: return f"{name} v{inst.vdst}, {addr}{off}{gds}"
if 'permute' in name: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}{off}{gds}"
if 'condxchg' in name: return f"{name} {_vreg(inst.vdst, 2)}, {addr}, {_vreg(inst.data0, 2)}{off}{gds}"
if _has(name, 'cmpstore', 'mskor', 'wrap'):
return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
def _disasm_vop3(inst: VOP3) -> str:
op, name = inst.op, inst.op_name.lower()
# VOP3SD (shared encoding)
if isinstance(op, VOP3SDOp):
sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
def src(v, neg, n): s = _fmt_src(v, n) if n > 1 else inst.lit(v); return f"-{s}" if neg else s
s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, inst.src_regs(2))
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {srcs}" + _omod(inst.omod)
# Detect 16-bit operand sizes (for .h/.l suffix handling)
is16_d = is16_s = is16_s2 = False
if 'cvt_pk' in name: is16_s = name.endswith('16')
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', name):
is16_d, is16_s = _has(m.group(1), 'f16','i16','u16','b16'), _has(m.group(2), 'f16','i16','u16','b16')
is16_s2 = is16_s
elif re.match(r'v_mad_[iu]32_[iu]16', name): is16_s = True
elif 'pack_b32' in name: is16_s = is16_s2 = True
else: is16_d = is16_s = is16_s2 = inst.is_16bit()
any_hi = inst.opsel != 0
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, inst.src_regs(0), is16_s, any_hi)
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, inst.src_regs(1), is16_s, any_hi)
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, inst.src_regs(2), is16_s2, any_hi)
# Destination
dn = inst.dst_regs()
if op == VOP3Op.V_READLANE_B32: dst = _fmt_sdst(inst.vdst, 1)
elif dn > 1: dst = _vreg(inst.vdst, dn)
elif is16_d: dst = f"v{inst.vdst}.h" if (inst.opsel & 8) else f"v{inst.vdst}.l" if any_hi else f"v{inst.vdst}"
else: dst = f"v{inst.vdst}"
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
nonvgpr_opsel = (inst.src0 < 256 and (inst.opsel & 1)) or (inst.src1 < 256 and (inst.opsel & 2)) or (inst.src2 < 256 and (inst.opsel & 4))
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
if inst.op < 256: # VOPC
return f"{name}_e64 {s0}, {s1}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}"
if inst.op < 384: # VOP2
n = inst.num_srcs()
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
if inst.op < 512: # VOP1
return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{_opsel_str(inst.opsel, 1, need_opsel, is16_d)}{cl}{om}"
# Native VOP3
n = inst.num_srcs()
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
def _disasm_vop3sd(inst: VOP3SD) -> str:
name = inst.op_name.lower()
def src(v, neg, n): s = _fmt_src(v, n) if n > 1 else inst.lit(v); return f"-{s}" if neg else s
s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, inst.src_regs(2))
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
suffix = "_e64" if name.startswith('v_') and 'co_' in name else ""
return f"{name}{suffix} {dst}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if inst.clmp else ''}{_omod(inst.omod)}"
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal or inst.literal
vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), VOPDOp(inst.opx).name.lower(), VOPDOp(inst.opy).name.lower()
def half(n, vd, s0, vs1): return f"{n} v{vd}, {inst.lit(s0)}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}" if 'mov' in n else f"{n} v{vd}, {inst.lit(s0)}, v{vs1}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
def _disasm_vop3p(inst: VOP3P) -> str:
name = inst.op_name.lower()
is_wmma, n, is_fma_mix = 'wmma' in name, inst.num_srcs(), 'fma_mix' in name
if is_wmma:
sc = 2 if 'iu4' in name else 4 if 'iu8' in name else 8
src0, src1, src2, dst = _fmt_src(inst.src0, sc), _fmt_src(inst.src1, sc), _fmt_src(inst.src2, 8), _vreg(inst.vdst, 8)
else: src0, src1, src2, dst = _fmt_src(inst.src0, 1), _fmt_src(inst.src1, 1), _fmt_src(inst.src2, 1), f"v{inst.vdst}"
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
if is_fma_mix:
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if inst.clmp else [])
else:
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != (7 if n == 3 else 3) else []) + \
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
def _disasm_buf(inst: MUBUF | MTBUF) -> str:
name = inst.op_name.lower()
if inst.op in (MUBUFOp.BUFFER_GL0_INV, MUBUFOp.BUFFER_GL1_INV): return name
w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
{'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], 1)
if inst.tfe: w += 1
vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else f"v{inst.vaddr}" if inst.offen or inst.idxen else "off"
srsrc = _sreg_or_ttmp(inst.srsrc*4, 4)
mods = ([f"format:{inst.format}"] if isinstance(inst, MTBUF) else []) + [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.glc,"glc"),(inst.dlc,"dlc"),(inst.slc,"slc"),(inst.tfe,"tfe")] if c]
return f"{name} {_vreg(inst.vdata, w)}, {vaddr}, {srsrc}, {decode_src(inst.soffset)}{' ' + ' '.join(mods) if mods else ''}"
def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
"""Calculate vaddr register count for MIMG sample/gather operations."""
# 1d,2d,3d,cube,1d_arr,2d_arr,2d_msaa,2d_msaa_arr
base = [1, 2, 3, 3, 2, 3, 3, 4][dim] # address coords
grad = [1, 2, 3, 2, 1, 2, 2, 2][dim] # gradient coords (for derivatives)
if 'get_resinfo' in name: return 1 # only mip level
packed, unpacked = 0, 0
if '_mip' in name: packed += 1
elif 'sample' in name or 'gather' in name:
if '_o' in name: unpacked += 1 # offset
if re.search(r'_c(_|$)', name): unpacked += 1 # compare (not _cl)
if '_d' in name: unpacked += (grad + 1) & ~1 if '_g16' in name else grad*2 # derivatives
if '_b' in name: unpacked += 1 # bias
if '_l' in name and '_cl' not in name and '_lz' not in name: packed += 1 # LOD
if '_cl' in name: packed += 1 # clamp
return (base + packed + 1) // 2 + unpacked if a16 else base + packed + unpacked
def _disasm_mimg(inst: MIMG) -> str:
name = inst.op_name.lower()
srsrc_base = inst.srsrc * 4
srsrc_str = _sreg_or_ttmp(srsrc_base, 8)
# BVH intersect ray: special case with 4 SGPR srsrc
if 'bvh' in name:
vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {_sreg_or_ttmp(srsrc_base, 4)}{' a16' if inst.a16 else ''}"
# vdata width from dmask (gather4/msaa_load always 4), d16 packs, tfe adds 1
vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
if inst.d16: vdata = (vdata + 1) // 2
if inst.tfe: vdata += 1
# vaddr width
dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"dim_{inst.dim}"
vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
vaddr_str = f"v{inst.vaddr}" if vaddr == 1 else _vreg(inst.vaddr, vaddr)
# modifiers
mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask and (inst.dmask != 15 or 'atomic' in name) else []
mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
for flag, mod in [(inst.unrm,"unorm"),(inst.glc,"glc"),(inst.slc,"slc"),(inst.dlc,"dlc"),(inst.r128,"r128"),
(inst.a16,"a16"),(inst.tfe,"tfe"),(inst.lwe,"lwe"),(inst.d16,"d16")]:
if flag: mods.append(mod)
# ssamp for sample/gather/get_lod
ssamp_str = ""
if 'sample' in name or 'gather' in name or 'get_lod' in name:
ssamp_str = ", " + _sreg_or_ttmp(inst.ssamp * 4, 4)
return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
def _disasm_sop1(inst: SOP1) -> str:
op, name = inst.op, inst.op_name.lower()
if op == SOP1Op.S_GETPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
if op in (SOP1Op.S_SETPC_B64, SOP1Op.S_RFE_B64): return f"{name} {_fmt_src(inst.ssrc0, 2)}"
if op == SOP1Op.S_SWAPPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {_fmt_src(inst.ssrc0, 2)}"
if op in (SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64): return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, sendmsg({MSG.get(inst.ssrc0, str(inst.ssrc0))})"
return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, {inst.lit(inst.ssrc0) if inst.src_regs(0) == 1 else _fmt_src(inst.ssrc0, inst.src_regs(0))}"
def _disasm_sop2(inst: SOP2) -> str:
return f"{inst.op_name.lower()} {_fmt_sdst(inst.sdst, inst.dst_regs())}, {inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0))}, {inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, inst.src_regs(1))}"
def _disasm_sopc(inst: SOPC) -> str:
return f"{inst.op_name.lower()} {_fmt_src(inst.ssrc0, inst.src_regs(0))}, {_fmt_src(inst.ssrc1, inst.src_regs(1))}"
def _disasm_sopk(inst: SOPK) -> str:
op, name = inst.op, inst.op_name.lower()
if op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
if op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32):
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
hs = f"0x{inst.simm16:x}" if hid in (16, 17) else f"hwreg({HWREG.get(hid, str(hid))}, {hoff}, {hsz})"
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1)}" if op == SOPKOp.S_SETREG_B32 else f"{name} {_fmt_sdst(inst.sdst, 1)}, {hs}"
return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, 0x{inst.simm16:x}"
def _disasm_vinterp(inst: VINTERP) -> str:
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
return f"{inst.op_name.lower()} v{inst.vdst}, {inst.lit(inst.src0, inst.neg & 1)}, {inst.lit(inst.src1, inst.neg & 2)}, {inst.lit(inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
DISASM_HANDLERS = {VOP1: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, MUBUF: _disasm_buf, MTBUF: _disasm_buf,
MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk}
def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
# ═══════════════════════════════════════════════════════════════════════════════
# ASSEMBLER
# ═══════════════════════════════════════════════════════════════════════════════
SPEC_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'vcc': RawImm(106), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125),
'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'exec': RawImm(126), 'scc': RawImm(253), 'src_scc': RawImm(253)}
FLOATS = {str(k): k for k in FLOAT_ENC} # Valid float literal strings: '0.5', '-0.5', '1.0', etc.
REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
SPEC_DSL = {'vcc_lo': 'VCC_LO', 'vcc_hi': 'VCC_HI', 'vcc': 'VCC_LO', 'null': 'NULL', 'off': 'OFF', 'm0': 'M0',
'exec_lo': 'EXEC_LO', 'exec_hi': 'EXEC_HI', 'exec': 'EXEC_LO', 'scc': 'SCC', 'src_scc': 'SCC'}
def _op2dsl(op: str) -> str:
op = op.strip()
neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
if neg: op = op[1:]
abs_ = (op.startswith('|') and op.endswith('|')) or (op.startswith('abs(') and op.endswith(')'))
if abs_: op = op[1:-1] if op.startswith('|') else op[4:-1]
hi = ".h" if op.endswith('.h') else ".l" if op.endswith('.l') else ""
if hi: op = op[:-2]
lo = op.lower()
def wrap(b): return f"{'-' if neg else ''}abs({b}){hi}" if abs_ else f"-{b}{hi}" if neg else f"{b}{hi}"
if lo in SPEC_DSL: return wrap(SPEC_DSL[lo])
if op in FLOATS: return wrap(op)
rp = {'s': 's', 'v': 'v', 't': 'ttmp', 'ttmp': 'ttmp'}
if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', lo): return wrap(f"{rp[m.group(1)]}[{m.group(2)}:{m.group(3)}]")
if m := re.match(r'^([svt](?:tmp)?)(\d+)$', lo): return wrap(f"{rp[m.group(1)]}[{m.group(2)}]")
if re.match(r'^-?\d+$|^-?0x[0-9a-fA-F]+$', op): return f"SrcMod({op}, neg={neg}, abs_={abs_})" if neg or abs_ else op
return wrap(op)
def _parse_ops(s: str) -> list[str]:
ops, cur, depth, pipe = [], "", 0, False
for c in s:
if c in '[(': depth += 1
elif c in '])': depth -= 1
elif c == '|': pipe = not pipe
if c == ',' and depth == 0 and not pipe: ops.append(cur.strip()); cur = ""
else: cur += c
if cur.strip(): ops.append(cur.strip())
return ops
def _extract(text: str, pat: str, flags=re.I):
if m := re.search(pat, text, flags): return m, text[:m.start()] + text[m.end():]
return None, text
def get_dsl(text: str) -> str:
text, kw = text.strip(), []
# Extract modifiers
for pat, val in [(r'\s+mul:2(?:\s|$)', 1), (r'\s+mul:4(?:\s|$)', 2), (r'\s+div:2(?:\s|$)', 3)]:
if (m := _extract(text, pat))[0]: kw.append(f'omod={val}'); text = m[1]; break
if (m := _extract(text, r'\s+clamp(?:\s|$)'))[0]: kw.append('clmp=1'); text = m[1]
opsel, m, text = None, *_extract(text, r'\s+op_sel:\[([^\]]+)\]')
if m:
bits, mn = [int(x.strip()) for x in m.group(1).split(',')], text.split()[0].lower()
is3p = mn.startswith(('v_pk_', 'v_wmma_', 'v_dot'))
opsel = (bits[0] | (bits[1] << 1) | (bits[2] << 2)) if len(bits) == 3 and is3p else \
(bits[0] | (bits[1] << 1) | (bits[2] << 3)) if len(bits) == 3 else sum(b << i for i, b in enumerate(bits))
m, text = _extract(text, r'\s+wait_exp:(\d+)'); waitexp = m.group(1) if m else None
m, text = _extract(text, r'\s+offset:(0x[0-9a-fA-F]+|-?\d+)'); off_val = m.group(1) if m else None
m, text = _extract(text, r'\s+dlc(?:\s|$)'); dlc = 1 if m else None
m, text = _extract(text, r'\s+glc(?:\s|$)'); glc = 1 if m else None
m, text = _extract(text, r'\s+slc(?:\s|$)'); slc = 1 if m else None
m, text = _extract(text, r'\s+neg_lo:\[([^\]]+)\]'); neg_lo = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
m, text = _extract(text, r'\s+neg_hi:\[([^\]]+)\]'); neg_hi = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
if waitexp: kw.append(f'waitexp={waitexp}')
parts = text.replace(',', ' ').split()
if not parts: raise ValueError("empty instruction")
mn, op_str = parts[0].lower(), text[len(parts[0]):].strip()
ops, args = _parse_ops(op_str), [_op2dsl(o) for o in _parse_ops(op_str)]
# s_waitcnt
if mn == 's_waitcnt':
vm, exp, lgkm = 0x3f, 0x7, 0x3f
for p in op_str.replace(',', ' ').split():
if m := re.match(r'vmcnt\((\d+)\)', p): vm = int(m.group(1))
elif m := re.match(r'expcnt\((\d+)\)', p): exp = int(m.group(1))
elif m := re.match(r'lgkmcnt\((\d+)\)', p): lgkm = int(m.group(1))
elif re.match(r'^0x[0-9a-f]+$|^\d+$', p): return f"s_waitcnt(simm16={int(p, 0)})"
return f"s_waitcnt(simm16={waitcnt(vm, exp, lgkm)})"
# VOPD
if '::' in text:
xp, yp = text.split('::')
xps, yps = xp.strip().replace(',', ' ').split(), yp.strip().replace(',', ' ').split()
xo, yo = [_op2dsl(p) for p in xps[1:]], [_op2dsl(p) for p in yps[1:]]
vdx, sx0, vsx1 = xo[0], xo[1] if len(xo) > 1 else '0', xo[2] if len(xo) > 2 else 'v[0]'
vdy, sy0, vsy1 = yo[0], yo[1] if len(yo) > 1 else '0', yo[2] if len(yo) > 2 else 'v[0]'
lit = xo[3] if 'fmaak' in xps[0].lower() and len(xo) > 3 else yo[3] if 'fmaak' in yps[0].lower() and len(yo) > 3 else None
if 'fmamk' in xps[0].lower() and len(xo) > 3: lit, vsx1 = xo[2], xo[3]
elif 'fmamk' in yps[0].lower() and len(yo) > 3: lit, vsy1 = yo[2], yo[3]
return f"VOPD(VOPDOp.{xps[0].upper()}, VOPDOp.{yps[0].upper()}, vdstx={vdx}, vdsty={vdy}, srcx0={sx0}, vsrcx1={vsx1}, srcy0={sy0}, vsrcy1={vsy1}{f', literal={lit}' if lit else ''})"
# Special instructions
if mn == 's_setreg_imm32_b32': raise ValueError(f"unsupported: {mn}")
if mn in ('s_setpc_b64', 's_rfe_b64'): return f"{mn}(ssrc0={args[0]})"
if mn in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'): return f"{mn}(sdst={args[0]}, ssrc0=RawImm({args[1].strip()}))"
if mn == 's_version': return f"{mn}(simm16={args[0]})"
if mn == 's_setreg_b32': return f"{mn}(simm16={args[0]}, sdst={args[1]})"
# SMEM
if mn in SMEM_OPS:
gs, ds = ", glc=1" if glc else "", ", dlc=1" if dlc else ""
if len(ops) >= 3 and re.match(r'^-?[0-9]|^-?0x', ops[2].strip().lower()):
return f"{mn}(sdata={args[0]}, sbase={args[1]}, offset={args[2]}, soffset=RawImm(124){gs}{ds})"
if off_val and len(ops) >= 3: return f"{mn}(sdata={args[0]}, sbase={args[1]}, offset={off_val}, soffset={args[2]}{gs}{ds})"
if len(ops) >= 3: return f"{mn}(sdata={args[0]}, sbase={args[1]}, soffset={args[2]}{gs}{ds})"
# Buffer
if mn.startswith('buffer_') and len(ops) >= 2 and ops[1].strip().lower() == 'off':
return f"{mn}(vdata={args[0]}, vaddr=0, srsrc={args[2]}, soffset={f'RawImm({args[3].strip()})' if len(args) > 3 else 'RawImm(0)'})"
# FLAT/GLOBAL/SCRATCH load/store/atomic - saddr needs RawImm(124) for off/null
def _saddr(a): return 'RawImm(124)' if a in ('OFF', 'NULL') else a
flat_mods = f"{f', offset={off_val}' if off_val else ''}{', glc=1' if glc else ''}{', slc=1' if slc else ''}{', dlc=1' if dlc else ''}"
for pre, flds in [('flat_load','vdst,addr,saddr'), ('global_load','vdst,addr,saddr'), ('scratch_load','vdst,addr,saddr'),
('flat_store','addr,data,saddr'), ('global_store','addr,data,saddr'), ('scratch_store','addr,data,saddr')]:
if mn.startswith(pre) and len(args) >= 2:
f0, f1, f2 = flds.split(',')
return f"{mn}({f0}={args[0]}, {f1}={args[1]}{f', {f2}={_saddr(args[2])}' if len(args) >= 3 else ', saddr=RawImm(124)'}{flat_mods})"
for pre in ('flat_atomic', 'global_atomic', 'scratch_atomic'):
if mn.startswith(pre):
if glc and len(args) >= 3: return f"{mn}(vdst={args[0]}, addr={args[1]}, data={args[2]}{f', saddr={_saddr(args[3])}' if len(args) >= 4 else ', saddr=RawImm(124)'}{flat_mods})"
if len(args) >= 2: return f"{mn}(addr={args[0]}, data={args[1]}{f', saddr={_saddr(args[2])}' if len(args) >= 3 else ', saddr=RawImm(124)'}{flat_mods})"
# DS instructions
if mn.startswith('ds_'):
off0, off1 = (str(int(off_val, 0) & 0xff), str((int(off_val, 0) >> 8) & 0xff)) if off_val else ("0", "0")
gds_s = ", gds=1" if 'gds' in text.lower().split()[-1:] else ""
off_kw = f", offset0={off0}, offset1={off1}{gds_s}"
if mn == 'ds_nop' or mn in ('ds_gws_sema_v', 'ds_gws_sema_p', 'ds_gws_sema_release_all'): return f"{mn}({off_kw.lstrip(', ')})"
if 'gws_' in mn: return f"{mn}(addr={args[0]}{off_kw})"
if 'consume' in mn or 'append' in mn: return f"{mn}(vdst={args[0]}{off_kw})"
if 'gs_reg' in mn: return f"{mn}(vdst={args[0]}, data0={args[1]}{off_kw})"
if '2addr' in mn:
if 'load' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
if 'store' in mn and 'xchg' not in mn: return f"{mn}(addr={args[0]}, data0={args[1]}, data1={args[2]}{off_kw})"
return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})"
if 'load' in mn: return f"{mn}(vdst={args[0]}{off_kw})" if 'addtid' in mn else f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
if 'store' in mn and not _has(mn, 'cmp', 'xchg'):
return f"{mn}(data0={args[0]}{off_kw})" if 'addtid' in mn else f"{mn}(addr={args[0]}, data0={args[1]}{off_kw})"
if 'swizzle' in mn or 'ordered_count' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}{off_kw})"
if 'permute' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})"
if 'bvh' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})"
if 'condxchg' in mn: return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})"
if _has(mn, 'cmpstore', 'mskor', 'wrap'):
return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}, data1={args[3]}{off_kw})" if '_rtn' in mn else f"{mn}(addr={args[0]}, data0={args[1]}, data1={args[2]}{off_kw})"
return f"{mn}(vdst={args[0]}, addr={args[1]}, data0={args[2]}{off_kw})" if '_rtn' in mn else f"{mn}(addr={args[0]}, data0={args[1]}{off_kw})"
# v_fmaak/v_fmamk literal extraction
lit_s = ""
if mn in ('v_fmaak_f32', 'v_fmaak_f16') and len(args) == 4: lit_s, args = f", literal={args[3].strip()}", args[:3]
elif mn in ('v_fmamk_f32', 'v_fmamk_f16') and len(args) == 4: lit_s, args = f", literal={args[2].strip()}", [args[0], args[1], args[3]]
# VCC ops cleanup
vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'}
if mn.replace('_e32', '') in vcc_ops and len(args) >= 5: mn, args = mn.replace('_e32', '') + '_e32', [args[0], args[2], args[3]]
if mn.replace('_e64', '') in vcc_ops and mn.endswith('_e64'): mn = mn.replace('_e64', '')
if mn.startswith('v_cmp') and not mn.endswith('_e64') and len(args) >= 3 and ops[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'): args = args[1:]
if 'cmpx' in mn and mn.endswith('_e64') and len(args) == 2: args = ['RawImm(126)'] + args
fn = mn.replace('.', '_')
if opsel is not None: args = [re.sub(r'\.[hl]$', '', a) for a in args]
# v_fma_mix*: extract inline neg/abs modifiers
if 'fma_mix' in mn and neg_lo is None and neg_hi is None:
inline_neg, inline_abs, clean_args = 0, 0, [args[0]]
for i, op in enumerate(ops[1:4]):
op = op.strip()
neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
if neg: op = op[1:]
abs_ = op.startswith('|') and op.endswith('|')
if abs_: op = op[1:-1]
if neg: inline_neg |= (1 << i)
if abs_: inline_abs |= (1 << i)
clean_args.append(_op2dsl(op))
args = clean_args + args[4:]
if inline_neg: neg_lo = inline_neg
if inline_abs: neg_hi = inline_abs
all_kw = list(kw)
if lit_s: all_kw.append(lit_s.lstrip(', '))
if opsel is not None: all_kw.append(f'opsel={opsel}')
if neg_lo is not None: all_kw.append(f'neg={neg_lo}')
if neg_hi is not None: all_kw.append(f'neg_hi={neg_hi}')
if 'bvh' in mn and 'intersect_ray' in mn: all_kw.extend(['dmask=15', 'unrm=1', 'r128=1'])
a_str, kw_str = ', '.join(args), ', '.join(all_kw)
return f"{fn}({a_str}, {kw_str})" if kw_str and a_str else f"{fn}({kw_str})" if kw_str else f"{fn}({a_str})"
def asm(text: str) -> Inst:
dsl = get_dsl(text)
ns = {n: getattr(ins, n) for n in dir(ins) if not n.startswith('_')}
ns.update({'s': s, 'v': v, 'ttmp': ttmp, 'abs': abs, 'RawImm': RawImm, 'SrcMod': SrcMod, 'VGPR': VGPR, 'SGPR': SGPR, 'TTMP': TTMP,
'VCC_LO': VCC_LO, 'VCC_HI': VCC_HI, 'VCC': VCC, 'EXEC_LO': EXEC_LO, 'EXEC_HI': EXEC_HI, 'EXEC': EXEC, 'SCC': SCC, 'M0': M0, 'NULL': NULL, 'OFF': OFF})
try: return eval(dsl, ns)
except NameError:
if m := re.match(r'^(v_\w+)(\(.*\))$', dsl): return eval(f"{m.group(1)}_e32{m.group(2)}", ns)
raise
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+533
View File
@@ -0,0 +1,533 @@
# library for RDNA3 assembly DSL
# mypy: ignore-errors
from __future__ import annotations
import struct, math, re
from enum import IntEnum
from functools import cache, cached_property
from typing import overload, Annotated, TypeVar, Generic
from extra.assembly.amd.autogen.rdna3.enum import (VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, SOP1Op, SOP2Op,
SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp, VINTERPOp)
# Common masks and bit conversion functions
MASK32, MASK64 = 0xffffffff, 0xffffffffffffffff
_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): return _struct_f.unpack(_struct_I.pack(i & MASK32))[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: return _struct_I.unpack(_struct_f.pack(f))[0]
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
def _i16(f):
if math.isnan(f): return 0x7e00
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
try: return _struct_H.unpack(_struct_e.pack(f))[0]
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
def _i64(f):
if math.isnan(f): return 0x7ff8000000000000
if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
try: return _struct_Q.unpack(_struct_d.pack(f))[0]
except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
# Instruction spec - register counts and dtypes derived from instruction names
_REGS = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16,
'F32': 1, 'I32': 1, 'U32': 1, 'F64': 2, 'I64': 2, 'U64': 2,
'F16': 1, 'I16': 1, 'U16': 1, 'B16': 1, 'I8': 1, 'U8': 1, 'B8': 1}
_CVT_RE = re.compile(r'CVT_([FIUB]\d+)_([FIUB]\d+)$')
_MAD_MUL_RE = re.compile(r'(?:MAD|MUL)_([IU]\d+)_([IU]\d+)$')
_PACK_RE = re.compile(r'PACK_([FIUB]\d+)_([FIUB]\d+)$')
_DST_SRC_RE = re.compile(r'_([FIUB]\d+)_([FIUB]\d+)$')
_SINGLE_RE = re.compile(r'_([FIUB](?:32|64|16|8|96|128|256|512))$')
@cache
def _suffix(name: str) -> tuple[str | None, str | None]:
name = name.upper()
if m := _CVT_RE.search(name): return m.group(1), m.group(2)
if m := _MAD_MUL_RE.search(name): return m.group(1), m.group(2)
if m := _PACK_RE.search(name): return m.group(1), m.group(2)
if m := _DST_SRC_RE.search(name): return m.group(1), m.group(2)
if m := _SINGLE_RE.search(name): return m.group(1), m.group(1)
return None, None
_SPECIAL_REGS = {
'V_LSHLREV_B64': (2, 1, 2, 1), 'V_LSHRREV_B64': (2, 1, 2, 1), 'V_ASHRREV_I64': (2, 1, 2, 1),
'S_LSHL_B64': (2, 2, 1, 1), 'S_LSHR_B64': (2, 2, 1, 1), 'S_ASHR_I64': (2, 2, 1, 1),
'S_BFE_U64': (2, 2, 1, 1), 'S_BFE_I64': (2, 2, 1, 1), 'S_BFM_B64': (2, 1, 1, 1),
'S_BITSET0_B64': (2, 1, 1, 1), 'S_BITSET1_B64': (2, 1, 1, 1),
'S_BITCMP0_B64': (1, 2, 1, 1), 'S_BITCMP1_B64': (1, 2, 1, 1),
'V_LDEXP_F64': (2, 2, 1, 1), 'V_TRIG_PREOP_F64': (2, 2, 1, 1),
'V_CMP_CLASS_F64': (1, 2, 1, 1), 'V_CMPX_CLASS_F64': (1, 2, 1, 1),
'V_CMP_CLASS_F32': (1, 1, 1, 1), 'V_CMPX_CLASS_F32': (1, 1, 1, 1),
'V_CMP_CLASS_F16': (1, 1, 1, 1), 'V_CMPX_CLASS_F16': (1, 1, 1, 1),
'V_MAD_U64_U32': (2, 1, 1, 2), 'V_MAD_I64_I32': (2, 1, 1, 2),
'V_QSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_U32_U8': (4, 2, 1, 4),
}
_SPECIAL_DTYPE = {
'V_LSHLREV_B64': ('B64', 'U32', 'B64', None), 'V_LSHRREV_B64': ('B64', 'U32', 'B64', None), 'V_ASHRREV_I64': ('I64', 'U32', 'I64', None),
'S_LSHL_B64': ('B64', 'B64', 'U32', None), 'S_LSHR_B64': ('B64', 'B64', 'U32', None), 'S_ASHR_I64': ('I64', 'I64', 'U32', None),
'S_BFE_U64': ('U64', 'U64', 'U32', None), 'S_BFE_I64': ('I64', 'I64', 'U32', None),
'S_BFM_B64': ('B64', 'U32', 'U32', None), 'S_BITSET0_B64': ('B64', 'U32', None, None), 'S_BITSET1_B64': ('B64', 'U32', None, None),
'S_BITCMP0_B64': ('SCC', 'B64', 'U32', None), 'S_BITCMP1_B64': ('SCC', 'B64', 'U32', None),
'V_LDEXP_F64': ('F64', 'F64', 'I32', None), 'V_TRIG_PREOP_F64': ('F64', 'F64', 'U32', None),
'V_CMP_CLASS_F64': ('VCC', 'F64', 'U32', None), 'V_CMPX_CLASS_F64': ('EXEC', 'F64', 'U32', None),
'V_CMP_CLASS_F32': ('VCC', 'F32', 'U32', None), 'V_CMPX_CLASS_F32': ('EXEC', 'F32', 'U32', None),
'V_CMP_CLASS_F16': ('VCC', 'F16', 'U32', None), 'V_CMPX_CLASS_F16': ('EXEC', 'F16', 'U32', None),
'V_MAD_U64_U32': ('U64', 'U32', 'U32', 'U64'), 'V_MAD_I64_I32': ('I64', 'I32', 'I32', 'I64'),
'V_QSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'), 'V_MQSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'),
'V_MQSAD_U32_U8': ('B128', 'B64', 'B64', 'B128'),
}
@cache
def spec_regs(name: str) -> tuple[int, int, int, int]:
uname = name.upper()
if uname in _SPECIAL_REGS: return _SPECIAL_REGS[uname]
if 'SAD' in uname and 'U8' in uname and 'QSAD' not in uname and 'MQSAD' not in uname: return 1, 1, 1, 1
dst_suf, src_suf = _suffix(name)
return _REGS.get(dst_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1)
@cache
def spec_dtype(name: str) -> tuple[str | None, str | None, str | None, str | None]:
uname = name.upper()
if uname in _SPECIAL_DTYPE: return _SPECIAL_DTYPE[uname]
if 'SAD' in uname and ('U8' in uname or 'U16' in uname) and 'QSAD' not in uname and 'MQSAD' not in uname: return 'U32', 'U32', 'U32', 'U32'
if '_CMP_' in uname or '_CMPX_' in uname:
dst_suf, src_suf = _suffix(name)
return 'EXEC' if '_CMPX_' in uname else 'VCC', src_suf, src_suf, None
dst_suf, src_suf = _suffix(name)
return dst_suf, src_suf, src_suf, src_suf
_F16_RE = re.compile(r'_[FIUB]16(?:_|$)')
_F64_RE = re.compile(r'_[FIUB]64(?:_|$)')
@cache
def spec_is_16bit(name: str) -> bool:
uname = name.upper()
if 'SAD' in uname or 'PACK' in uname or '_PK_' in uname or 'SAT_PK' in uname or 'DOT2' in uname: return False
if '_F32' in uname or '_I32' in uname or '_U32' in uname or '_B32' in uname: return False
return bool(_F16_RE.search(uname))
@cache
def spec_is_64bit(name: str) -> bool: return bool(_F64_RE.search(name.upper()))
_3SRC = {'FMA', 'MAD', 'MIN3', 'MAX3', 'MED3', 'DIV_FIX', 'DIV_FMAS', 'DIV_SCALE', 'SAD', 'LERP', 'ALIGN', 'CUBE', 'BFE', 'BFI',
'PERM_B32', 'PERMLANE', 'CNDMASK', 'XOR3', 'OR3', 'ADD3', 'LSHL_OR', 'AND_OR', 'LSHL_ADD', 'ADD_LSHL', 'XAD', 'MAXMIN',
'MINMAX', 'DOT2', 'DOT4', 'DOT8', 'WMMA', 'CVT_PK_U8', 'MULLIT', 'CO_CI'}
_2SRC = {'FMAC'} # FMAC uses dst as implicit accumulator, so only 2 explicit sources
def spec_num_srcs(name: str) -> int:
name = name.upper()
if any(k in name for k in _2SRC): return 2
return 3 if any(k in name for k in _3SRC) else 2
def is_dtype_16(dt: str | None) -> bool: return dt is not None and '16' in dt
def is_dtype_64(dt: str | None) -> bool: return dt is not None and '64' in dt
# Bit field DSL
class BitField:
def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name, self._marker = hi, lo, name, None
def __set_name__(self, owner, name):
import typing
self.name, self._owner = name, owner
# Cache marker at class definition time
hints = typing.get_type_hints(owner, include_extras=True)
if name in hints:
hint = hints[name]
if typing.get_origin(hint) is Annotated:
args = typing.get_args(hint)
self._marker = args[1] if len(args) > 1 else None
def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val) # type: ignore
def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
@property
def marker(self) -> type | None: return self._marker
@overload
def __get__(self, obj: None, objtype: type) -> BitField: ...
@overload
def __get__(self, obj: object, objtype: type | None = None) -> int: ...
def __get__(self, obj, objtype=None):
if obj is None: return self
val = unwrap(obj._values.get(self.name, 0))
# Convert to IntEnum if marker is an IntEnum subclass
if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
# VOP3 with VOPC opcodes (0-255) -> VOPCOp, VOP3SD opcodes -> VOP3SDOp
if self.marker is VOP3Op:
if val < 256: return VOPCOp(val)
if val in Inst._VOP3SD_OPS: return VOP3SDOp(val)
try: return self.marker(val)
except ValueError: pass
return val
class _Bits:
def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
bits = _Bits()
# Source operand with modifiers - base class for anything that can be a src with neg/abs
class SrcMod:
__slots__ = ('val', 'neg', 'abs_')
def __init__(self, val: int, neg: bool = False, abs_: bool = False): self.val, self.neg, self.abs_ = val, neg, abs_
def __repr__(self): return f"{'-' if self.neg else ''}{'|' if self.abs_ else ''}{self.val}{'|' if self.abs_ else ''}"
def __neg__(self): return SrcMod(self.val, not self.neg, self.abs_)
def __abs__(self): return SrcMod(self.val, self.neg, True)
# Register types
class Reg(SrcMod):
__slots__ = ('idx', 'count', 'hi')
def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False, abs_: bool = False):
self.idx, self.count, self.hi = idx, count, hi
super().__init__(idx, neg, abs_)
def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
def __neg__(self): return self.__class__(self.idx, self.count, self.hi, not self.neg, self.abs_)
def __abs__(self): return self.__class__(self.idx, self.count, self.hi, self.neg, True)
@property
def l(self): return self.__class__(self.idx, self.count, False, self.neg, self.abs_)
@property
def h(self): return self.__class__(self.idx, self.count, True, self.neg, self.abs_)
T = TypeVar('T', bound=Reg)
class _RegFactory(Generic[T]):
def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
@overload
def __getitem__(self, key: int) -> Reg: ...
@overload
def __getitem__(self, key: slice) -> Reg: ...
def __getitem__(self, key: int | slice) -> Reg:
return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
def __repr__(self): return f"<{self._name} factory>"
class SGPR(Reg): pass
class VGPR(Reg): pass
class TTMP(Reg): pass
s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
# Special registers as SrcMod objects (support -VCC_LO, abs(EXEC_LO), etc.)
VCC_LO, VCC_HI, VCC = SrcMod(106), SrcMod(107), SrcMod(106)
EXEC_LO, EXEC_HI, EXEC = SrcMod(126), SrcMod(127), SrcMod(126)
SCC, M0, NULL, OFF = SrcMod(253), SrcMod(125), SrcMod(124), SrcMod(124)
# Field type markers (runtime classes for validation)
class _SSrc: pass
class _Src: pass
class _Imm: pass
class _SImm: pass
class _VDSTYEnc: pass # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
class _SGPRField: pass
class _VGPRField: pass
# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
SSrc = Annotated[BitField, _SSrc]
Src = Annotated[BitField, _Src]
Imm = Annotated[BitField, _Imm]
SImm = Annotated[BitField, _SImm]
VDSTYEnc = Annotated[BitField, _VDSTYEnc]
SGPRField = Annotated[BitField, _SGPRField]
VGPRField = Annotated[BitField, _VGPRField]
class RawImm:
def __init__(self, val: int): self.val = val
def __repr__(self): return f"RawImm({self.val})"
def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
def unwrap(val) -> int:
if isinstance(val, RawImm): return val.val
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special registers like VCC_LO, NULL
if hasattr(val, 'value'): return val.value # IntEnum
if hasattr(val, 'idx'): return val.idx # Reg
return val
# Encoding/decoding constants
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
FLOAT_DEC = {v: str(k) for k, v in FLOAT_ENC.items()}
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata', 'vsrc1'}
def _encode_reg(val: Reg) -> int: return (108 if isinstance(val, TTMP) else 0) + val.idx
def _is_inline_const(v: int) -> bool: return 0 <= v <= 127 or 128 <= v <= 208 or 240 <= v <= 255
def encode_src(val) -> int:
if isinstance(val, VGPR): return 256 + _encode_reg(val)
if isinstance(val, Reg): return _encode_reg(val)
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val if _is_inline_const(val.val) else 255
if hasattr(val, 'value'): return val.value # IntEnum
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
if isinstance(val, int): return 128 + val if 0 <= val <= 64 else 192 - val if -16 <= val <= -1 else 255
return 255
def decode_src(val: int) -> str:
if val <= 105: return f"s{val}"
if val in SPECIAL_GPRS: return SPECIAL_GPRS[val]
if val in FLOAT_DEC: return FLOAT_DEC[val]
if 108 <= val <= 123: return f"ttmp{val - 108}"
if 128 <= val <= 192: return str(val - 128)
if 193 <= val <= 208: return str(-(val - 192))
if 256 <= val <= 511: return f"v{val - 256}"
return "lit" if val == 255 else f"?{val}"
# Instruction base class
class Inst:
_fields: dict[str, BitField]
_encoding: tuple[BitField, int] | None = None
_defaults: dict[str, int] = {}
_values: dict[str, int | RawImm]
_words: int # size in 32-bit words, set by decode_program
_literal: int | None
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
def _or_field(self, name: str, bit: int):
cur = self._values.get(name, 0)
self._values[name] = (cur.val if isinstance(cur, RawImm) else cur) | bit
def _encode_src(self, name: str, val):
"""Encode a source field, handling modifiers and literals."""
encoded = encode_src(val)
has_opsel = 'opsel' in self._fields
if isinstance(val, Reg) and val.hi and not has_opsel: encoded |= 0x80 # hi bit in src for VOP1/2/C
self._values[name] = RawImm(encoded)
# Handle neg/abs/opsel modifiers
if isinstance(val, SrcMod):
mod_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
if val.neg and 'neg' in self._fields: self._or_field('neg', mod_bit)
if val.abs_ and 'abs' in self._fields: self._or_field('abs', mod_bit)
if isinstance(val, Reg) and val.hi and has_opsel:
self._or_field('opsel', {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0))
# Track literal value if needed
if encoded == 255 and self._literal is None:
import struct
# Check if THIS source uses 64-bit encoding (not just src0)
src_idx = {'src0': 0, 'src1': 1, 'src2': 2, 'ssrc0': 0, 'ssrc1': 1}.get(name, 0)
src_regs = self.src_regs(src_idx)
is_64 = src_regs == 2
if isinstance(val, SrcMod) and not isinstance(val, Reg): lit32 = val.val & MASK32
elif isinstance(val, int) and not isinstance(val, IntEnum): lit32 = val & MASK32
elif isinstance(val, float): lit32 = (_i64(val) >> 32) if is_64 else _i32(val) # f64: high 32 bits of f64 repr
else: return
self._literal = (lit32 << 32) if is_64 else lit32
def _encode_raw(self, name: str, val):
"""Encode a raw register field (vdst, vdata, etc.)."""
if isinstance(val, Reg):
encoded = _encode_reg(val)
if val.hi and 'opsel' not in self._fields: encoded |= 0x80
self._values[name] = encoded
if name == 'vdst' and val.hi and 'opsel' in self._fields: self._or_field('opsel', 8)
elif hasattr(val, 'value'): self._values[name] = val.value
def _validate(self, orig_args: dict):
"""Format-specific validation. Override in subclass or check by class name."""
cls_name, op = self.__class__.__name__, orig_args.get('op')
if hasattr(op, 'value'): op = op.value
# SMEM: register count must match opcode
if cls_name == 'SMEM' and op is not None:
expected = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op)
sdata = orig_args.get('sdata')
if expected and isinstance(sdata, Reg) and sdata.count != expected:
raise ValueError(f"SMEM op {op} expects {expected} registers, got {sdata.count}")
# SOP1: b32=1 reg, b64=2 regs
if cls_name == 'SOP1' and hasattr(orig_args.get('op'), 'name'):
expected = 2 if orig_args['op'].name.endswith('_B64') else 1
for fld in ('sdst', 'ssrc0'):
if isinstance(orig_args.get(fld), Reg) and orig_args[fld].count != expected:
raise ValueError(f"SOP1 {orig_args['op'].name} expects {expected} register(s) for {fld}, got {orig_args[fld].count}")
def __init__(self, *args, literal: int | None = None, **kwargs):
self._values, self._literal = dict(self._defaults), None
field_names = [n for n in self._fields if n != 'encoding']
orig_args = dict(zip(field_names, args)) | kwargs
self._values.update(orig_args)
self._validate(orig_args)
# Pre-shift literal for 64-bit sources (literal param is always raw 32-bit value from user)
if literal is not None:
# Find which source uses the literal (255) and check its register count
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
v = orig_args.get(n)
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
self._literal = (literal << 32) if self.src_regs(idx) == 2 else literal
break
else:
self._literal = literal # fallback if no literal source found
cls_name = self.__class__.__name__
# Format-specific setup
if cls_name == 'FLAT' and 'sve' in self._fields:
seg = self._values.get('seg', 0)
if (seg.val if isinstance(seg, RawImm) else seg) == 1 and isinstance(orig_args.get('addr'), VGPR): self._values['sve'] = 1
if cls_name == 'VOP3P':
op = orig_args.get('op')
if hasattr(op, 'value'): op = op.value
if op in (32, 33, 34) and 'opsel_hi' not in orig_args: self._values['opsel_hi'] = self._values['opsel_hi2'] = 0
# Encode all fields
for name, val in list(self._values.items()):
if name == 'encoding': continue
if isinstance(val, RawImm):
if name in RAW_FIELDS: self._values[name] = val.val
continue
field = self._fields.get(name)
marker = field.marker if field else None
# Type validation
if marker is _SGPRField and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
if marker is _VGPRField and not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
# Encode by field type
if name in SRC_FIELDS: self._encode_src(name, val)
elif name in RAW_FIELDS: self._encode_raw(name, val)
elif name == 'sbase': self._values[name] = (val.idx if isinstance(val, Reg) else val.val if isinstance(val, SrcMod) else val * 2) // 2
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg): self._values[name] = val.idx // 4
elif marker is _VDSTYEnc and isinstance(val, VGPR): self._values[name] = val.idx >> 1
def _encode_field(self, name: str, val) -> int:
if isinstance(val, RawImm): return val.val
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special regs like VCC_LO
if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val.val // 2 if isinstance(val, SrcMod) else val
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
return val.value if hasattr(val, 'value') else val
def to_int(self) -> int:
word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
for n, bf in self._fields.items():
if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
return word
def _get_literal(self) -> int | None:
for n in SRC_FIELDS:
if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
return None
def _is_64bit_op(self) -> bool:
"""Check if this instruction uses 64-bit operands (and thus 64-bit literals)."""
op = self._values.get('op')
if op is None: return False
op_name = op.name if hasattr(op, 'name') else None
# Look up op name from int if needed (happens in from_bytes path)
if op_name is None and self.__class__.__name__ == 'VOP3':
try: op_name = VOP3Op(op).name
except ValueError: pass
if op_name is None and self.__class__.__name__ == 'VOPC':
try: op_name = VOPCOp(op).name
except ValueError: pass
if op_name is None: return False
# V_LDEXP_F64 has 32-bit integer src1, so literal is 32-bit
return op_name != 'V_LDEXP_F64' and op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
def to_bytes(self) -> bytes:
result = self.to_int().to_bytes(self._size(), 'little')
lit = self._get_literal() or getattr(self, '_literal', None)
if lit is None: return result
# For 64-bit sources, literal is stored in high 32 bits internally, but encoded as 4 bytes
# Find which source uses the literal (255) and check its register count
lit_src_is_64 = False
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
if n not in self._values: continue
v = self._values[n]
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
lit_src_is_64 = self.is_src_64(idx)
break
lit32 = (lit >> 32) if lit_src_is_64 else lit
return result + (lit32 & MASK32).to_bytes(4, 'little')
@classmethod
def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
def size(self) -> int:
# Literal is always 4 bytes in the binary (for 64-bit ops, it's in high 32 bits)
return self._size() + (4 if self._literal is not None else 0)
@classmethod
def from_int(cls, word: int):
inst = object.__new__(cls)
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
inst._literal = None
return inst
@classmethod
def from_bytes(cls, data: bytes):
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
op_val = inst._values.get('op', 0)
has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
# VOPD fmaak/fmamk always have a literal (opx/opy value 1 or 2)
opx, opy = inst._values.get('opx', 0), inst._values.get('opy', 0)
has_literal = has_literal or (cls.__name__ == 'VOPD' and (opx in (1, 2) or opy in (1, 2)))
for n in SRC_FIELDS:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
if has_literal:
# For 64-bit ops, the literal is 32 bits placed in the HIGH 32 bits of the 64-bit value
# (low 32 bits are zero). This is how AMD hardware interprets 32-bit literals for 64-bit ops.
# Check which source uses the literal and whether THAT source is 64-bit
if len(data) >= cls._size() + 4:
lit32 = int.from_bytes(data[cls._size():cls._size()+4], 'little')
# Find which source has literal (255) and check its register count
lit_src_is_64 = False
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2)]:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255:
lit_src_is_64 = inst.src_regs(idx) == 2
break
inst._literal = (lit32 << 32) if lit_src_is_64 else lit32
return inst
def __repr__(self):
# Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
and not (is_zero(self._values[k]) and k not in {'op'})]
lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
def __getattr__(self, name: str):
if name.startswith('_'): raise AttributeError(name)
return unwrap(self._values.get(name, 0))
def lit(self, v: int, neg: bool = False) -> str:
s = f"0x{self._literal:x}" if v == 255 and self._literal else decode_src(v)
return f"-{s}" if neg else s
def __eq__(self, other):
if not isinstance(other, Inst): return NotImplemented
return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
def disasm(self) -> str:
from extra.assembly.amd.asm import disasm
return disasm(self)
_enum_map = {'VOP1': VOP1Op, 'VOP2': VOP2Op, 'VOP3': VOP3Op, 'VOP3SD': VOP3SDOp, 'VOP3P': VOP3POp, 'VOPC': VOPCOp,
'SOP1': SOP1Op, 'SOP2': SOP2Op, 'SOPC': SOPCOp, 'SOPK': SOPKOp, 'SOPP': SOPPOp,
'SMEM': SMEMOp, 'DS': DSOp, 'FLAT': FLATOp, 'MUBUF': MUBUFOp, 'MTBUF': MTBUFOp, 'MIMG': MIMGOp,
'VOPD': VOPDOp, 'VINTERP': VINTERPOp}
_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
@property
def op(self):
"""Return the op as an enum (e.g., VOP1Op.V_MOV_B32). VOP3 returns VOPCOp/VOP3SDOp for those op ranges."""
val = self._values.get('op')
if val is None: return None
if hasattr(val, 'name'): return val # already an enum
cls_name = self.__class__.__name__
assert cls_name in self._enum_map, f"no enum map for {cls_name}"
return self._enum_map[cls_name](val)
@cached_property
def op_name(self) -> str:
op = self.op
return op.name if hasattr(op, 'name') else ''
@cached_property
def _spec_regs(self) -> tuple[int, int, int, int]: return spec_regs(self.op_name)
@cached_property
def _spec_dtype(self) -> tuple[str | None, str | None, str | None, str | None]: return spec_dtype(self.op_name)
def dst_regs(self) -> int: return self._spec_regs[0]
def src_regs(self, n: int) -> int: return self._spec_regs[n + 1]
def num_srcs(self) -> int: return spec_num_srcs(self.op_name)
def dst_dtype(self) -> str | None: return self._spec_dtype[0]
def src_dtype(self, n: int) -> str | None: return self._spec_dtype[n + 1]
def is_src_16(self, n: int) -> bool: return self._spec_regs[n + 1] == 1 and is_dtype_16(self._spec_dtype[n + 1])
def is_src_64(self, n: int) -> bool: return self._spec_regs[n + 1] == 2
def is_16bit(self) -> bool: return spec_is_16bit(self.op_name)
def is_64bit(self) -> bool: return spec_is_64bit(self.op_name)
def is_dst_16(self) -> bool: return self._spec_regs[0] == 1 and is_dtype_16(self._spec_dtype[0])
class Inst32(Inst): pass
class Inst64(Inst): pass
+462
View File
@@ -0,0 +1,462 @@
# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
# mypy: ignore-errors
from __future__ import annotations
import ctypes
from extra.assembly.amd.dsl import Inst, unwrap, FLOAT_ENC, MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
from extra.assembly.amd.pcode import Reg
from extra.assembly.amd.asm import detect_format
from extra.assembly.amd.autogen.rdna3.gen_pcode import get_compiled_functions
from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD,
SrcEnum, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, GLOBALOp, FLATOp, DSOp, VOPDOp)
Program = dict[int, Inst]
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
VCC_LO, VCC_HI, NULL, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535} # INV_2PI
def _build_inline_consts(mask, to_bits):
tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
return tbl
_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
# Helper: extract/write 16-bit half from/to 32-bit value
def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
def _vgpr_hi(src: int) -> bool: return src >= 256 and ((src - 256) & 0x80) != 0
def _vgpr_masked(src: int) -> int: return ((src - 256) & 0x7f) + 256 if src >= 256 else src
# 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_uint32).from_address(addr)
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
def mem_write(addr: int, size: int, val: int) -> None:
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
# Memory op tables (not pseudocode - these are format descriptions)
def _mem_ops(ops, suffix_map):
return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
FLAT_LOAD, FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP), _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
# D16 ops: load/store 16-bit to lower or upper half of VGPR. Format: (size, sign, hi) where hi=1 means upper 16 bits
_D16_LOAD_MAP = {'LOAD_D16_U8': (1,0,0), 'LOAD_D16_I8': (1,1,0), 'LOAD_D16_B16': (2,0,0),
'LOAD_D16_HI_U8': (1,0,1), 'LOAD_D16_HI_I8': (1,1,1), 'LOAD_D16_HI_B16': (2,0,1)}
_D16_STORE_MAP = {'STORE_D16_HI_B8': (1,1), 'STORE_D16_HI_B16': (2,1)} # (size, hi)
FLAT_D16_LOAD = _mem_ops([GLOBALOp, FLATOp], _D16_LOAD_MAP)
FLAT_D16_STORE = _mem_ops([GLOBALOp, FLATOp], _D16_STORE_MAP)
SMEM_LOAD = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
_VOPD_TO_VOP = {
VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32,
VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32,
VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32,
VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32,
VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32,
VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32,
}
# Compiled pseudocode functions (lazy loaded)
_COMPILED: dict | None = None
def _get_compiled() -> dict:
global _COMPILED
if _COMPILED is None: _COMPILED = get_compiled_functions()
return _COMPILED
class WaveState:
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
def __init__(self):
self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
self.sgpr[EXEC_LO], self.scc, self.pc, self.literal, self._pend_sgpr = 0xffffffff, 0, 0, 0, {}
@property
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
@vcc.setter
def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
@property
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
@exec_mask.setter
def exec_mask(self, v: int): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
def rsgpr(self, i: int) -> int: return 0 if i == NULL else self.scc if i == SCC else self.sgpr[i] if i < SGPR_COUNT else 0
def wsgpr(self, i: int, v: int):
if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & MASK32
def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
def wsgpr64(self, i: int, v: int): self.wsgpr(i, v & MASK32); self.wsgpr(i+1, (v >> 32) & MASK32)
def _rsrc_base(self, v: int, lane: int, consts):
if v < SGPR_COUNT: return self.sgpr[v]
if v == SCC: return self.scc
if v < 255: return consts[v - 128]
if v == 255: return self.literal
return self.vgpr[lane][v - 256] if v <= 511 else 0
def rsrc(self, v: int, lane: int) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS)
def rsrc_f16(self, v: int, lane: int) -> int: return self._rsrc_base(v, lane, _INLINE_CONSTS_F16)
def rsrc64(self, v: int, lane: int) -> int:
if 128 <= v < 255: return _INLINE_CONSTS_F64[v - 128]
if v == 255: return self.literal # literal is already shifted in from_bytes for 64-bit ops
return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
def pend_sgpr_lane(self, reg: int, lane: int, val: int):
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()
def decode_program(data: bytes) -> Program:
result: Program = {}
i = 0
while i < len(data):
try: inst_class = detect_format(data[i:])
except ValueError: break # stop at invalid instruction (padding/metadata after code)
if inst_class is None: i += 4; continue
base_size = inst_class._size()
# Pass enough data for potential 64-bit literal (base + 8 bytes max)
inst = inst_class.from_bytes(data[i:i+base_size+8])
for name, val in inst._values.items():
if name != 'op': setattr(inst, name, unwrap(val)) # skip op to preserve property access
inst._words = inst.size() // 4
result[i // 4] = inst
i += inst._words * 4
return result
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION - All ALU ops use pseudocode from PDF
# ═══════════════════════════════════════════════════════════════════════════════
def exec_scalar(st: WaveState, inst: Inst) -> int:
"""Execute scalar instruction. Returns PC delta or negative for special cases."""
compiled = _get_compiled()
# SOPP: special cases for control flow that has no pseudocode
if isinstance(inst, SOPP):
if inst.op == SOPPOp.S_ENDPGM: return -1
if inst.op == SOPPOp.S_BARRIER: return -2
# SMEM: memory loads (not ALU)
if isinstance(inst, SMEM):
addr = st.rsgpr64(inst.sbase * 2) + _sext(inst.offset, 21)
if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & MASK64, 4))
return 0
# Get op enum and lookup compiled function
if 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)}")
# SOPP has gaps in the opcode enum - treat unknown opcodes as no-ops
try: op = inst.op
except ValueError:
if isinstance(inst, SOPP): return 0
raise
fn = compiled.get(type(op), {}).get(op)
if fn is None:
# SOPP instructions without pseudocode (waits, hints, nops) are no-ops
if isinstance(inst, SOPP): return 0
raise NotImplementedError(f"{op.name} not in pseudocode")
# Build context - use inst methods to determine operand sizes
s0 = st.rsrc64(ssrc0, 0) if inst.is_src_64(0) else (st.rsrc(ssrc0, 0) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
s1 = st.rsrc64(inst.ssrc1, 0) if inst.is_src_64(1) else (st.rsrc(inst.ssrc1, 0) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
d0 = st.rsgpr64(sdst) if inst.dst_regs() == 2 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else st.literal
# Create Reg objects for compiled function - mask VCC/EXEC to 32 bits for wave32
result = fn(Reg(s0), Reg(s1), None, Reg(d0), Reg(st.scc), Reg(st.vcc & MASK32), 0, Reg(st.exec_mask & MASK32), literal, None, PC=Reg(st.pc * 4))
# Apply results - extract values from returned Reg objects
if sdst is not None and 'D0' in result:
(st.wsgpr64 if inst.dst_regs() == 2 else st.wsgpr)(sdst, result['D0']._val)
if 'SCC' in result: st.scc = result['SCC']._val & 1
if 'EXEC' in result: st.exec_mask = result['EXEC']._val
if 'PC' in result:
# Convert absolute byte address to word delta
pc_val = result['PC']._val
new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
new_pc_words = new_pc // 4
return new_pc_words - st.pc - 1 # -1 because emulator adds inst_words (1 for scalar)
return 0
def exec_vector(st: WaveState, inst: Inst, lane: int, lds: bytearray | None = None) -> None:
"""Execute vector instruction for one lane."""
compiled = _get_compiled()
V = st.vgpr[lane]
# Memory ops (not ALU pseudocode)
if isinstance(inst, FLAT):
op, addr_reg, data_reg, vdst, offset, saddr = inst.op, inst.addr, inst.data, inst.vdst, _sext(inst.offset, 13), inst.saddr
addr = V[addr_reg] | (V[addr_reg+1] << 32)
addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & MASK64 if saddr not in (NULL, 0x7f) else (addr + offset) & MASK64
if op in FLAT_LOAD:
cnt, sz, sign = FLAT_LOAD[op]
for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = _sext(val, sz * 8) & MASK32 if sign else val
elif op in FLAT_STORE:
cnt, sz = FLAT_STORE[op]
for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
elif op in FLAT_D16_LOAD:
sz, sign, hi = FLAT_D16_LOAD[op]
val = mem_read(addr, sz)
if sign: val = _sext(val, sz * 8) & 0xffff
V[vdst] = _dst16(V[vdst], val, hi)
elif op in FLAT_D16_STORE:
sz, hi = FLAT_D16_STORE[op]
mem_write(addr, sz, _src16(V[data_reg], hi) & ((1 << (sz * 8)) - 1))
else: raise NotImplementedError(f"FLAT op {op}")
return
if isinstance(inst, DS):
fn = compiled.get(DSOp, {}).get(inst.op)
if fn is None: raise NotImplementedError(f"DS op {inst.op.name} not in pseudocode")
# Prepare data registers as lists of dwords
data0 = [V[inst.data0 + i] for i in range(4)] # up to 4 dwords
data1 = [V[inst.data1 + i] for i in range(4)] if inst.data1 else [0, 0, 0, 0]
result = fn(lds, V[inst.addr], data0, data1, inst.vdst, inst.offset0, inst.offset1)
# Write results for loads
if 'vdst' in result:
for i, val in enumerate(result['vdst']): V[inst.vdst + i] = val & MASK32
return
# VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)
if isinstance(inst, VOPD):
vdsty = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1)
inputs = [(inst.opx, st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], V[inst.vdstx], inst.vdstx),
(inst.opy, st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], V[vdsty], vdsty)]
def exec_vopd(vopd_op, s0, s1, d0):
op = _VOPD_TO_VOP[vopd_op]
return compiled[type(op)][op](Reg(s0), Reg(s1), None, Reg(d0), Reg(st.scc), Reg(st.vcc), lane, Reg(st.exec_mask), st.literal, None)['D0']._val
for vopd_op, s0, s1, d0, dst in inputs: V[dst] = exec_vopd(vopd_op, s0, s1, d0)
return
# VOP3SD: has extra scalar dest for carry output
if isinstance(inst, VOP3SD):
fn = compiled[VOP3SDOp][inst.op]
# Read sources based on register counts from inst properties
def rsrc_n(src, regs): return st.rsrc64(src, lane) if regs == 2 else st.rsrc(src, lane)
s0, s1, s2 = rsrc_n(inst.src0, inst.src_regs(0)), rsrc_n(inst.src1, inst.src_regs(1)), rsrc_n(inst.src2, inst.src_regs(2))
# Carry-in ops use src2 as carry bitmask instead of VCC
vcc = st.rsgpr64(inst.src2) if 'CO_CI' in inst.op_name else st.vcc
result = fn(Reg(s0), Reg(s1), Reg(s2), Reg(V[inst.vdst]), Reg(st.scc), Reg(vcc), lane, Reg(st.exec_mask), st.literal, None)
d0_val = result['D0']._val
V[inst.vdst] = d0_val & MASK32
if inst.dst_regs() == 2: V[inst.vdst + 1] = (d0_val >> 32) & MASK32
if 'VCC' in result: st.pend_sgpr_lane(inst.sdst, lane, (result['VCC']._val >> lane) & 1)
return
# Get op enum and sources (None means "no source" for that operand)
# dst_hi: for VOP1/VOP2 16-bit dst ops, bit 7 of vdst indicates .h (high 16-bit) destination
dst_hi = False
if isinstance(inst, VOP1):
if inst.op == VOP1Op.V_NOP: return
src0, src1, src2 = inst.src0, None, None
dst_hi = (inst.vdst & 0x80) != 0 and inst.is_dst_16()
vdst = inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
elif isinstance(inst, VOP2):
src0, src1, src2 = inst.src0, inst.vsrc1 + 256, None
dst_hi = (inst.vdst & 0x80) != 0 and inst.is_dst_16()
vdst = inst.vdst & 0x7f if inst.is_dst_16() else inst.vdst
elif isinstance(inst, VOP3):
# VOP3 ops 0-255 are VOPC comparisons encoded as VOP3 - inst.op returns VOPCOp for these
src0, src1, src2, vdst = inst.src0, inst.src1, (None if inst.op.value < 256 else inst.src2), inst.vdst
elif isinstance(inst, VOPC):
# For 16-bit VOPC, vsrc1 uses same encoding as VOP2 16-bit: bit 7 selects hi(1) or lo(0) half
# vsrc1 field is 8 bits: [6:0] = VGPR index, [7] = hi flag
src0, src1, src2, vdst = inst.src0, inst.vsrc1 + 256, None, VCC_LO
elif isinstance(inst, VOP3P):
# VOP3P: Packed 16-bit operations using compiled functions
# WMMA: wave-level matrix multiply-accumulate (special handling - needs cross-lane access)
if 'WMMA' in inst.op_name:
if lane == 0: # Only execute once per wave, write results for all lanes
exec_wmma(st, inst, inst.op)
return
# V_FMA_MIX: Mixed precision FMA - opsel_hi controls f32(0) vs f16(1), opsel selects which f16 half
# Handle inline because abs/neg must be applied AFTER type conversion
if inst.op in (VOP3POp.V_FMA_MIX_F32, VOP3POp.V_FMA_MIXLO_F16, VOP3POp.V_FMA_MIXHI_F16):
opsel, opsel_hi, opsel_hi2 = getattr(inst, 'opsel', 0), getattr(inst, 'opsel_hi', 0), getattr(inst, 'opsel_hi2', 0)
neg, abs_ = getattr(inst, 'neg', 0), getattr(inst, 'neg_hi', 0) # neg_hi reused as abs for FMA_MIX
raws = [st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane) if inst.src2 is not None else 0]
is_f16 = [opsel_hi & 1, opsel_hi & 2, opsel_hi2]
srcs = [_f16(_src16(raws[i], bool(opsel & (1<<i)))) if is_f16[i] else _f32(raws[i]) for i in range(3)]
for i in range(3):
if abs_ & (1<<i): srcs[i] = abs(srcs[i])
if neg & (1<<i): srcs[i] = -srcs[i]
result_f = srcs[0] * srcs[1] + srcs[2]
V = st.vgpr[lane]
V[inst.vdst] = _i32(result_f) if inst.op == VOP3POp.V_FMA_MIX_F32 else _dst16(V[inst.vdst], _i16(result_f), inst.op == VOP3POp.V_FMA_MIXHI_F16)
return
# VOP3P packed ops: opsel selects halves for lo, opsel_hi for hi; neg toggles f16 sign
raws = [st.rsrc_f16(inst.src0, lane), st.rsrc_f16(inst.src1, lane), st.rsrc_f16(inst.src2, lane) if inst.src2 is not None else 0]
opsel, opsel_hi, opsel_hi2 = getattr(inst, 'opsel', 0), getattr(inst, 'opsel_hi', 3), getattr(inst, 'opsel_hi2', 1)
neg, neg_hi = getattr(inst, 'neg', 0), getattr(inst, 'neg_hi', 0)
hi_sels = [opsel_hi & 1, opsel_hi & 2, opsel_hi2]
srcs = [((_src16(raws[i], hi_sels[i]) ^ (0x8000 if neg_hi & (1<<i) else 0)) << 16) |
(_src16(raws[i], opsel & (1<<i)) ^ (0x8000 if neg & (1<<i) else 0)) for i in range(3)]
result = compiled[VOP3POp][inst.op](Reg(srcs[0]), Reg(srcs[1]), Reg(srcs[2]), Reg(0), Reg(st.scc), Reg(st.vcc), lane, Reg(st.exec_mask), st.literal, None)
st.vgpr[lane][inst.vdst] = result['D0']._val & MASK32
return
else: raise NotImplementedError(f"Unknown vector type {type(inst)}")
op_cls = type(inst.op)
if (fn := compiled.get(op_cls, {}).get(inst.op)) is None: raise NotImplementedError(f"{inst.op_name} not in pseudocode")
# Read sources (with VOP3 modifiers if applicable)
neg, abs_ = (getattr(inst, 'neg', 0), getattr(inst, 'abs', 0)) if isinstance(inst, VOP3) else (0, 0)
opsel = getattr(inst, 'opsel', 0) if isinstance(inst, VOP3) else 0
def mod_src(val: int, idx: int, is64=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
# Use inst methods to determine operand sizes (inst.is_src_16, inst.is_src_64, etc.)
is_vop2_16bit = isinstance(inst, VOP2) and inst.is_16bit()
# Read sources based on register counts and dtypes from inst properties
def read_src(src, idx, regs, is_src_16):
if src is None: return 0
if regs == 2: return mod_src(st.rsrc64(src, lane), idx, is64=True)
if is_src_16 and isinstance(inst, VOP3):
raw = st.rsrc_f16(src, lane) if 128 <= src < 255 else st.rsrc(src, lane)
val = _src16(raw, bool(opsel & (1 << idx)))
if abs_ & (1 << idx): val &= 0x7fff
if neg & (1 << idx): val ^= 0x8000
return val
if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
if src >= 256: return _src16(mod_src(st.rsrc(_vgpr_masked(src), lane), idx), _vgpr_hi(src))
return mod_src(st.rsrc_f16(src, lane), idx) & 0xffff
return mod_src(st.rsrc(src, lane), idx)
s0 = read_src(src0, 0, inst.src_regs(0), inst.is_src_16(0))
s1 = read_src(src1, 1, inst.src_regs(1), inst.is_src_16(1)) if src1 is not None else 0
s2 = read_src(src2, 2, inst.src_regs(2), inst.is_src_16(2)) if src2 is not None else 0
# Read destination (accumulator for VOP2 f16, 64-bit for 64-bit ops)
d0 = _src16(V[vdst], dst_hi) if is_vop2_16bit else (V[vdst] | (V[vdst + 1] << 32)) if inst.dst_regs() == 2 else V[vdst]
# V_CNDMASK_B32/B16: VOP3 encoding uses src2 as mask (not VCC); VOP2 uses VCC implicitly
# Pass the correct mask as vcc to the function so pseudocode VCC.u64[laneId] works correctly
vcc_for_fn = st.rsgpr64(src2) if inst.op in (VOP3Op.V_CNDMASK_B32, VOP3Op.V_CNDMASK_B16) and isinstance(inst, VOP3) and src2 is not None and src2 < 256 else st.vcc
# Execute compiled function - pass src0_idx and vdst_idx for lane instructions
# For VGPR access: src0 index is the VGPR number (src0 - 256 if VGPR, else src0 for SGPR)
src0_idx = (src0 - 256) if src0 is not None and src0 >= 256 else (src0 if src0 is not None else 0)
result = fn(Reg(s0), Reg(s1), Reg(s2), Reg(d0), Reg(st.scc), Reg(vcc_for_fn), lane, Reg(st.exec_mask), st.literal, st.vgpr, src0_idx, vdst)
# Apply results - extract values from returned Reg objects
if 'vgpr_write' in result:
# Lane instruction wrote to VGPR: (lane, vgpr_idx, value)
wr_lane, wr_idx, wr_val = result['vgpr_write']
st.vgpr[wr_lane][wr_idx] = wr_val
if 'VCC' in result:
# VOP2 carry ops write to VCC implicitly; VOPC/VOP3 write to vdst
st.pend_sgpr_lane(VCC_LO if isinstance(inst, VOP2) and 'CO_CI' in inst.op_name else vdst, lane, (result['VCC']._val >> lane) & 1)
if 'EXEC' in result:
# V_CMPX instructions write to EXEC per-lane (not to vdst)
st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC']._val >> lane) & 1)
elif op_cls is VOPCOp:
# VOPC comparison result stored in D0 bitmask, extract lane bit (non-CMPX only)
st.pend_sgpr_lane(vdst, lane, (result['D0']._val >> lane) & 1)
if op_cls is not VOPCOp and 'vgpr_write' not in result:
writes_to_sgpr = 'READFIRSTLANE' in inst.op_name or 'READLANE' in inst.op_name
d0_val = result['D0']._val
if writes_to_sgpr: st.wsgpr(vdst, d0_val & MASK32)
elif inst.dst_regs() == 2: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
elif inst.is_dst_16(): V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
else: V[vdst] = d0_val & MASK32
# ═══════════════════════════════════════════════════════════════════════════════
# WMMA (Wave Matrix Multiply-Accumulate)
# ═══════════════════════════════════════════════════════════════════════════════
def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
"""Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
src0, src1, src2, vdst = inst.src0, inst.src1, inst.src2, inst.vdst
# Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
def read_f16_mat(src):
return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
# Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
# Compute D = A × B + C (16x16 matrix multiply)
mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
# Write result - f16 packed or f32
if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
for i in range(0, 256, 2):
st.vgpr[(i//2) % 32][vdst + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
else:
for i in range(256): st.vgpr[i % 32][vdst + i//32] = _i32(mat_d[i])
# ═══════════════════════════════════════════════════════════════════════════════
# MAIN EXECUTION LOOP
# ═══════════════════════════════════════════════════════════════════════════════
def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
inst = program.get(st.pc)
if inst is None: return 1
inst_words, st.literal = inst._words, getattr(inst, '_literal', None) or 0
if isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)):
delta = exec_scalar(st, inst)
if delta == -1: return -1 # endpgm
if delta == -2: st.pc += inst_words; return -2 # barrier
st.pc += inst_words + delta
else:
# V_READFIRSTLANE/V_READLANE write to SGPR, execute once; others execute per-lane with exec_mask
is_readlane = isinstance(inst, (VOP1, VOP3)) and ('READFIRSTLANE' in inst.op_name or 'READLANE' in inst.op_name)
exec_mask = 1 if is_readlane else st.exec_mask
for lane in range(1 if is_readlane else n_lanes):
if exec_mask & (1 << lane): exec_vector(st, inst, lane, lds)
st.commit_pends()
st.pc += inst_words
return 0
def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
while st.pc in program:
result = step_wave(program, st, lds, n_lanes)
if result == -1: return 0
if result == -2: return -2
return 0
def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int,
wg_id_sgpr_base: int, wg_id_enables: tuple[bool, bool, bool]) -> None:
lx, ly, lz = local_size
total_threads, lds = lx * ly * lz, bytearray(65536)
waves: list[tuple[WaveState, int, int]] = []
for wave_start in range(0, total_threads, WAVE_SIZE):
n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
st.exec_mask = (1 << n_lanes) - 1
st.wsgpr64(0, args_ptr)
# Set workgroup IDs in SGPRs based on USER_SGPR_COUNT and enable flags from COMPUTE_PGM_RSRC2
sgpr_idx = wg_id_sgpr_base
for wg_id, enabled in zip(workgroup_id, wg_id_enables):
if enabled: st.sgpr[sgpr_idx] = wg_id; sgpr_idx += 1
# Set workitem IDs in VGPR0 using packed method: v0 = (Z << 20) | (Y << 10) | X
for i in range(n_lanes):
tid = wave_start + i
st.vgpr[i][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
waves.append((st, n_lanes, wave_start))
has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
for _ in range(2 if has_barrier else 1):
for st, n_lanes, _ in waves: exec_wave(program, st, lds, n_lanes)
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)
if not program: return -1
wg_id_enables = tuple(bool((rsrc2 >> (7+i)) & 1) for i in range(3))
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 >> 1) & 0x1f, wg_id_enables)
return 0
+542
View File
@@ -0,0 +1,542 @@
# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
import struct, math
from extra.assembly.amd.dsl import MASK32, MASK64, _f32, _i32, _sext, _f16, _i16, _f64, _i64
# ═══════════════════════════════════════════════════════════════════════════════
# HELPER FUNCTIONS
# ═══════════════════════════════════════════════════════════════════════════════
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 _to_f16_bits(v): return v if isinstance(v, int) else _i16(v)
def _isnan(x):
try: return math.isnan(float(x))
except (TypeError, ValueError): return False
def _check_nan_type(x, quiet_bit_expected, default):
"""Check NaN type by examining quiet bit. Returns default if can't determine."""
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)
# NaN format: exponent all 1s, quiet bit, mantissa != 0
# f16: exp[14:10]=31, quiet=bit9, mant[8:0] | f32: exp[30:23]=255, quiet=bit22, mant[22:0] | f64: exp[62:52]=2047, quiet=bit51, mant[51:0]
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 _isquietnan(x): return _check_nan_type(x, 1, True) # quiet NaN has quiet bit = 1
def _issignalnan(x): return _check_nan_type(x, 0, False) # signaling NaN has quiet bit = 0
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 _fma(a, b, c): return a * b + c
def _signext(v): return v
def _fpop(fn): return lambda x: (x := float(x), x if math.isnan(x) or math.isinf(x) else float(fn(x)))[1]
trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
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))
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"))
i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
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 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_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
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 u8_to_u32(v): return int(v) & 0xff
def u4_to_u32(v): return int(v) & 0xf
def _sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
def _mantissa_f32(f): return struct.unpack("<I", struct.pack("<f", f))[0] & 0x7fffff if not (math.isinf(f) or math.isnan(f)) else 0
def _ldexp(m, e): return math.ldexp(m, e)
def isEven(x):
x = float(x)
if math.isinf(x) or math.isnan(x): return False
return int(x) % 2 == 0
def fract(x): return x - math.floor(x)
PI = math.pi
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))
return fn(frac_cycles * 2 * math.pi)
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 _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
def _brev32(v): return _brev(v, 32)
def _brev64(v): return _brev(v, 64)
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 _ctz32(v): return _ctz(v, 32)
def _ctz64(v): return _ctz(v, 64)
def _exponent(f):
# Handle TypedView (f16/f32/f64) to get correct exponent for that type
if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
raw = f._val
if f._bits == 16: return (raw >> 10) & 0x1f # f16: 5-bit exponent
if f._bits == 32: return (raw >> 23) & 0xff # f32: 8-bit exponent
if f._bits == 64: return (raw >> 52) & 0x7ff # f64: 11-bit exponent
# Fallback: convert to f32 and get exponent
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 _is_denorm_f32(f):
if not isinstance(f, float): f = _f32(int(f) & 0xffffffff)
if math.isinf(f) or math.isnan(f) or f == 0.0: return False
bits = struct.unpack("<I", struct.pack("<f", float(f)))[0]
return (bits >> 23) & 0xff == 0
def _is_denorm_f64(f):
if not isinstance(f, float): f = _f64(int(f) & 0xffffffffffffffff)
if math.isinf(f) or math.isnan(f) or f == 0.0: return False
bits = struct.unpack("<Q", struct.pack("<d", float(f)))[0]
return (bits >> 52) & 0x7ff == 0
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)
def ABSDIFF(a, b): return abs(int(a) - int(b))
# BF16 (bfloat16) conversion functions
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 bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
def f32_to_bf16(f): return _ibf16(f)
# BYTE_PERMUTE for V_PERM_B32 - select bytes from 64-bit data based on selector
def BYTE_PERMUTE(data, sel):
"""Select a byte from 64-bit data based on selector value.
sel 0-7: select byte from data (S1 is bytes 0-3, S0 is bytes 4-7 in {S0,S1})
sel 8-11: sign-extend from specific bytes (8->byte1, 9->byte3, 10->byte5, 11->byte7)
sel 12: constant 0x00
sel >= 13: constant 0xFF"""
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 # sign of byte 1
if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00 # sign of byte 3
if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00 # sign of byte 5
if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00 # sign of byte 7
if sel == 12: return 0x00
return 0xff # sel >= 13
# v_sad_u8 helper for V_SAD instructions (sum of absolute differences of 4 bytes)
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
# v_msad_u8 helper (masked SAD - skip when reference byte is 0)
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 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 u32_to_u16(u): return int(u) & 0xffff
def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
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
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
# ═══════════════════════════════════════════════════════════════════════════════
# DSL EXPORTS
# ═══════════════════════════════════════════════════════════════════════════════
__all__ = [
# Classes
'Reg', 'SliceProxy', 'TypedView',
# Pack functions
'_pack', '_pack32', 'pack', 'pack32',
# Constants
'WAVE32', 'WAVE64', 'MASK32', 'MASK64', 'WAVE_MODE', 'DENORM', 'OVERFLOW_F32', 'UNDERFLOW_F32',
'OVERFLOW_F64', 'UNDERFLOW_F64', 'MAX_FLOAT_F32', 'ROUND_MODE', 'cvtToQuietNAN', 'DST', 'INF', 'PI',
'TWO_OVER_PI_1201',
# Aliases for pseudocode
's_ff1_i32_b32', 's_ff1_i32_b64', 'GT_NEG_ZERO', 'LT_NEG_ZERO',
'isNAN', 'isQuietNAN', 'isSignalNAN', 'fma', 'ldexp', 'sign', 'exponent', 'F', 'signext',
# Conversion functions
'_f32', '_i32', '_f16', '_i16', '_f64', '_i64', '_sext', '_to_f16_bits', '_f16_to_f32_bits',
'i32_to_f32', 'u32_to_f32', 'i32_to_f64', 'u32_to_f64', 'f32_to_f64', 'f64_to_f32',
'f32_to_i32', 'f32_to_u32', 'f64_to_i32', 'f64_to_u32', 'f32_to_f16', 'f16_to_f32',
'i16_to_f16', 'u16_to_f16', 'f16_to_i16', 'f16_to_u16', 'u32_to_u16', 'i32_to_i16',
'f16_to_snorm', 'f16_to_unorm', 'f32_to_snorm', 'f32_to_unorm', 'v_cvt_i16_f32', 'v_cvt_u16_f32',
'SAT8', 'f32_to_u8', 'u8_to_u32', 'u4_to_u32',
# BF16 conversion functions
'_bf16', '_ibf16', 'bf16_to_f32', 'f32_to_bf16',
# Math functions
'trunc', 'floor', 'ceil', 'sqrt', 'log2', 'sin', 'cos', 'pow', 'fract', 'isEven', 'mantissa',
# Min/max functions
'v_min_f32', 'v_max_f32', 'v_min_i32', 'v_max_i32', 'v_min_u32', 'v_max_u32',
'v_min_f16', 'v_max_f16', 'v_min_i16', 'v_max_i16', 'v_min_u16', 'v_max_u16',
'v_min3_f32', 'v_max3_f32', 'v_min3_i32', 'v_max3_i32', 'v_min3_u32', 'v_max3_u32',
'v_min3_f16', 'v_max3_f16', 'v_min3_i16', 'v_max3_i16', 'v_min3_u16', 'v_max3_u16',
'ABSDIFF',
# Byte/SAD helper functions
'BYTE_PERMUTE', 'v_sad_u8', 'v_msad_u8',
# Bit manipulation
'_brev32', '_brev64', '_ctz32', '_ctz64', '_exponent', '_is_denorm_f32', '_is_denorm_f64',
'_sign', '_mantissa_f32', '_div', '_isnan', '_isquietnan', '_issignalnan', '_gt_neg_zero', '_lt_neg_zero', '_fma', '_ldexp', '_signext',
'signext_from_bit',
]
# Aliases used in pseudocode
s_ff1_i32_b32, s_ff1_i32_b64 = _ctz32, _ctz64
GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
isNAN = _isnan
isQuietNAN = _isquietnan
isSignalNAN = _issignalnan
fma, ldexp, sign, exponent = _fma, _ldexp, _sign, _exponent
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) # int -> interpret as f32 bits
if isinstance(x, TypedView): return x # preserve TypedView for bit-pattern checks
return float(x) # already a float or float-like
signext = lambda x: int(x) # sign-extend to full width - already handled by Python's arbitrary precision ints
pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
_pack, _pack32 = pack, pack32 # Aliases for internal use
WAVE32, WAVE64 = True, False
# Float overflow/underflow constants
OVERFLOW_F32 = float('inf')
UNDERFLOW_F32 = 0.0
OVERFLOW_F64 = float('inf')
UNDERFLOW_F64 = 0.0
MAX_FLOAT_F32 = 3.4028235e+38 # Largest finite float32
# INF object that supports .f16/.f32/.f64 access and comparison with floats
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)
INF = _Inf()
# Rounding mode placeholder
class _RoundMode:
NEAREST_EVEN = 0
ROUND_MODE = _RoundMode()
# Helper functions for pseudocode
def cvtToQuietNAN(x): return float('nan')
DST = None # Placeholder, will be set in context
# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
# Computed as: int((2/pi) * 2^1201) - this is the fractional part of 2/pi scaled to integer
# The MSB (bit 1200) corresponds to 2^0 position in the fraction 0.b1200 b1199 ... b1 b0
_TWO_OVER_PI_1201_RAW = 0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6
class _BigInt:
"""Wrapper for large integers that supports bit slicing [high:low]."""
__slots__ = ('_val',)
def __init__(self, val): self._val = val
def __getitem__(self, key):
if isinstance(key, slice):
high, low = key.start, key.stop
if high < low: high, low = low, high # Handle reversed slice
mask = (1 << (high - low + 1)) - 1
return (self._val >> low) & mask
return (self._val >> key) & 1
def __int__(self): return self._val
def __index__(self): return self._val
def __lshift__(self, n): return self._val << int(n)
def __rshift__(self, n): return self._val >> int(n)
def __and__(self, n): return self._val & int(n)
def __or__(self, n): return self._val | int(n)
TWO_OVER_PI_1201 = _BigInt(_TWO_OVER_PI_1201_RAW)
class _WaveMode:
IEEE = False
WAVE_MODE = _WaveMode()
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):
return _is_denorm_f64(float(other)) if self._bits == 64 else _is_denorm_f32(float(other))
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)
DENORM = _Denorm()
def _brev(v, bits):
"""Bit-reverse a value."""
result = 0
for i in range(bits): result |= ((v >> i) & 1) << (bits - 1 - i)
return result
class SliceProxy:
"""Proxy for D0[31:16] that supports .f16/.u16 etc getters and setters."""
__slots__ = ('_reg', '_high', '_low', '_reversed')
def __init__(self, reg, high, low):
self._reg = reg
# Handle reversed slices like [0:31] which means bit-reverse
if high < low: self._high, self._low, self._reversed = low, high, True
else: self._high, self._low, self._reversed = high, low, False
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)
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
def __int__(self): return self._get()
def __index__(self): return self._get()
# Comparison operators (compare as integers)
def __eq__(s, o): return s._get() == int(o)
def __ne__(s, o): return s._get() != int(o)
def __lt__(s, o): return s._get() < int(o)
def __le__(s, o): return s._get() <= int(o)
def __gt__(s, o): return s._get() > int(o)
def __ge__(s, o): return s._get() >= int(o)
class TypedView:
"""View for S0.u32 that supports [4:0] slicing and [bit] access."""
__slots__ = ('_reg', '_bits', '_signed', '_float', '_bf16')
def __init__(self, reg, bits, signed=False, is_float=False, is_bf16=False):
self._reg, self._bits, self._signed, self._float, self._bf16 = reg, bits, signed, is_float, is_bf16
@property
def _val(self):
mask = MASK64 if self._bits == 64 else MASK32 if self._bits == 32 else (1 << self._bits) - 1
return self._reg._val & mask
def __getitem__(self, key):
if isinstance(key, slice):
high, low = int(key.start), int(key.stop)
return SliceProxy(self._reg, high, low)
return (self._val >> 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._val, self._bits) if self._signed else self._val
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._val) # bf16 uses different conversion
return _f16(self._val) if self._bits == 16 else _f32(self._val) if self._bits == 32 else _f64(self._val)
return float(int(self))
# 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 else 0
def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 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)
def __bool__(s): return bool(int(s))
# Allow chained type access like jump_addr.i64 when jump_addr is already a TypedView
# These just return self or convert appropriately
@property
def i64(s): return s if s._bits == 64 and s._signed else int(s)
@property
def u64(s): return s if s._bits == 64 and not s._signed else int(s) & MASK64
@property
def i32(s): return s if s._bits == 32 and s._signed else _sext(int(s) & MASK32, 32)
@property
def u32(s): return s if s._bits == 32 and not s._signed else int(s) & MASK32
class Reg:
"""GPU register: D0.f32 = S0.f32 + S1.f32 just works."""
__slots__ = ('_val',)
def __init__(self, val=0): self._val = int(val) & MASK64
# Typed views
u64 = property(lambda s: TypedView(s, 64), lambda s, v: setattr(s, '_val', int(v) & MASK64))
i64 = property(lambda s: TypedView(s, 64, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
b64 = property(lambda s: TypedView(s, 64), lambda s, v: setattr(s, '_val', int(v) & MASK64))
f64 = property(lambda s: TypedView(s, 64, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
u32 = property(lambda s: TypedView(s, 32), lambda s, v: setattr(s, '_val', int(v) & MASK32))
i32 = property(lambda s: TypedView(s, 32, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
b32 = property(lambda s: TypedView(s, 32), lambda s, v: setattr(s, '_val', int(v) & MASK32))
f32 = property(lambda s: TypedView(s, 32, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
u24 = property(lambda s: TypedView(s, 24))
i24 = property(lambda s: TypedView(s, 24, signed=True))
u16 = property(lambda s: TypedView(s, 16), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
i16 = property(lambda s: TypedView(s, 16, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
b16 = property(lambda s: TypedView(s, 16), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
f16 = property(lambda s: TypedView(s, 16, 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, 16, 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, 8))
i8 = property(lambda s: TypedView(s, 8, signed=True))
u1 = property(lambda s: TypedView(s, 1)) # single bit
def __getitem__(s, key):
if isinstance(key, slice): return SliceProxy(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)
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)
+670
View File
@@ -0,0 +1,670 @@
# Generate AMD ISA autogen files from PDF documentation
# Combines format/enum generation (previously in dsl.py) and pseudocode compilation (previously in pcode.py)
# Usage: python -m extra.assembly.amd.pdf [--arch rdna3|rdna4|cdna|all]
import re, functools
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
PDF_URLS = {
"rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",
"rdna4": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content",
"cdna": ["https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-mi300-cdna3-instruction-set-architecture.pdf",
"https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf"],
}
# Field type mappings and ordering
FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField',
'VDATA': 'VGPRField', 'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField',
'SIMM16': 'SImm', 'OFFSET': 'Imm', 'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src',
'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
FIELD_ORDER = {
'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
'4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
INST_PATTERN = re.compile(r'^([SVD]S?_[A-Z0-9_]+)\s+(\d+)\s*$', re.M)
# Patterns that can't be handled by the DSL (require special handling in emu.py)
UNSUPPORTED = ['SGPR[', 'V_SWAP', 'eval ', 'FATAL_HALT', 'HW_REGISTERS',
'vscnt', 'vmcnt', 'expcnt', 'lgkmcnt',
'CVT_OFF_TABLE', 'ThreadMask',
'S1[i', 'C.i32',
'if n.', 'DST.u32', 'addrd = DST', 'addr = DST',
'BARRIER_STATE', 'ReallocVgprs',
'GPR_IDX', 'VSKIP', 'specified in', 'TTBL',
'fp6', 'bf6'] # Malformed pseudocode from PDF
# ═══════════════════════════════════════════════════════════════════════════════
# COMPILER: pseudocode -> Python (minimal transforms)
# ═══════════════════════════════════════════════════════════════════════════════
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
declared_arrays: dict[str, int] = {} # Track declared arrays: name -> size
for line in joined_lines:
line = line.strip()
if not line or line.startswith('//'): continue
if line.startswith('if '):
lines.append(' ' * indent + f"if {_expr(line[3:].rstrip(' then'), declared_arrays)}:")
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'), declared_arrays)}:")
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 m := re.match(r'declare\s+(\w+)\s*:\s*\d+\'[FBU]\[(\d+)\]', line):
# Handle array declarations: declare in : 32'F[3] or declare S : 32'B[3]
arr_name, arr_size = m[1], int(m[2])
declared_arrays[arr_name] = arr_size
py_name = f"{arr_name}_" if arr_name == 'in' else arr_name # 'in' is Python keyword
if arr_name == 'S':
lines.append(' ' * indent + f"{py_name} = [S0, S1, S2]") # Map to source registers
else:
lines.append(' ' * indent + f"{py_name} = [Reg(0) for _ in range({arr_size})]")
elif line.startswith('declare '):
pass # Ignore other declare statements
elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
start, end = _expr(m[2].strip(), declared_arrays), _expr(m[3].strip(), declared_arrays)
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], declared_arrays)
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)
lhs_s = _expr(lhs.strip(), declared_arrays) # Transform LHS too for array access
lines.append(' ' * indent + f"{lhs_s} {op} {_expr(rhs.strip(), declared_arrays)}")
break
else:
lhs, rhs = line.split('=', 1)
lhs_s, rhs_s = lhs.strip(), rhs.strip()
lhs_t = _expr(lhs_s, declared_arrays) # Transform LHS for array access
stmt = _assign(lhs_t, _expr(rhs_s, declared_arrays), declared_arrays)
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, declared_arrays: dict[str, int] | None = None) -> str:
# Check for array element assignment: in_[i] should not wrap in Reg()
if declared_arrays and re.match(r'\w+_?\[\w+\]', lhs):
return f"{lhs} = {rhs}"
if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
return f"{lhs} = Reg({rhs})"
return f"{lhs} = {rhs}"
def _expr(e: str, declared_arrays: dict[str, int] | None = None) -> str:
e = e.strip()
# Handle OPSEL_HI.u3[i] and OPSEL.u3[i] - bit extraction from opsel fields
e = re.sub(r'(OPSEL(?:_HI)?)\.u\d+\[(\w+)\]', r'((\1 >> \2) & 1)', e)
# Rename 'in' to 'in_' to avoid Python keyword conflict
e = re.sub(r'\bin\[', 'in_[', e)
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(), declared_arrays), _expr(m[2].strip(), declared_arrays)
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], declared_arrays)
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
# ═══════════════════════════════════════════════════════════════════════════════
# PDF PARSING WITH PAGE CACHING
# ═══════════════════════════════════════════════════════════════════════════════
class CachedPDF:
"""PDF wrapper with page text/table caching for faster repeated access."""
def __init__(self, pdf):
self._pdf, self._text_cache, self._table_cache = pdf, {}, {}
def __len__(self): return len(self._pdf.pages)
def text(self, i):
if i not in self._text_cache: self._text_cache[i] = self._pdf.pages[i].extract_text() or ''
return self._text_cache[i]
def tables(self, i):
if i not in self._table_cache: self._table_cache[i] = [t.extract() for t in self._pdf.pages[i].find_tables()]
return self._table_cache[i]
def _parse_bits(s: str) -> tuple[int, int] | None:
return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
def _parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
fields = []
for row in table[1:]:
if not row or not row[0]: continue
name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
if not (bits := _parse_bits(bits_str)): continue
enc_val, hi, lo = None, bits[0], bits[1]
if name == 'ENCODING' and row[2]:
if m := re.search(r"(?:'b|Must be:\s*)([01_]+)", row[2]):
enc_bits = m.group(1).replace('_', '')
enc_val, declared_width, actual_width = int(enc_bits, 2), hi - lo + 1, len(enc_bits)
if actual_width > declared_width: lo = hi - actual_width + 1
ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
fields.append((name, hi, lo, enc_val, ftype))
return fields
def _parse_single_pdf(url: str):
"""Parse a single PDF and return (formats, enums, src_enum, doc_name, instructions)."""
import pdfplumber
from tinygrad.helpers import fetch
pdf = CachedPDF(pdfplumber.open(fetch(url)))
total_pages = len(pdf)
# Auto-detect document type
first_page = pdf.text(0)
is_cdna4, is_cdna3 = 'CDNA4' in first_page or 'CDNA 4' in first_page, 'CDNA3' in first_page or 'MI300' in first_page
is_cdna, is_rdna4 = is_cdna3 or is_cdna4, 'RDNA4' in first_page or 'RDNA 4' in first_page
is_rdna35, is_rdna3 = 'RDNA3.5' in first_page or 'RDNA 3.5' in first_page, 'RDNA3' in first_page and 'RDNA3.5' not in first_page
doc_name = "CDNA4" if is_cdna4 else "CDNA3" if is_cdna3 else "RDNA4" if is_rdna4 else "RDNA3.5" if is_rdna35 else "RDNA3" if is_rdna3 else "Unknown"
# Find Microcode Formats section (for formats/enums)
microcode_start = next((i for i in range(int(total_pages * 0.2), total_pages)
if re.search(r'\d+\.\d+\.\d+\.\s+SOP2\b|Chapter \d+\.\s+Microcode Formats', pdf.text(i))), int(total_pages * 0.9))
# Find Instructions section (for pseudocode)
instr_start = next((i for i in range(int(total_pages * 0.1), int(total_pages * 0.5))
if re.search(r'Chapter \d+\.\s+Instructions\b', pdf.text(i))), total_pages // 3)
instr_end = next((i for start in [int(total_pages * 0.6), int(total_pages * 0.5), instr_start]
for i in range(start, min(start + 100, total_pages))
if re.search(r'Chapter \d+\.\s+Microcode Formats', pdf.text(i))), total_pages)
# Parse src enum from SSRC encoding table
src_enum = dict(SRC_EXTRAS)
for i in range(microcode_start, min(microcode_start + 10, total_pages)):
text = pdf.text(i)
if 'SSRC0' in text and 'VCC_LO' in text:
for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
val, name = int(m.group(1)), m.group(2).rstrip('.:')
if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
break
# Parse opcode tables
full_text = '\n'.join(pdf.text(i) for i in range(microcode_start, min(microcode_start + 50, total_pages)))
enums: dict[str, dict[int, str]] = {}
for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
enums[m.group(1) + "Op"] = ops
if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
enums["VOPDOp"] = ops
enum_names = set(enums.keys())
# Parse instruction formats
def is_fields_table(t): return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
def has_encoding(fields): return any(f[0] == 'ENCODING' for f in fields)
def has_header_before_fields(text): return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
format_headers = []
for i in range(50):
if microcode_start + i >= total_pages: break
text = pdf.text(microcode_start + i)
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
fmt_name = m.group(1)
if is_cdna and fmt_name.isupper() and len(fmt_name) >= 2: format_headers.append((fmt_name, i, m.start()))
elif m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < 50:
next_text = pdf.text(microcode_start + i + 1).lstrip()
if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
format_headers.append((fmt_name, i, m.start()))
formats: dict[str, list] = {}
for fmt_name, rel_idx, header_pos in format_headers:
if fmt_name in formats: continue
page_idx = microcode_start + rel_idx
text = pdf.text(page_idx)
field_pos = text.find('Field Name', header_pos)
fields = None
for offset in range(3):
if page_idx + offset >= total_pages: break
if offset > 0 and has_header_before_fields(pdf.text(page_idx + offset)): break
for t in pdf.tables(page_idx + offset) if offset > 0 or field_pos > header_pos else []:
if is_fields_table(t) and (f := _parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f): fields = f; break
if fields: break
if not fields and field_pos > header_pos:
for t in pdf.tables(page_idx):
if is_fields_table(t) and (f := _parse_fields_table(t, fmt_name, enum_names)): fields = f; break
if not fields: continue
field_names = {f[0] for f in fields}
for pg_offset in range(1, 3):
if page_idx + pg_offset >= total_pages or has_header_before_fields(pdf.text(page_idx + pg_offset)): break
for t in pdf.tables(page_idx + pg_offset):
if is_fields_table(t) and (extra := _parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
for ef in extra:
if ef[0] not in field_names: fields.append(ef); field_names.add(ef[0])
break
formats[fmt_name] = fields
# Fix known PDF errors
if 'SMEM' in formats:
formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
for n, h, l, e, t in formats['SMEM']]
if doc_name in ('RDNA3', 'RDNA3.5'):
if 'SOPPOp' in enums: assert 8 not in enums['SOPPOp']; enums['SOPPOp'][8] = 'S_WAITCNT_DEPCTR'
if 'DSOp' in enums:
for k, v in {24: 'DS_GWS_SEMA_RELEASE_ALL', 25: 'DS_GWS_INIT', 26: 'DS_GWS_SEMA_V', 27: 'DS_GWS_SEMA_BR', 28: 'DS_GWS_SEMA_P', 29: 'DS_GWS_BARRIER'}.items():
assert k not in enums['DSOp']; enums['DSOp'][k] = v
if 'FLATOp' in enums:
for k, v in {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}.items():
assert k not in enums['FLATOp']; enums['FLATOp'][k] = v
# Extract pseudocode for instructions
all_text = '\n'.join(pdf.text(i) for i in range(instr_start, instr_end))
matches = list(INST_PATTERN.finditer(all_text))
raw_pseudocode: dict[tuple[str, int], str] = {}
for i, match in enumerate(matches):
name, opcode = match.group(1), int(match.group(2))
start, end = match.end(), matches[i + 1].start() if i + 1 < len(matches) else match.end() + 2000
snippet = all_text[start:end].strip()
if pseudocode := _extract_pseudocode(snippet): raw_pseudocode[(name, opcode)] = pseudocode
return {"formats": formats, "enums": enums, "src_enum": src_enum, "doc_name": doc_name, "pseudocode": raw_pseudocode, "is_cdna": is_cdna}
def _extract_pseudocode(text: str) -> str | None:
"""Extract pseudocode from an instruction description snippet."""
lines, result, depth, in_lambda = text.split('\n'), [], 0, 0
for line in lines:
s = line.strip()
if not s or re.match(r'^\d+ of \d+$', s) or re.match(r'^\d+\.\d+\..*Instructions', s): continue
if s.startswith(('Notes', 'Functional examples')): break
if s.startswith(('"RDNA', 'AMD ', 'CDNA')): continue
if '= lambda(' in s: in_lambda += 1; continue
if in_lambda > 0:
if s.endswith(');'): in_lambda -= 1
continue
if s.startswith('if '): depth += 1
elif s.startswith('endif'): depth = max(0, depth - 1)
if s.endswith('.') and not any(p in s for p in ['D0', 'D1', 'S0', 'S1', 'S2', 'SCC', 'VCC', 'tmp', '=']): continue
if re.match(r'^[a-z].*\.$', s) and '=' not in s: continue
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', 'DATA.', 'DATA0', 'DATA1', 'ADDR']) 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 is_code: result.append(s)
return '\n'.join(result) if result else None
def _merge_results(results: list[dict]) -> dict:
"""Merge multiple PDF parse results into a superset."""
merged = {"formats": {}, "enums": {}, "src_enum": dict(SRC_EXTRAS), "doc_names": [], "pseudocode": {}, "is_cdna": False}
for r in results:
merged["doc_names"].append(r["doc_name"])
merged["is_cdna"] = merged["is_cdna"] or r["is_cdna"]
for val, name in r["src_enum"].items():
if val in merged["src_enum"]: assert merged["src_enum"][val] == name
else: merged["src_enum"][val] = name
for enum_name, ops in r["enums"].items():
if enum_name not in merged["enums"]: merged["enums"][enum_name] = {}
for val, name in ops.items():
if val in merged["enums"][enum_name]: assert merged["enums"][enum_name][val] == name
else: merged["enums"][enum_name][val] = name
for fmt_name, fields in r["formats"].items():
if fmt_name not in merged["formats"]: merged["formats"][fmt_name] = list(fields)
else:
existing = {f[0]: (f[1], f[2]) for f in merged["formats"][fmt_name]}
for f in fields:
if f[0] in existing: assert existing[f[0]] == (f[1], f[2])
else: merged["formats"][fmt_name].append(f)
for key, pc in r["pseudocode"].items():
if key not in merged["pseudocode"]: merged["pseudocode"][key] = pc
return merged
# ═══════════════════════════════════════════════════════════════════════════════
# CODE GENERATION
# ═══════════════════════════════════════════════════════════════════════════════
def _generate_enum_py(enums, src_enum, doc_name) -> str:
"""Generate enum.py content (just enums, no dsl.py dependency)."""
def enum_lines(name, items): return [f"class {name}(IntEnum):"] + [f" {n} = {v}" for v, n in sorted(items.items())] + [""]
lines = [f"# autogenerated from AMD {doc_name} ISA PDF by pdf.py - do not edit", "from enum import IntEnum", ""]
lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
return '\n'.join(lines)
def _generate_ins_py(formats, enums, src_enum, doc_name) -> str:
"""Generate ins.py content (instruction formats and helpers, imports dsl.py and enum.py)."""
def field_key(f, order): return order.index(f[0].lower()) if f[0].lower() in order else 1000
lines = [f"# autogenerated from AMD {doc_name} ISA PDF by pdf.py - do not edit",
"# ruff: noqa: F401,F403", "from typing import Annotated",
"from extra.assembly.amd.dsl import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
"from extra.assembly.amd.autogen.{arch}.enum import *",
"import functools", ""]
format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
lines.append("# instruction formats")
for fmt_name, fields in sorted(formats.items()):
base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
order = FIELD_ORDER.get(fmt_name, [])
lines.append(f"class {fmt_name}({base}):")
if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
lines.append(f" encoding = bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f" encoding = bits[{enc[1]}] == {enc[3]}")
if defaults := format_defaults.get(fmt_name): lines.append(f" _defaults = {defaults}")
for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=lambda f: field_key(f, order)):
ann = f":Annotated[BitField, {ftype}]" if ftype and ftype.endswith('Op') else f":{ftype}" if ftype else ""
lines.append(f" {name.lower()}{ann} = bits[{hi}]" if hi == lo else f" {name.lower()}{ann} = bits[{hi}:{lo}]")
lines.append("")
lines.append("# instruction helpers")
for cls_name, ops in sorted(enums.items()):
fmt = cls_name[:-2]
for op_val, name in sorted(ops.items()):
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=1"}.get(fmt, "")
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
suffix = "_e32" if fmt in ("VOP1", "VOP2", "VOPC") else "_e64" if fmt == "VOP3" and op_val < 512 else ""
if name in ('V_FMAMK_F32', 'V_FMAMK_F16'):
lines.append(f"def {name.lower()}{suffix}(vdst, src0, K, vsrc1): return {fmt}({cls_name}.{name}, vdst, src0, vsrc1, literal=K)")
elif name in ('V_FMAAK_F32', 'V_FMAAK_F16'):
lines.append(f"def {name.lower()}{suffix}(vdst, src0, vsrc1, K): return {fmt}({cls_name}.{name}, vdst, src0, vsrc1, literal=K)")
else: lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
src_names = {name for _, name in src_enum.items()}
lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in {'DPP8', 'DPP16'}]
if "NULL" in src_names: lines.append("OFF = NULL\n")
return '\n'.join(lines)
def _generate_gen_pcode_py(enums, pseudocode, arch) -> str:
"""Generate gen_pcode.py content (compiled pseudocode functions)."""
# Get op enums for this arch (import from .ins which re-exports from .enum)
import importlib
autogen = importlib.import_module(f"extra.assembly.amd.autogen.{arch}.ins")
OP_ENUMS = [getattr(autogen, name) for name in ['SOP1Op', 'SOP2Op', 'SOPCOp', 'SOPKOp', 'SOPPOp', 'VOP1Op', 'VOP2Op', 'VOP3Op', 'VOP3SDOp', 'VOP3POp', 'VOPCOp', 'VOP3AOp', 'VOP3BOp', 'DSOp'] if hasattr(autogen, name)]
# Build defined ops mapping
defined_ops: dict[tuple, list] = {}
for enum_cls in OP_ENUMS:
for op in enum_cls:
if op.name.startswith(('S_', 'V_', 'DS_')): defined_ops.setdefault((op.name, op.value), []).append((enum_cls, op))
enum_names = [e.__name__ for e in OP_ENUMS]
lines = [f'''# autogenerated by pdf.py - do not edit
# to regenerate: python -m extra.assembly.amd.pdf --arch {arch}
# ruff: noqa: E501,F405,F403
# mypy: ignore-errors
from extra.assembly.amd.autogen.{arch}.enum import {", ".join(enum_names)}
from extra.assembly.amd.pcode import *
''']
instructions: dict = {cls: {} for cls in OP_ENUMS}
for key, pc in pseudocode.items():
if key in defined_ops:
for enum_cls, enum_val in defined_ops[key]: instructions[enum_cls][enum_val] = pc
for enum_cls in OP_ENUMS:
cls_name = enum_cls.__name__
if not instructions.get(enum_cls): continue
fn_entries = []
for op, pc in instructions[enum_cls].items():
if any(p in pc for p in UNSUPPORTED): continue
try:
code = compile_pseudocode(pc)
code = _apply_pseudocode_fixes(op, code)
fn_name, fn_code = _generate_function(cls_name, op, pc, code)
lines.append(fn_code)
fn_entries.append((op, fn_name))
except Exception as e: print(f" Warning: Failed to compile {op.name}: {e}")
if fn_entries:
lines.append(f'{cls_name}_FUNCTIONS = {{')
for op, fn_name in fn_entries: lines.append(f" {cls_name}.{op.name}: {fn_name},")
lines.append('}\n')
# Add V_WRITELANE_B32 if VOP3Op exists
if 'VOP3Op' in enum_names:
lines.append('''
# V_WRITELANE_B32: Write scalar to specific lane's VGPR (not in PDF pseudocode)
def _VOP3Op_V_WRITELANE_B32(s0, s1, s2, d0, scc, vcc, lane, exec_mask, literal, VGPR, _vars, src0_idx=0, vdst_idx=0):
wr_lane = s1 & 0x1f
return {'d0': d0, 'scc': scc, 'vgpr_write': (wr_lane, vdst_idx, s0 & 0xffffffff)}
VOP3Op_FUNCTIONS[VOP3Op.V_WRITELANE_B32] = _VOP3Op_V_WRITELANE_B32
''')
lines.append('COMPILED_FUNCTIONS = {')
for enum_cls in OP_ENUMS:
if instructions.get(enum_cls): lines.append(f' {enum_cls.__name__}: {enum_cls.__name__}_FUNCTIONS,')
lines.append('}\n\ndef get_compiled_functions(): return COMPILED_FUNCTIONS')
return '\n'.join(lines)
def _apply_pseudocode_fixes(op, code: str) -> str:
"""Apply known fixes for PDF pseudocode bugs."""
if op.name == 'V_DIV_FMAS_F32':
code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
if op.name == 'V_DIV_FMAS_F64':
code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
if op.name == 'V_DIV_SCALE_F32':
code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(0x1); D0.f32 = float("nan")')
code = code.replace('elif S1.f32 == DENORM.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n pass')
code += '\nif S1.f32 == DENORM.f32:\n D0.f32 = float("nan")'
code = code.replace('elif exponent(S2.f32) <= 23:\n D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n VCC = Reg(0x1); D0.f32 = ldexp(S0.f32, 64)')
code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)\n if S0.f32 == S2.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)')
if op.name == 'V_DIV_SCALE_F64':
code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(0x1); D0.f64 = float("nan")')
code = code.replace('elif S1.f64 == DENORM.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n pass')
code += '\nif S1.f64 == DENORM.f64:\n D0.f64 = float("nan")'
code = code.replace('elif exponent(S2.f64) <= 52:\n D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n VCC = Reg(0x1); D0.f64 = ldexp(S0.f64, 128)')
code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)\n if S0.f64 == S2.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)')
if op.name == 'V_DIV_FIXUP_F32':
code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
if op.name == 'V_DIV_FIXUP_F64':
code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
if op.name == 'V_TRIG_PREOP_F64':
code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
return code
def _generate_function(cls_name: str, op, pc: str, code: str) -> tuple[str, str]:
"""Generate a single compiled pseudocode function."""
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)
has_opsel = 'OPSEL' in pc # FMA_MIX and similar instructions need OPSEL/OPSEL_HI
combined = code + pc
fn_name = f"_{cls_name}_{op.name}"
# Function accepts Reg objects directly (uppercase names), laneId is passed directly as int
params = "S0, S1, S2, D0, SCC, VCC, laneId, EXEC, literal, VGPR, src0_idx=0, vdst_idx=0, PC=None"
if has_opsel: params += ", OPSEL=0, OPSEL_HI=0"
lines = [f"def {fn_name}({params}):"]
# Registers that need special handling (not passed directly)
# Only init if used but not first assigned as `name = Reg(...)` in the compiled code
def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
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)'))
used = {name for name, _ in special_regs if name in combined}
# Detect which registers are modified (not just read) - look for assignments
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 init code for special registers
init_lines = []
if is_div_scale: init_lines.append(" D0 = Reg(S0._val)")
for name, init in special_regs:
if name in used: init_lines.append(f" {name} = {init}")
if 'EXEC_LO' in code: init_lines.append(" EXEC_LO = SliceProxy(EXEC, 31, 0)")
if 'EXEC_HI' in code: init_lines.append(" EXEC_HI = SliceProxy(EXEC, 63, 32)")
if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_lines.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_lines.append(" EXECZ = Reg(1 if EXEC._val == 0 else 0)")
code_lines = [line for line in code.split('\n') if line.strip()]
if init_lines:
lines.extend(init_lines)
if code_lines: lines.append(" # --- compiled pseudocode ---")
for line in code_lines:
lines.append(f" {line}")
# Build result dict - only include registers that are modified
result_items = []
if modifies_d0: result_items.append("'D0': D0")
if modifies_scc: result_items.append("'SCC': SCC")
if modifies_vcc: result_items.append("'VCC': VCC")
if modifies_exec: result_items.append("'EXEC': EXEC")
if has_d1: result_items.append("'D1': D1")
if modifies_pc: result_items.append("'PC': PC")
lines.append(f" return {{{', '.join(result_items)}}}\n")
return fn_name, '\n'.join(lines)
# ═══════════════════════════════════════════════════════════════════════════════
# MAIN GENERATION
# ═══════════════════════════════════════════════════════════════════════════════
def generate_arch(arch: str) -> dict:
"""Generate enum.py, ins.py and gen_pcode.py for a single architecture."""
urls = PDF_URLS[arch]
if isinstance(urls, str): urls = [urls]
print(f"\n{'='*60}\nGenerating {arch}...")
print(f"Parsing {len(urls)} PDF(s)...")
results = [_parse_single_pdf(url) for url in urls]
merged = _merge_results(results) if len(results) > 1 else results[0]
doc_name = "+".join(merged["doc_names"]) if len(results) > 1 else merged["doc_name"]
base_path = Path(f"extra/assembly/amd/autogen/{arch}")
base_path.mkdir(parents=True, exist_ok=True)
(base_path / "__init__.py").touch()
# Write enum.py (enums only, no dsl.py dependency)
enum_path = base_path / "enum.py"
enum_content = _generate_enum_py(merged["enums"], merged["src_enum"], doc_name)
enum_path.write_text(enum_content)
print(f"Generated {enum_path}: SrcEnum ({len(merged['src_enum'])}) + {len(merged['enums'])} enums")
# Write ins.py (instruction formats and helpers, imports dsl.py and enum.py)
ins_path = base_path / "ins.py"
ins_content = _generate_ins_py(merged["formats"], merged["enums"], merged["src_enum"], doc_name).replace("{arch}", arch)
ins_path.write_text(ins_content)
print(f"Generated {ins_path}: {len(merged['formats'])} formats")
# Write gen_pcode.py (needs enum.py to exist first for imports)
pcode_path = base_path / "gen_pcode.py"
pcode_content = _generate_gen_pcode_py(merged["enums"], merged["pseudocode"], arch)
pcode_path.write_text(pcode_content)
print(f"Generated {pcode_path}: {len(merged['pseudocode'])} instructions")
return merged
def _generate_arch_wrapper(arch: str):
"""Wrapper for multiprocessing - returns arch name for ordering."""
generate_arch(arch)
return arch
def generate_all():
"""Generate all architectures in parallel."""
with ProcessPoolExecutor() as executor:
list(executor.map(_generate_arch_wrapper, PDF_URLS.keys()))
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Generate AMD ISA autogen files from PDF documentation")
parser.add_argument("--arch", choices=list(PDF_URLS.keys()) + ["all"], default="rdna3")
args = parser.parse_args()
if args.arch == "all": generate_all()
else: generate_arch(args.arch)
@@ -7,7 +7,7 @@ from typing import Callable
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.rdna3.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
@@ -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()
+66
View File
@@ -0,0 +1,66 @@
"""Shared test helpers for RDNA3 tests."""
import shutil
from dataclasses import dataclass
@dataclass
class KernelInfo:
code: bytes
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
buf_sizes: list[int] # sizes for each buffer index
# LLVM tool detection (shared across test files)
def get_llvm_mc():
"""Find llvm-mc executable, preferring newer versions."""
for p in ['llvm-mc', 'llvm-mc-21', 'llvm-mc-20']:
if shutil.which(p): return p
raise FileNotFoundError("llvm-mc not found")
def get_llvm_objdump():
"""Find llvm-objdump executable, preferring newer versions."""
for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
if shutil.which(p): return p
raise FileNotFoundError("llvm-objdump not found")
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION CONTEXT (for testing compiled pseudocode)
# ═══════════════════════════════════════════════════════════════════════════════
class ExecContext:
"""Context for running compiled pseudocode in tests."""
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, SliceProxy
self._Reg, self._MASK64, self._SliceProxy = Reg, MASK64, SliceProxy
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._SliceProxy(self.EXEC, 31, 0), 'EXEC_HI': self._SliceProxy(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}
@@ -7,18 +7,21 @@ from pathlib import Path
# This allows generating AMD GPU kernels without requiring real hardware
os.environ["AMD"] = "1"
os.environ["MOCKGPU"] = "1"
os.environ["PYTHON_REMU"] = "1"
from extra.assembly.rdna3.emu import WaveState, decode_program, step_wave, WAVE_SIZE
from extra.assembly.amd.emu import WaveState, decode_program, step_wave, WAVE_SIZE, set_valid_mem_ranges
from extra.assembly.amd.test.helpers import KernelInfo
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
@dataclass
class KernelInfo:
code: bytes
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
buf_sizes: list[int] # sizes for each buffer index
def _is_f32_nan(bits: int) -> bool:
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
def _vals_equal(a: int, b: int) -> bool:
"""Compare two 32-bit values, treating all NaN bit patterns as equal."""
if a == b: return True
return _is_f32_nan(a) and _is_f32_nan(b)
@dataclass
class StateSnapshot:
@@ -29,20 +32,20 @@ class StateSnapshot:
sgpr: list[int]
vgpr: list[list[int]]
def diff(self, other: 'StateSnapshot', n_lanes: int) -> list[str]:
def diff(self, other: 'StateSnapshot', n_lanes: int, arrow: str = " vs ") -> list[str]:
"""Return list of differences between two states."""
diffs = []
if self.pc != other.pc: diffs.append(f"pc: {self.pc} vs {other.pc}")
if self.scc != other.scc: diffs.append(f"scc: {self.scc} vs {other.scc}")
if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x} vs 0x{other.vcc:08x}")
if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x} vs 0x{other.exec_mask:08x}")
if self.pc != other.pc: diffs.append(f"pc: {self.pc}{arrow}{other.pc}")
if self.scc != other.scc: diffs.append(f"scc: {self.scc}{arrow}{other.scc}")
if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x}{arrow}0x{other.vcc:08x}")
if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x}{arrow}0x{other.exec_mask:08x}")
for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
# Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
if i in (106, 107, 126, 127): continue
if a != b: diffs.append(f"sgpr[{i}]: 0x{a:08x} vs 0x{b:08x}")
if not _vals_equal(a, b): diffs.append(f"sgpr[{i}]: 0x{a:08x}{arrow}0x{b:08x}")
for lane in range(n_lanes):
for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
if a != b: diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x} vs 0x{b:08x}")
if not _vals_equal(a, b): diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x}{arrow}0x{b:08x}")
return diffs
class CStateSnapshot(ctypes.Structure):
@@ -117,7 +120,7 @@ class PythonEmulator:
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
max_workgroups: int = 64) -> tuple[bool, str, int]:
max_workgroups: int = 8) -> tuple[bool, str, int]:
"""Run a single kernel through both emulators. Returns (success, message, total_steps)."""
gx, gy, gz = global_size
total_steps = 0
@@ -157,25 +160,52 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
# v_div_scale/v_div_fixup: Rust has different VCC handling
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
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'))
diffs = rust_before.diff(python_before, n_lanes)
if diffs:
trace_lines = []
for s, pc, d, rb, pb in trace[:-1]:
for idx, (s, pc, d, rb, pb) in enumerate(trace):
trace_lines.append(f" step {s}: PC={pc:3d} {d}")
if trace.index((s, pc, d, rb, pb)) < len(trace) - 2:
next_rb, next_pb = trace[trace.index((s, pc, d, rb, pb)) + 1][3:5]
inst_diffs = rb.diff(next_rb, n_lanes)
if inst_diffs: trace_lines.append(f" rust changes: {', '.join(inst_diffs[:3])}")
if idx < len(trace) - 1:
next_rb, next_pb = trace[idx + 1][3:5]
rust_diffs = rb.diff(next_rb, n_lanes, "->")
python_diffs = pb.diff(next_pb, n_lanes, "->")
if rust_diffs: trace_lines.append(f" rust: {', '.join(rust_diffs[:5])}")
if python_diffs: trace_lines.append(f" python: {', '.join(python_diffs[:5])}")
elif rust_diffs: trace_lines.append(f" python: (no changes)")
else:
# Last traced instruction - compare with current state
rust_diffs = rb.diff(rust_before, n_lanes, "->")
python_diffs = pb.diff(python_before, n_lanes, "->")
if rust_diffs: trace_lines.append(f" rust: {', '.join(rust_diffs[:5])}")
if python_diffs: trace_lines.append(f" python: {', '.join(python_diffs[:5])}")
elif rust_diffs: trace_lines.append(f" python: (no changes)")
trace_str = "\n".join(trace_lines)
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ:\n " + "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}", total_steps
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n " + "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}", total_steps
rust_result = rust.step()
python_result = python.step()
if rust_result != python_result:
# Rust returns 1 for unsupported instructions - skip test
if rust_result == 1 and python_result == 0:
raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
trace_str = "\n".join(f" step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n Recent instructions:\n{trace_str}", total_steps
# Sync Python state to Rust after instructions with known Rust emulator differences
if sync_after:
rust_after = rust.get_snapshot()
for i in range(128): python.set_sgpr(i, rust_after.sgpr[i])
for lane in range(n_lanes):
for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
assert python.state is not None
python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
if rust_result == -1:
total_steps += step + 1
break
@@ -196,7 +226,6 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
"""Run all kernels through both emulators with shared buffer pool."""
from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
if buf_data is None: buf_data = {}
# Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
@@ -240,8 +269,6 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
"""Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
# Allocate buffers
buffers = []
for size in buf_sizes:
@@ -315,7 +342,6 @@ def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[
k = kernels[-1]
return k.code, k.global_size, k.local_size, k.buf_sizes
@unittest.skipUnless(REMU_PATH.exists(), "libremu.so not found")
class TestTinygradKernels(unittest.TestCase):
"""Compare emulators on real tinygrad-compiled kernels."""
@@ -324,143 +350,53 @@ class TestTinygradKernels(unittest.TestCase):
ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
self.assertTrue(ok, msg)
# Basic unary ops
def test_neg(self): self._test_kernel(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Basic ops - consolidated tests covering key instruction patterns
def test_unary_ops(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu().exp().log().sqrt().reciprocal())
def test_binary_ops(self): self._test_kernel(lambda T: (T([1.0, 2.0]) + T([3.0, 4.0])) * T([0.5, 0.5]) - T([1.0, 1.0]))
def test_trig(self): self._test_kernel(lambda T: T([0.1, 1.0, 3.14, -1.0]*8).sin() + T([0.1, 1.0, 3.14, -1.0]*8).cos())
def test_compare(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_bitwise(self): self._test_kernel(lambda T: (T([0xF0, 0x0F, 0xFF]*11).int() & T([0x0F, 0x0F, 0x00]*11).int()) | T([1]*33).int())
def test_int_ops(self): self._test_kernel(lambda T: ((T.empty(64).int() + T.empty(64).int()) * T.empty(64).int()).float())
# Reductions
def test_sum_reduce(self): self._test_kernel(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel(lambda T: T.empty(32).mean())
def test_reduce(self): self._test_kernel(lambda T: T.empty(64).sum() + T.empty(64).max())
def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
# Matmul - various sizes
def test_gemm_4x4(self): self._test_kernel(lambda T: T.empty(4, 4) @ T.empty(4, 4), max_steps=100000)
def test_gemm_8x8(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=200000)
@unittest.skip("too slow")
def test_gemm_16x16(self): self._test_kernel(lambda T: T.empty(16, 16) @ T.empty(16, 16), max_steps=500000)
def test_gemv(self): self._test_kernel(lambda T: T.empty(1, 16) @ T.empty(16, 16), max_steps=100000)
# Matmul
def test_gemm(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=100000)
@unittest.skip("Rust emulator crashes on this kernel (assertion failure in thread.rs)")
def test_gemm_fp16(self): self._test_kernel(lambda T: T.empty(16, 16).half() @ T.empty(16, 16).half(), max_steps=100000)
# Complex ops
def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
def test_memory(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous() + T.empty(4, 1).expand(4, 4))
# Cast ops
def test_cast_int(self): self._test_kernel(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel(lambda T: T.empty(16).half().float())
def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
# Min/max (uses comparison internally)
def test_min_binary(self): self._test_kernel(lambda T: T([1.0, 5.0, 3.0]).minimum(T([3.0, 2.0, 4.0])))
# Pooling - regression for VCC wave32 mode
def test_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
# Comparison ops (test VOPC instructions) - use 32+ elements to force vector instructions
def test_cmp_lt(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_cmp_eq(self): self._test_kernel(lambda T: (T.empty(64) == T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Convolution
def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
# Bitwise ops
def test_bitwise_and(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0xFF]).int() & T([0x0F, 0x0F, 0x00]).int())
def test_bitwise_or(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0x00]).int() | T([0x0F, 0x0F, 0xFF]).int())
def test_bitwise_xor(self): self._test_kernel(lambda T: T([0xFF, 0x0F, 0xF0]).int() ^ T([0x0F, 0xF0, 0xF0]).int())
# Integer ops - use 32+ elements to force vector instructions
def test_int_add(self): self._test_kernel(lambda T: (T.empty(64).int() + T.empty(64).int()).float())
def test_int_mul(self): self._test_kernel(lambda T: (T.empty(64).int() * T.empty(64).int()).float())
def test_int_mod(self): self._test_kernel(lambda T: (T.empty(64).int().abs() % (T.empty(64).int().abs() + 1)).float())
# More math ops - use 32+ elements to force vector instructions
def test_abs(self): self._test_kernel(lambda T: T.empty(64).abs())
def test_floor(self): self._test_kernel(lambda T: T.empty(64).floor())
def test_ceil(self): self._test_kernel(lambda T: T.empty(64).ceil())
def test_trunc(self): self._test_kernel(lambda T: T.empty(64).trunc())
# Fused ops
def test_fma(self): self._test_kernel(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
# Argmax/argmin (tests different reduction pattern) - use 32+ elements to force vector instructions
def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
def test_argmin(self): self._test_kernel(lambda T: T.empty(64).argmin())
# Exact value tests - use 32+ elements to force vector instructions (small tensors use scalar ops which Rust emu doesn't fully support)
def test_abs_exact(self): self._test_kernel(lambda T: T([-1., 0., 1.]*11).abs()) # 33 elements
def test_neg_exact(self): self._test_kernel(lambda T: -T([-1., 0., 1.]*11))
def test_log_special(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).log())
def test_exp_exact(self): self._test_kernel(lambda T: T([0., 1., -1.]*11).exp())
def test_reciprocal_exact(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).reciprocal())
# Integer division and mod - use 32+ elements
def test_int_div(self): self._test_kernel(lambda T: (T([10, 20, 30]*11).int() // T([3, 4, 5]*11).int()).float())
def test_int_neg(self): self._test_kernel(lambda T: (-T([1, -2, 3]*11).int()).float())
# Mixed precision - use 32+ elements
def test_half_add(self): self._test_kernel(lambda T: (T([1., 2.]*16).half() + T([3., 4.]*16).half()).float())
def test_half_mul(self): self._test_kernel(lambda T: (T([2., 3.]*16).half() * T([4., 5.]*16).half()).float())
# Matrix ops - patterns from test_ops.py failures
def test_cat(self): self._test_kernel(lambda T: T.empty(32, 64).cat(T.empty(32, 64), dim=1))
def test_gather(self): self._test_kernel(lambda T: T.empty(64).gather(0, T.arange(32).int()))
# Tests from test_ops.py that are failing
def test_permute(self): self._test_kernel(lambda T: T.empty(3, 4, 5, 6).permute((3, 2, 1, 0)).contiguous())
def test_cat_large(self): self._test_kernel(lambda T: T.empty(45, 65, 9).cat(T.empty(45, 65, 9), T.empty(45, 65, 9), dim=1))
def test_gather_small(self): self._test_kernel(lambda T: T.empty(10).gather(0, T.arange(5).int()))
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
def test_cross_entropy(self): self._test_kernel(lambda T: T.randn(32, 10).softmax().log().sum())
def test_cross_entropy_class(self):
import numpy as np
np.random.seed(0)
classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
x_np = np.random.randn(32, 10).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(32,10) + 0).cross_entropy((T(classes).int().reshape(32) + 0)))
# Regression tests for BFE operations with width=0 (walrus operator bug)
# Regression tests
def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
def test_interpolate_uint8(self): self._test_kernel(lambda T: T.empty(2,3,64,64).relu().cast('uint8').interpolate((10,10), mode="linear"))
# Regression test for 64-bit comparison (V_CMP_GT_I64, V_CMP_LT_U64, etc.) with rsrc64
def test_interpolate(self): self._test_kernel(lambda T: T.empty(1,2,16,16).relu().cast('uint8').interpolate((8,8), mode="linear"))
def test_index_int64(self):
from tinygrad import dtypes
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
@unittest.skip("only works with mock GPU")
def test_index_int64_2d(self):
from tinygrad import dtypes
# Tests 64-bit compare with inline constants (comparing against 0)
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), T.arange(4).cast(dtypes.int64)])
# Pooling operations - regression test for VCC wave32 mode (S_CBRANCH_VCCZ should only check VCC_LO)
def test_avg_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4), stride=2))
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
def test_avg_pool3d(self):
def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
def test_cross_entropy(self):
import numpy as np
np.random.seed(0)
self._test_kernel(lambda T: T(np.random.randn(1, 1, 16, 16, 16).astype(np.float32).tolist()).avg_pool2d(kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False))
def test_max_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4), stride=2))
# Convolution operations - multi-kernel tests
def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv2d(T.empty(4, 4, 3, 3)), max_steps=100000)
def test_conv_transpose2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv_transpose2d(T.empty(4, 4, 3, 3)), max_steps=200000)
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
def test_conv_transpose3d(self):
import numpy as np
np.random.seed(0)
self._test_kernel(lambda T: T(np.random.randn(2, 4, 9, 9, 9).astype(np.float32).tolist()).conv_transpose2d(
T(np.random.randn(4, 4, 3, 3, 3).astype(np.float32).tolist())), max_steps=500000)
classes = np.random.randint(0, 10, (16,), dtype=np.int32).tolist()
x_np = np.random.randn(16, 10).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
if __name__ == "__main__":
unittest.main()
File diff suppressed because it is too large Load Diff
@@ -1,8 +1,9 @@
#!/usr/bin/env python3
"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
import unittest
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import encode_src
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import encode_src, RawImm
from extra.assembly.amd.asm import detect_format
class TestMUBUF(unittest.TestCase):
"""Test MUBUF (buffer) instructions."""
@@ -308,7 +309,7 @@ class TestVOP3Literal(unittest.TestCase):
def test_vop3_with_literal(self):
# v_add3_u32 v5, vcc_hi, 0xaf123456, v255
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
from extra.assembly.rdna3.lib import RawImm
from extra.assembly.amd.dsl import RawImm
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
self.assertEqual(inst.to_bytes(), expected)
@@ -316,17 +317,91 @@ class TestVOP3Literal(unittest.TestCase):
def test_vop3_literal_null_operand(self):
# v_add3_u32 v5, null, exec_lo, 0xaf123456
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
from extra.assembly.rdna3.lib import RawImm
from extra.assembly.amd.dsl import RawImm
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
self.assertEqual(inst.to_bytes(), expected)
def test_vop3p_with_literal(self):
# Test VOP3P literal encoding (also uses Inst64)
from extra.assembly.rdna3.lib import RawImm
from extra.assembly.amd.dsl import RawImm
inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
self.assertEqual(len(inst.to_bytes()), 12) # 8 bytes + 4 byte literal
class TestDetectFormat(unittest.TestCase):
"""Test detect_format uses encoding from autogen classes."""
def test_detect_sopp(self):
self.assertEqual(detect_format(s_endpgm().to_bytes()), SOPP)
self.assertEqual(detect_format(s_nop(0).to_bytes()), SOPP)
self.assertEqual(detect_format(s_barrier().to_bytes()), SOPP)
def test_detect_sop1(self):
self.assertEqual(detect_format(s_mov_b32(s[0], 0).to_bytes()), SOP1)
self.assertEqual(detect_format(s_mov_b64(s[0:1], 0).to_bytes()), SOP1)
def test_detect_sop2(self):
self.assertEqual(detect_format(s_add_u32(s[0], s[1], s[2]).to_bytes()), SOP2)
self.assertEqual(detect_format(s_mul_i32(s[0], s[1], s[2]).to_bytes()), SOP2)
def test_detect_sopc(self):
self.assertEqual(detect_format(s_cmp_eq_i32(s[0], s[1]).to_bytes()), SOPC)
def test_detect_sopk(self):
self.assertEqual(detect_format(s_movk_i32(s[0], 0x1234).to_bytes()), SOPK)
def test_detect_vop1(self):
self.assertEqual(detect_format(v_mov_b32_e32(v[0], 0).to_bytes()), VOP1)
self.assertEqual(detect_format(v_rcp_f32_e32(v[0], v[1]).to_bytes()), VOP1)
def test_detect_vop2(self):
self.assertEqual(detect_format(v_add_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
self.assertEqual(detect_format(v_mul_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
def test_detect_vopc(self):
self.assertEqual(detect_format(v_cmp_eq_f32_e32(v[0], v[1]).to_bytes()), VOPC)
self.assertEqual(detect_format(v_cmp_lt_i32_e32(v[0], v[1]).to_bytes()), VOPC)
def test_detect_vop3(self):
self.assertEqual(detect_format(v_add_f32_e64(v[0], v[1], v[2]).to_bytes()), VOP3)
self.assertEqual(detect_format(v_fma_f32(v[0], v[1], v[2], v[3]).to_bytes()), VOP3)
def test_detect_vop3p(self):
self.assertEqual(detect_format(VOP3P(VOP3POp.V_PK_ADD_F16, v[0], v[1], v[2], v[3]).to_bytes()), VOP3P)
def test_detect_smem(self):
self.assertEqual(detect_format(s_load_b32(s[0], s[2:3], 0).to_bytes()), SMEM)
self.assertEqual(detect_format(s_load_b64(s[0:1], s[2:3], s[5]).to_bytes()), SMEM)
def test_detect_ds(self):
self.assertEqual(detect_format(ds_load_b32(v[0], v[1]).to_bytes()), DS)
self.assertEqual(detect_format(ds_store_b32(v[0], v[1]).to_bytes()), DS)
def test_detect_flat(self):
self.assertEqual(detect_format(global_load_b32(v[0], v[1:3], RawImm(124)).to_bytes()), FLAT)
self.assertEqual(detect_format(global_store_b32(v[0:2], v[2], RawImm(124)).to_bytes()), FLAT)
def test_detect_mubuf(self):
self.assertEqual(detect_format(buffer_load_b32(v[0], v[1], s[0:4], s[5]).to_bytes()), MUBUF)
def test_detect_mtbuf(self):
self.assertEqual(detect_format(tbuffer_load_format_x(v[0], v[1], s[0:4], s[5], format=22).to_bytes()), MTBUF)
def test_detect_mimg(self):
self.assertEqual(detect_format(image_load(v[0:4], v[4:6], s[0:8], dmask=0xf, dim=1).to_bytes()), MIMG)
def test_detect_exp(self):
self.assertEqual(detect_format(EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3]).to_bytes()), EXP)
def test_detect_vopd(self):
inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=0, srcy0=0)
self.assertEqual(detect_format(inst.to_bytes()), VOPD)
def test_detect_vinterp(self):
inst = VINTERP(VINTERPOp.V_INTERP_P10_F32, vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(detect_format(inst.to_bytes()), VINTERP)
if __name__ == "__main__":
unittest.main()
@@ -2,10 +2,10 @@
# the Inst constructor should be looking at the types of the fields to correctly set the value
import unittest, struct
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import Inst
from extra.assembly.rdna3.asm import asm
from extra.assembly.rdna3.test.test_roundtrip import compile_asm
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.asm import asm
from extra.assembly.amd.test.test_roundtrip import compile_asm
class TestIntegration(unittest.TestCase):
inst: Inst
@@ -1,17 +1,9 @@
#!/usr/bin/env python3
"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
import unittest, re, io, sys
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.asm import waitcnt, asm
def get_amd_toolchain():
"""Check if AMD toolchain is available."""
try:
from tinygrad.runtime.support.compiler_amd import HIPCompiler
HIPCompiler("gfx1100").compile(".text\ns_endpgm")
return True
except Exception:
return False
import unittest, re, io, sys, subprocess
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.asm import waitcnt, asm
from extra.assembly.amd.test.helpers import get_llvm_mc
def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
"""Disassemble ELF binary using tinygrad's compiler, return raw output."""
@@ -47,7 +39,6 @@ def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[
lib = HIPCompiler(arch).compile(asm_src)
return parse_disassembly(disassemble(lib, arch))
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
class TestIntegration(unittest.TestCase):
"""Test our assembler output matches LLVM disassembly."""
@@ -157,7 +148,6 @@ class TestIntegration(unittest.TestCase):
return
self.fail("Could not find s_mov_b32 in disassembly")
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
class TestAsm(unittest.TestCase):
"""Test asm() string parsing."""
@@ -212,10 +202,8 @@ class TestAsm(unittest.TestCase):
def test_asm_vop3_modifiers(self):
"""Test asm() with VOP3 modifiers (neg, abs, clamp)."""
import subprocess, re
def get_llvm_encoding(instr: str) -> str:
result = subprocess.run(['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
input=instr, capture_output=True, text=True)
if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
return m.group(1).replace('0x','').replace(',','').replace(' ','')
@@ -233,7 +221,6 @@ class TestAsm(unittest.TestCase):
llvm_hex = get_llvm_encoding(t)
self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
class TestTinygradIntegration(unittest.TestCase):
"""Test that we can parse disassembled tinygrad kernels."""
@@ -1,10 +1,10 @@
#!/usr/bin/env python3
"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
import unittest, re
import unittest, re, subprocess
from tinygrad.helpers import fetch
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.asm import asm
from extra.assembly.rdna3.test.test_roundtrip import compile_asm, disassemble_lib
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.asm import asm
from extra.assembly.amd.test.helpers import get_llvm_mc
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
@@ -65,12 +65,18 @@ def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
if not asm_text: continue
for j in range(i, min(i + 3, len(lines))):
# Match GFX11, W32, or W64 encodings (all valid for gfx11)
# Format 1: "// GFX11: v_foo ... ; encoding: [0x01,0x02,...]"
# Format 2: "// GFX11: [0x01,0x02,...]" (used by DS, older files)
if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
if hex_bytes:
try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
except ValueError: pass
break
elif m := re.search(r'(?:GFX11|W32|W64)[^:]*:\s*\[(0x[0-9a-fA-F,x\s]+)\]', lines[j]):
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
else:
continue
if hex_bytes:
try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
except ValueError: pass
break
return tests
def try_assemble(text: str):
@@ -78,6 +84,24 @@ def try_assemble(text: str):
try: return asm(text).to_bytes()
except: return None
def compile_asm_batch(instrs: list[str]) -> list[bytes]:
"""Compile multiple instructions with a single llvm-mc call."""
if not instrs: return []
asm_text = ".text\n" + "\n".join(instrs) + "\n"
result = subprocess.run(
[get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=asm_text, capture_output=True, text=True, timeout=30)
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
# Parse all encodings from output
results = []
for line in result.stdout.split('\n'):
if 'encoding:' not in line: continue
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
results.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
if len(results) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(results)}")
return results
class TestLLVM(unittest.TestCase):
"""Test assembler and disassembler against all LLVM test vectors."""
tests: dict[str, list[tuple[str, bytes]]] = {}
@@ -107,59 +131,63 @@ def _make_asm_test(name):
def _make_disasm_test(name):
def test(self):
from tinygrad.runtime.support.compiler_amd import HIPCompiler
compiler = HIPCompiler('gfx1100')
_, fmt_cls, op_enum = LLVM_TEST_FILES[name]
passed, failed, skipped, failures = 0, 0, 0, []
# VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
# Note: opcodes 0-255 are VOPC promoted to VOP3, never VOP3SD
vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
# vop3_from_vopc/vopcx tests have VOPC opcodes 0-255, not VOP3SD - don't detect as VOP3SD
is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
# Undocumented opcodes not in AMD ISA PDF - skip these
undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}} # s_atc_probe*, s_subvector_loop*, s_waitcnt_depctr, unknown
undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}
# First pass: decode all instructions and collect disasm strings
to_test: list[tuple[str, bytes, str | None, str | None]] = [] # (asm_text, data, disasm_str, error)
skipped = 0
for asm_text, data in self.tests.get(name, []):
if len(data) > fmt_cls._size(): continue # skip literals (need different handling)
# Skip undocumented opcodes
if len(data) > fmt_cls._size(): continue
temp_inst = fmt_cls.from_bytes(data)
temp_op = temp_inst._values.get('op', 0)
temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
if temp_op in undocumented.get(name, set()): skipped += 1; continue
# Skip SOPP no-imm instructions with non-zero simm16 (can't roundtrip through LLVM)
if name == 'sopp':
simm16 = temp_inst._values.get('simm16', 0)
simm16 = simm16.val if hasattr(simm16, 'val') else simm16
sopp_no_imm = {48, 54, 53, 55, 60, 61, 62} # s_endpgm, s_barrier, s_wakeup, s_icache_inv, s_wait_idle, s_endpgm_saved, s_code_end
sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}
if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
try:
# VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use
if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
temp = VOP3.from_bytes(data)
op_val = temp._values.get('op', 0)
op_val = op_val.val if hasattr(op_val, 'val') else op_val
is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
# Validate opcode with appropriate enum
if is_vop3sd:
VOP3SDOp(op_val)
else:
VOP3Op(op_val)
if is_vop3sd: VOP3SDOp(op_val)
else: VOP3Op(op_val)
else:
decoded = fmt_cls.from_bytes(data)
op_val = decoded._values.get('op', 0)
op_val = op_val.val if hasattr(op_val, 'val') else op_val
op_enum(op_val) # validate opcode
op_enum(op_val)
if decoded.to_bytes()[:len(data)] != data:
failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue
disasm_str = decoded.disasm()
# Test: LLVM should assemble our disasm output to the same bytes
llvm_bytes = compile_asm(disasm_str, compiler)
if llvm_bytes is None:
failed += 1; failures.append(f"LLVM failed to assemble: '{disasm_str}' (from '{asm_text}')")
elif llvm_bytes == data: passed += 1
else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
to_test.append((asm_text, data, None, "decode roundtrip failed"))
continue
to_test.append((asm_text, data, decoded.disasm(), None))
except Exception as e:
failed += 1; failures.append(f"exception for {data.hex()}: {e}")
to_test.append((asm_text, data, None, f"exception: {e}"))
# Batch compile all disasm strings with single llvm-mc call
disasm_strs = [(i, t[2]) for i, t in enumerate(to_test) if t[2] is not None]
llvm_results = compile_asm_batch([s for _, s in disasm_strs]) if disasm_strs else []
llvm_map = {i: llvm_results[j] for j, (i, _) in enumerate(disasm_strs)}
# Match results back
passed, failed = 0, 0
failures: list[str] = []
for idx, (asm_text, data, disasm_str, error) in enumerate(to_test):
if error:
failed += 1; failures.append(f"{error} for {data.hex()}")
elif disasm_str is not None and idx in llvm_map:
llvm_bytes = llvm_map[idx]
if llvm_bytes is not None and llvm_bytes == data: passed += 1
elif llvm_bytes is not None: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
if failures[:10]: print(" " + "\n ".join(failures[:10]))
self.assertEqual(failed, 0)
@@ -46,9 +46,10 @@ dev.synchronize()
elapsed = time.perf_counter() - st
self.assertNotEqual(result.returncode, 0, "should have raised")
self.assertIn("NotImplementedError", result.stderr)
self.assertTrue("NotImplementedError" in result.stderr or "ValueError" in result.stderr,
f"expected NotImplementedError or ValueError in stderr")
# Should exit immediately, not wait for the full timeout
self.assertLess(elapsed, 5.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
if __name__ == "__main__":
unittest.main()
+404
View File
@@ -0,0 +1,404 @@
#!/usr/bin/env python3
"""Tests for the RDNA3 pseudocode DSL."""
import unittest
from extra.assembly.amd.pcode import (Reg, TypedView, SliceProxy, MASK32, MASK64,
_f32, _i32, _f16, _i16, f32_to_f16, _isnan, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
BYTE_PERMUTE, v_sad_u8, v_msad_u8)
from extra.assembly.amd.pdf import compile_pseudocode, _expr
from extra.assembly.amd.test.helpers import ExecContext
from extra.assembly.amd.autogen.rdna3.gen_pcode import _VOP3SDOp_V_DIV_SCALE_F32, _VOPCOp_V_CMP_CLASS_F32
class TestReg(unittest.TestCase):
def test_u32_read(self):
r = Reg(0xDEADBEEF)
self.assertEqual(int(r.u32), 0xDEADBEEF)
def test_u32_write(self):
r = Reg(0)
r.u32 = 0x12345678
self.assertEqual(r._val, 0x12345678)
def test_f32_read(self):
r = Reg(0x40400000) # 3.0f
self.assertAlmostEqual(float(r.f32), 3.0)
def test_f32_write(self):
r = Reg(0)
r.f32 = 3.0
self.assertEqual(r._val, 0x40400000)
def test_i32_signed(self):
r = Reg(0xFFFFFFFF) # -1 as signed
self.assertEqual(int(r.i32), -1)
def test_u64(self):
r = Reg(0xDEADBEEFCAFEBABE)
self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
def test_f64(self):
r = Reg(0x4008000000000000) # 3.0 as f64
self.assertAlmostEqual(float(r.f64), 3.0)
class TestTypedView(unittest.TestCase):
def test_bit_slice(self):
r = Reg(0xDEADBEEF)
# Slices return SliceProxy 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] # SliceProxy, 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 TestSliceProxy(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 = Reg(0x3f800000) # 1.0
S1 = Reg(0x40400000) # 3.0
S2 = Reg(0x3f800000) # 1.0 (numerator)
D0, SCC, VCC, EXEC = Reg(0), Reg(0), Reg(0), Reg(0xffffffff)
result = _VOP3SDOp_V_DIV_SCALE_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
# Must always have VCC in result
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
self.assertEqual(result['VCC']._val & 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
S0, S1, S2, D0, SCC, VCC, EXEC = Reg(quiet_nan), Reg(s1_quiet), Reg(0), Reg(0), Reg(0), Reg(0), Reg(0xffffffff)
result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
self.assertEqual(result['D0']._val & 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
S0, S1 = Reg(signal_nan), Reg(s1_signal)
result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
self.assertEqual(result['D0']._val & 1, 1, "Should detect signaling NaN with signaling NaN mask")
# Test that quiet NaN doesn't match signaling NaN mask
S0, S1 = Reg(quiet_nan), Reg(s1_signal)
result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
self.assertEqual(result['D0']._val & 1, 0, "Quiet NaN should not match signaling NaN mask")
# Test that signaling NaN doesn't match quiet NaN mask
S0, S1 = Reg(signal_nan), Reg(s1_quiet)
result = _VOPCOp_V_CMP_CLASS_F32(S0, S1, S2, D0, SCC, VCC, 0, EXEC, 0, None)
self.assertEqual(result['D0']._val & 1, 0, "Signaling NaN should not match quiet NaN mask")
def test_isnan_with_typed_view(self):
"""_isnan must work with TypedView objects, not just Python floats.
Bug: _isnan checked isinstance(x, float) which returned False for TypedView."""
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 SliceProxy.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,10 +1,7 @@
#!/usr/bin/env python3
"""Test that PDF parser correctly extracts format fields."""
import unittest
from extra.assembly.rdna3.autogen import (
SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
SOP1Op, SOP2Op, VOP1Op, VOP3Op
)
import unittest, os
from extra.assembly.amd.autogen.rdna3.ins import SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD, SOP1Op, SOP2Op, VOP1Op, VOP3Op
# expected formats with key fields and whether they have ENCODING
EXPECTED_FORMATS = {
@@ -33,34 +30,32 @@ EXPECTED_FORMATS = {
'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
}
# Skip PDF parsing tests by default - only run with TEST_PDF_PARSER=1
# These are slow (~5s) and only needed when regenerating autogen/
@unittest.skipUnless(os.environ.get("TEST_PDF_PARSER"), "set TEST_PDF_PARSER=1 to run PDF parser tests")
class TestPDFParserGenerate(unittest.TestCase):
"""Test the PDF parser by running generate() and checking results."""
result: dict
@classmethod
def setUpClass(cls):
from extra.assembly.rdna3.gen import generate
cls.result = generate()
def test_pdf_parser(self):
"""Single test that validates all PDF parser outputs."""
from extra.assembly.amd.dsl import generate
result = generate()
def test_all_formats_present(self):
"""All expected formats should be parsed."""
# test_all_formats_present
for fmt_name in EXPECTED_FORMATS:
self.assertIn(fmt_name, self.result["formats"], f"missing format {fmt_name}")
self.assertIn(fmt_name, result["formats"], f"missing format {fmt_name}")
def test_format_count(self):
"""Should have exactly 23 formats."""
self.assertEqual(len(self.result["formats"]), 23)
# test_format_count
self.assertEqual(len(result["formats"]), 23)
def test_no_duplicate_fields(self):
"""No format should have duplicate field names."""
for fmt_name, fields in self.result["formats"].items():
# test_no_duplicate_fields
for fmt_name, fields in result["formats"].items():
field_names = [f[0] for f in fields]
self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
def test_expected_fields(self):
"""Each format should have its expected key fields."""
# test_expected_fields
for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
fields = {f[0] for f in self.result["formats"].get(fmt_name, [])}
fields = {f[0] for f in result["formats"].get(fmt_name, [])}
for field in expected_fields:
self.assertIn(field, fields, f"{fmt_name} missing {field}")
if has_encoding:
@@ -68,21 +63,18 @@ class TestPDFParserGenerate(unittest.TestCase):
else:
self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
def test_vopd_no_dpp16_fields(self):
"""VOPD should not have DPP16-specific fields (parser boundary bug)."""
vopd_fields = {f[0] for f in self.result["formats"].get("VOPD", [])}
# test_vopd_no_dpp16_fields
vopd_fields = {f[0] for f in result["formats"].get("VOPD", [])}
for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
def test_dpp16_no_vinterp_fields(self):
"""DPP16 should not have VINTERP-specific fields."""
dpp16_fields = {f[0] for f in self.result["formats"].get("DPP16", [])}
# test_dpp16_no_vinterp_fields
dpp16_fields = {f[0] for f in result["formats"].get("DPP16", [])}
for field in ['VDST', 'WAITEXP']:
self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
def test_sopp_no_smem_fields(self):
"""SOPP should not have SMEM fields (page break bug)."""
sopp_fields = {f[0] for f in self.result["formats"].get("SOPP", [])}
# test_sopp_no_smem_fields
sopp_fields = {f[0] for f in result["formats"].get("SOPP", [])}
for field in ['SBASE', 'SDATA']:
self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
@@ -1,11 +1,12 @@
#!/usr/bin/env python3
import unittest, subprocess
from extra.assembly.rdna3.autogen import *
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.test.helpers import get_llvm_mc
def llvm_assemble(asm: str) -> bytes:
"""Assemble using llvm-mc and return bytes."""
result = subprocess.run(
["llvm-mc", "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
[get_llvm_mc(), "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
input=asm, capture_output=True, text=True
)
out = b''
+263
View File
@@ -0,0 +1,263 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re, subprocess, os
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.asm import asm
from extra.assembly.amd.asm import detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
old_stdout = sys.stdout
sys.stdout = io.StringIO()
compiler.disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
results = []
for line in output.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if not instr: continue
comment = line.split('//')[1].strip()
if ':' not in comment: continue
hex_str = comment.split(':')[1].strip().split()[0]
try:
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
results.append((instr, machine_bytes))
except ValueError:
continue
return results
def compile_asm(instr: str, compiler=None) -> bytes:
"""Compile a single instruction with llvm-mc and return the machine code bytes."""
llvm_mc = get_llvm_mc()
result = subprocess.run(
[llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=f".text\n{instr}\n", capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed for '{instr}': {result.stderr.strip()}")
# Parse encoding: [0x01,0x39,0x0a,0x7e]
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
return bytes.fromhex(hex_vals)
raise RuntimeError(f"no encoding found in llvm-mc output for: {instr}")
def compile_asm_batch(instrs: list[str]) -> list[bytes]:
"""Compile multiple instructions with a single llvm-mc call."""
if not instrs: return []
llvm_mc = get_llvm_mc()
src = ".text\n" + "\n".join(instrs) + "\n"
result = subprocess.run(
[llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=src, capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
# Parse all encodings in order
encodings = []
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
encodings.append(bytes.fromhex(hex_vals))
if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
return encodings
def compile_and_disasm_batch(instrs: list[str], compiler) -> list[str]:
"""Compile instructions with LLVM and get LLVM's disassembly."""
import tempfile, os
if not instrs: return []
# Build assembly source with all instructions
src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n"
src += "\n".join(f" {instr}" for instr in instrs) + "\n"
# Use llvm-mc to assemble to object file
with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
obj_path = f.name
try:
result = subprocess.run(
[get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-filetype=obj', '-o', obj_path],
input=src, capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
# Disassemble with llvm-objdump
result = subprocess.run([get_llvm_objdump(), '-d', '--mcpu=gfx1100', obj_path], capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
# Parse disassembly output
results: list[str] = []
for line in result.stdout.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if instr: results.append(instr)
return results[:len(instrs)]
finally:
os.unlink(obj_path)
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
def _test_kernel_roundtrip(self, op_fn):
"""Generate kernel from op_fn, test:
1. decode -> reencode matches original bytes
2. asm(disasm()) matches LLVM output
3. our disasm() matches LLVM's disassembly string exactly
"""
from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.compiler_amd import HIPCompiler
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
compiler = HIPCompiler('gfx1100')
# First pass: decode all instructions and collect info
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
for ki, kernel in enumerate(kernels):
offset = 0
while offset < len(kernel.code):
remaining = kernel.code[offset:]
fmt = detect_format(remaining)
if fmt is None:
decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
offset += 4
continue
base_size = fmt._size()
if len(remaining) < base_size:
break
try:
decoded = fmt.from_bytes(remaining) # pass all remaining bytes so from_bytes can read literal
size = decoded.size() # actual size including literal
orig_bytes = remaining[:size]
reencoded = decoded.to_bytes()
our_disasm = decoded.disasm()
decode_ok = reencoded == orig_bytes
decode_err: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
except Exception as e:
decoded_instrs.append((ki, offset, remaining[:base_size], None, None, False, str(e)))
size = base_size
offset += size
# Collect disasm strings for batched LLVM calls - skip unknown opcodes (op_X) that LLVM can't compile
asm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for asm test
disasm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for disasm comparison test
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
if our_disasm is None: continue
# Skip unknown opcodes and malformed instructions for both tests
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
asm_test_instrs.append((idx, our_disasm))
disasm_test_instrs.append((idx, our_disasm))
# Batch compile for asm test
asm_llvm_results = compile_asm_batch([d for _, d in asm_test_instrs])
asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
# Batch compile+disasm for disasm comparison test
disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], compiler)
disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
# Now evaluate results
decode_passed, decode_failed, decode_skipped = 0, 0, 0
asm_passed, asm_failed, asm_skipped = 0, 0, 0
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
decode_failures: list[str] = []
asm_failures: list[str] = []
disasm_failures: list[str] = []
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
# Decode test
if decode_ok:
decode_passed += 1
elif decode_err == "no format":
decode_skipped += 1
else:
decode_failed += 1
decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
# Asm test
if our_disasm is None:
asm_skipped += 1
elif idx in asm_llvm_map:
llvm_bytes = asm_llvm_map[idx]
try:
our_bytes = asm(our_disasm).to_bytes()
if our_bytes[:len(llvm_bytes)] == llvm_bytes:
asm_passed += 1
else:
asm_failed += 1
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
except Exception:
asm_skipped += 1
else:
asm_skipped += 1
# Disasm comparison test
if our_disasm is None:
disasm_skipped += 1
elif idx in disasm_llvm_map:
llvm_disasm = disasm_llvm_map[idx]
if our_disasm == llvm_disasm:
disasm_passed += 1
else:
disasm_failed += 1
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
else:
disasm_skipped += 1
print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
# Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
# Basic unary ops
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Reductions
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
# Matmul
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
# Complex ops
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
# Cast ops
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
# Comparison ops
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
if __name__ == "__main__":
unittest.main()
-254
View File
@@ -1,254 +0,0 @@
# Pure combinational ALU functions for RDNA3 emulation
from __future__ import annotations
import struct, math
from typing import Callable
from extra.assembly.rdna3.autogen import SOP1Op, SOP2Op, SOPCOp, SOPKOp, VOP1Op, VOP2Op, VOP3Op
# Format base offsets for unified opcode space
SOP2_BASE, SOP1_BASE, SOPC_BASE, SOPK_BASE = 0x000, 0x100, 0x200, 0x300
VOP2_BASE, VOP1_BASE = 0x100, 0x180
# Float conversion helpers
_I, _f, _H, _e = struct.Struct('<I'), struct.Struct('<f'), struct.Struct('<H'), struct.Struct('<e')
def f32(i: int) -> float: return _f.unpack(_I.pack(i & 0xffffffff))[0]
def i32(f: float) -> int:
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
try: return _I.unpack(_f.pack(f))[0]
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
def f16(i: int) -> float: return _e.unpack(_H.pack(i & 0xffff))[0]
def i16(f: float) -> int:
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
try: return _H.unpack(_e.pack(f))[0]
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
def sext(v: int, b: int) -> int: return v - (1 << b) if v & (1 << (b-1)) else v
def clz(x: int) -> int: return 32 - x.bit_length() if x else 32
def cls(x: int) -> int: x &= 0xffffffff; return 31 if x in (0, 0xffffffff) else clz(~x & 0xffffffff if x >> 31 else x) - 1
def _cvt_i32_f32(v): return (0x7fffffff if v > 0 else 0x80000000) if math.isinf(v) else (0 if math.isnan(v) else max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff)
def _cvt_u32_f32(v): return (0xffffffff if v > 0 else 0) if math.isinf(v) else (0 if math.isnan(v) or v < 0 else min(0xffffffff, int(v)))
# SALU: op -> fn(s0, s1, scc_in) -> (result, scc_out)
SALU: dict[int, Callable] = {
# SOP2
SOP2_BASE + SOP2Op.S_ADD_U32: lambda a, b, scc: ((a + b) & 0xffffffff, int((a + b) >= 0x100000000)),
SOP2_BASE + SOP2Op.S_SUB_U32: lambda a, b, scc: ((a - b) & 0xffffffff, int(b > a)),
SOP2_BASE + SOP2Op.S_ADDC_U32: lambda a, b, scc: ((r := a + b + scc) & 0xffffffff, int(r >= 0x100000000)),
SOP2_BASE + SOP2Op.S_SUBB_U32: lambda a, b, scc: ((a - b - scc) & 0xffffffff, int((b + scc) > a)),
SOP2_BASE + SOP2Op.S_ADD_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 32)) & 0xffffffff, int(((a >> 31) == (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
SOP2_BASE + SOP2Op.S_SUB_I32: lambda a, b, scc: ((r := sext(a, 32) - sext(b, 32)) & 0xffffffff, int(((a >> 31) != (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
SOP2_BASE + SOP2Op.S_AND_B32: lambda a, b, scc: ((r := a & b), int(r != 0)),
SOP2_BASE + SOP2Op.S_OR_B32: lambda a, b, scc: ((r := a | b), int(r != 0)),
SOP2_BASE + SOP2Op.S_XOR_B32: lambda a, b, scc: ((r := a ^ b), int(r != 0)),
SOP2_BASE + SOP2Op.S_AND_NOT1_B32: lambda a, b, scc: ((r := a & (~b & 0xffffffff)), int(r != 0)),
SOP2_BASE + SOP2Op.S_OR_NOT1_B32: lambda a, b, scc: ((r := a | (~b & 0xffffffff)), int(r != 0)),
SOP2_BASE + SOP2Op.S_LSHL_B32: lambda a, b, scc: ((r := (a << (b & 0x1f)) & 0xffffffff), int(r != 0)),
SOP2_BASE + SOP2Op.S_LSHR_B32: lambda a, b, scc: ((r := a >> (b & 0x1f)), int(r != 0)),
SOP2_BASE + SOP2Op.S_ASHR_I32: lambda a, b, scc: ((r := sext(a, 32) >> (b & 0x1f)) & 0xffffffff, int(r != 0)),
SOP2_BASE + SOP2Op.S_MUL_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 32)) & 0xffffffff, scc),
SOP2_BASE + SOP2Op.S_MUL_HI_U32: lambda a, b, scc: (((a * b) >> 32) & 0xffffffff, scc),
SOP2_BASE + SOP2Op.S_MUL_HI_I32: lambda a, b, scc: (((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff, scc),
SOP2_BASE + SOP2Op.S_MIN_I32: lambda a, b, scc: (a, 1) if sext(a, 32) < sext(b, 32) else (b, 0),
SOP2_BASE + SOP2Op.S_MIN_U32: lambda a, b, scc: (a, 1) if a < b else (b, 0),
SOP2_BASE + SOP2Op.S_MAX_I32: lambda a, b, scc: (a, 1) if sext(a, 32) > sext(b, 32) else (b, 0),
SOP2_BASE + SOP2Op.S_MAX_U32: lambda a, b, scc: (a, 1) if a > b else (b, 0),
SOP2_BASE + SOP2Op.S_CSELECT_B32: lambda a, b, scc: (a if scc else b, scc),
SOP2_BASE + SOP2Op.S_BFE_U32: lambda a, b, scc: ((r := ((a >> (b & 0x1f)) & ((1 << ((b >> 16) & 0x7f)) - 1)) if (b >> 16) & 0x7f else 0), int(r != 0)),
SOP2_BASE + SOP2Op.S_BFE_I32: lambda a, b, scc: ((r := sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := (b >> 16) & 0x7f) else 0), int(r != 0)),
SOP2_BASE + SOP2Op.S_PACK_LL_B32_B16: lambda a, b, scc: ((a & 0xffff) | ((b & 0xffff) << 16), scc),
SOP2_BASE + SOP2Op.S_PACK_LH_B32_B16: lambda a, b, scc: ((a & 0xffff) | (b & 0xffff0000), scc),
SOP2_BASE + SOP2Op.S_PACK_HH_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | (b & 0xffff0000), scc),
SOP2_BASE + SOP2Op.S_PACK_HL_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | ((b & 0xffff) << 16), scc),
SOP2_BASE + SOP2Op.S_ADD_F32: lambda a, b, scc: (i32(f32(a) + f32(b)), scc),
SOP2_BASE + SOP2Op.S_SUB_F32: lambda a, b, scc: (i32(f32(a) - f32(b)), scc),
SOP2_BASE + SOP2Op.S_MUL_F32: lambda a, b, scc: (i32(f32(a) * f32(b)), scc),
# SOP1
SOP1_BASE + SOP1Op.S_MOV_B32: lambda a, b, scc: (a, scc),
SOP1_BASE + SOP1Op.S_NOT_B32: lambda a, b, scc: ((r := (~a) & 0xffffffff), int(r != 0)),
SOP1_BASE + SOP1Op.S_BREV_B32: lambda a, b, scc: (int(f'{a & 0xffffffff:032b}'[::-1], 2), scc),
SOP1_BASE + SOP1Op.S_CLZ_I32_U32: lambda a, b, scc: (clz(a), scc),
SOP1_BASE + SOP1Op.S_CLS_I32: lambda a, b, scc: (cls(a), scc),
SOP1_BASE + SOP1Op.S_SEXT_I32_I8: lambda a, b, scc: (sext(a & 0xff, 8) & 0xffffffff, scc),
SOP1_BASE + SOP1Op.S_SEXT_I32_I16: lambda a, b, scc: (sext(a & 0xffff, 16) & 0xffffffff, scc),
SOP1_BASE + SOP1Op.S_ABS_I32: lambda a, b, scc: ((r := abs(sext(a, 32)) & 0xffffffff), int(r != 0)),
SOP1_BASE + SOP1Op.S_CVT_F32_I32: lambda a, b, scc: (i32(float(sext(a, 32))), scc),
SOP1_BASE + SOP1Op.S_CVT_F32_U32: lambda a, b, scc: (i32(float(a)), scc),
SOP1_BASE + SOP1Op.S_CVT_I32_F32: lambda a, b, scc: (_cvt_i32_f32(f32(a)), scc),
SOP1_BASE + SOP1Op.S_CVT_U32_F32: lambda a, b, scc: (_cvt_u32_f32(f32(a)), scc),
SOP1_BASE + SOP1Op.S_CEIL_F32: lambda a, b, scc: (i32(math.ceil(f32(a))), scc),
SOP1_BASE + SOP1Op.S_FLOOR_F32: lambda a, b, scc: (i32(math.floor(f32(a))), scc),
SOP1_BASE + SOP1Op.S_TRUNC_F32: lambda a, b, scc: (i32(math.trunc(f32(a))), scc),
SOP1_BASE + SOP1Op.S_RNDNE_F32: lambda a, b, scc: (i32(round(f32(a))), scc),
SOP1_BASE + SOP1Op.S_CVT_F16_F32: lambda a, b, scc: (i16(f32(a)), scc),
SOP1_BASE + SOP1Op.S_CVT_F32_F16: lambda a, b, scc: (i32(f16(a)), scc),
# SOPC
SOPC_BASE + SOPCOp.S_CMP_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_EQ_U32: lambda a, b, scc: (0, int(a == b)),
SOPC_BASE + SOPCOp.S_CMP_LG_U32: lambda a, b, scc: (0, int(a != b)),
SOPC_BASE + SOPCOp.S_CMP_GT_U32: lambda a, b, scc: (0, int(a > b)),
SOPC_BASE + SOPCOp.S_CMP_GE_U32: lambda a, b, scc: (0, int(a >= b)),
SOPC_BASE + SOPCOp.S_CMP_LT_U32: lambda a, b, scc: (0, int(a < b)),
SOPC_BASE + SOPCOp.S_CMP_LE_U32: lambda a, b, scc: (0, int(a <= b)),
SOPC_BASE + SOPCOp.S_BITCMP0_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) == 0)),
SOPC_BASE + SOPCOp.S_BITCMP1_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) != 0)),
# SOPK
SOPK_BASE + SOPKOp.S_MOVK_I32: lambda a, b, scc: (sext(b, 16) & 0xffffffff, scc),
SOPK_BASE + SOPKOp.S_CMOVK_I32: lambda a, b, scc: ((sext(b, 16) & 0xffffffff) if scc else a, scc),
SOPK_BASE + SOPKOp.S_ADDK_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 16)) & 0xffffffff, int(((a >> 31) == ((b >> 15) & 1)) and ((a >> 31) != ((r >> 31) & 1)))),
SOPK_BASE + SOPKOp.S_MULK_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 16)) & 0xffffffff, scc),
SOPK_BASE + SOPKOp.S_CMPK_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_EQ_U32: lambda a, b, scc: (0, int(a == (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_LG_U32: lambda a, b, scc: (0, int(a != (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_GT_U32: lambda a, b, scc: (0, int(a > (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_GE_U32: lambda a, b, scc: (0, int(a >= (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_LT_U32: lambda a, b, scc: (0, int(a < (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_LE_U32: lambda a, b, scc: (0, int(a <= (b & 0xffff))),
}
# VALU: op -> fn(s0, s1, s2) -> result
VALU: dict[int, Callable] = {
# VOP2
VOP2_BASE + VOP2Op.V_ADD_F32: lambda a, b, c: i32(f32(a) + f32(b)),
VOP2_BASE + VOP2Op.V_SUB_F32: lambda a, b, c: i32(f32(a) - f32(b)),
VOP2_BASE + VOP2Op.V_SUBREV_F32: lambda a, b, c: i32(f32(b) - f32(a)),
VOP2_BASE + VOP2Op.V_MUL_F32: lambda a, b, c: i32(f32(a) * f32(b)),
VOP2_BASE + VOP2Op.V_MIN_F32: lambda a, b, c: i32(min(f32(a), f32(b))),
VOP2_BASE + VOP2Op.V_MAX_F32: lambda a, b, c: i32(max(f32(a), f32(b))),
VOP2_BASE + VOP2Op.V_ADD_NC_U32: lambda a, b, c: (a + b) & 0xffffffff,
VOP2_BASE + VOP2Op.V_SUB_NC_U32: lambda a, b, c: (a - b) & 0xffffffff,
VOP2_BASE + VOP2Op.V_SUBREV_NC_U32: lambda a, b, c: (b - a) & 0xffffffff,
VOP2_BASE + VOP2Op.V_AND_B32: lambda a, b, c: a & b,
VOP2_BASE + VOP2Op.V_OR_B32: lambda a, b, c: a | b,
VOP2_BASE + VOP2Op.V_XOR_B32: lambda a, b, c: a ^ b,
VOP2_BASE + VOP2Op.V_XNOR_B32: lambda a, b, c: (~(a ^ b)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_LSHLREV_B32: lambda a, b, c: (b << (a & 0x1f)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_LSHRREV_B32: lambda a, b, c: b >> (a & 0x1f),
VOP2_BASE + VOP2Op.V_ASHRREV_I32: lambda a, b, c: (sext(b, 32) >> (a & 0x1f)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MIN_I32: lambda a, b, c: a if sext(a, 32) < sext(b, 32) else b,
VOP2_BASE + VOP2Op.V_MAX_I32: lambda a, b, c: a if sext(a, 32) > sext(b, 32) else b,
VOP2_BASE + VOP2Op.V_MIN_U32: lambda a, b, c: min(a, b),
VOP2_BASE + VOP2Op.V_MAX_U32: lambda a, b, c: max(a, b),
VOP2_BASE + VOP2Op.V_MUL_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MUL_HI_I32_I24: lambda a, b, c: ((sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) >> 32) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MUL_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MUL_HI_U32_U24: lambda a, b, c: (((a & 0xffffff) * (b & 0xffffff)) >> 32) & 0xffffffff,
VOP2_BASE + VOP2Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
VOP2_BASE + VOP2Op.V_LDEXP_F16: lambda a, b, c: i16(math.ldexp(f16(a), sext(b, 32))),
VOP2_BASE + VOP2Op.V_ADD_F16: lambda a, b, c: i16(f16(a) + f16(b)),
VOP2_BASE + VOP2Op.V_SUB_F16: lambda a, b, c: i16(f16(a) - f16(b)),
VOP2_BASE + VOP2Op.V_MUL_F16: lambda a, b, c: i16(f16(a) * f16(b)),
VOP2_BASE + VOP2Op.V_MIN_F16: lambda a, b, c: i16(min(f16(a), f16(b))),
VOP2_BASE + VOP2Op.V_MAX_F16: lambda a, b, c: i16(max(f16(a), f16(b))),
# VOP1
VOP1_BASE + VOP1Op.V_MOV_B32: lambda a, b, c: a,
VOP1_BASE + VOP1Op.V_NOT_B32: lambda a, b, c: (~a) & 0xffffffff,
VOP1_BASE + VOP1Op.V_BFREV_B32: lambda a, b, c: int(f'{a & 0xffffffff:032b}'[::-1], 2),
VOP1_BASE + VOP1Op.V_CLZ_I32_U32: lambda a, b, c: clz(a),
VOP1_BASE + VOP1Op.V_CLS_I32: lambda a, b, c: cls(a),
VOP1_BASE + VOP1Op.V_CVT_F32_I32: lambda a, b, c: i32(float(sext(a, 32))),
VOP1_BASE + VOP1Op.V_CVT_F32_U32: lambda a, b, c: i32(float(a)),
VOP1_BASE + VOP1Op.V_CVT_I32_F32: lambda a, b, c: _cvt_i32_f32(f32(a)),
VOP1_BASE + VOP1Op.V_CVT_U32_F32: lambda a, b, c: _cvt_u32_f32(f32(a)),
VOP1_BASE + VOP1Op.V_CVT_F16_F32: lambda a, b, c: i16(f32(a)),
VOP1_BASE + VOP1Op.V_CVT_F32_F16: lambda a, b, c: i32(f16(a)),
VOP1_BASE + VOP1Op.V_RCP_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
VOP1_BASE + VOP1Op.V_RCP_IFLAG_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
VOP1_BASE + VOP1Op.V_RSQ_F32: lambda a, b, c: i32(1.0 / math.sqrt(f32(a)) if f32(a) > 0 else (float('nan') if f32(a) < 0 else float('inf'))),
VOP1_BASE + VOP1Op.V_SQRT_F32: lambda a, b, c: i32(math.sqrt(f32(a)) if f32(a) >= 0 else float('nan')),
VOP1_BASE + VOP1Op.V_LOG_F32: lambda a, b, c: i32(math.log2(f32(a)) if f32(a) > 0 else (float('-inf') if f32(a) == 0 else float('nan'))),
VOP1_BASE + VOP1Op.V_EXP_F32: lambda a, b, c: i32(float('inf') if f32(a) > 128 else (0.0 if f32(a) < -150 else math.pow(2.0, f32(a)))),
VOP1_BASE + VOP1Op.V_SIN_F32: lambda a, b, c: i32(math.sin(f32(a) * 2 * math.pi)),
VOP1_BASE + VOP1Op.V_COS_F32: lambda a, b, c: i32(math.cos(f32(a) * 2 * math.pi)),
VOP1_BASE + VOP1Op.V_FLOOR_F32: lambda a, b, c: i32(math.floor(f32(a))),
VOP1_BASE + VOP1Op.V_CEIL_F32: lambda a, b, c: i32(math.ceil(f32(a))),
VOP1_BASE + VOP1Op.V_TRUNC_F32: lambda a, b, c: i32(math.trunc(f32(a))),
VOP1_BASE + VOP1Op.V_RNDNE_F32: lambda a, b, c: i32(round(f32(a))),
VOP1_BASE + VOP1Op.V_FRACT_F32: lambda a, b, c: i32((v := f32(a)) - math.floor(v)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE0: lambda a, b, c: i32(float(a & 0xff)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE1: lambda a, b, c: i32(float((a >> 8) & 0xff)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE2: lambda a, b, c: i32(float((a >> 16) & 0xff)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE3: lambda a, b, c: i32(float((a >> 24) & 0xff)),
VOP1_BASE + VOP1Op.V_FREXP_MANT_F32: lambda a, b, c: i32(math.frexp(v)[0] if (v := f32(a)) != 0 else 0.0),
VOP1_BASE + VOP1Op.V_FREXP_EXP_I32_F32: lambda a, b, c: (math.frexp(v)[1] if (v := f32(a)) != 0 else 0) & 0xffffffff,
# VOP3
VOP3Op.V_FMA_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
VOP3Op.V_DIV_FMAS_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
VOP3Op.V_ADD3_U32: lambda a, b, c: (a + b + c) & 0xffffffff,
VOP3Op.V_LSHL_ADD_U32: lambda a, b, c: ((a << (b & 0x1f)) + c) & 0xffffffff,
VOP3Op.V_ADD_LSHL_U32: lambda a, b, c: ((a + b) << (c & 0x1f)) & 0xffffffff,
VOP3Op.V_XOR3_B32: lambda a, b, c: a ^ b ^ c,
VOP3Op.V_OR3_B32: lambda a, b, c: a | b | c,
VOP3Op.V_AND_OR_B32: lambda a, b, c: (a & b) | c,
VOP3Op.V_LSHL_OR_B32: lambda a, b, c: ((a << (b & 0x1f)) | c) & 0xffffffff,
VOP3Op.V_XAD_U32: lambda a, b, c: ((a ^ b) + c) & 0xffffffff,
VOP3Op.V_MAD_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff) + c) & 0xffffffff,
VOP3Op.V_MAD_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24) + sext(c, 32)) & 0xffffffff,
VOP3Op.V_BFE_U32: lambda a, b, c: (a >> (b & 0x1f)) & ((1 << (c & 0x1f)) - 1) if c & 0x1f else 0,
VOP3Op.V_BFE_I32: lambda a, b, c: sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := c & 0x1f) else 0,
VOP3Op.V_ALIGNBIT_B32: lambda a, b, c: (((a << 32) | b) >> (c & 0x1f)) & 0xffffffff,
VOP3Op.V_MUL_LO_U32: lambda a, b, c: (a * b) & 0xffffffff,
VOP3Op.V_MUL_HI_U32: lambda a, b, c: ((a * b) >> 32) & 0xffffffff,
VOP3Op.V_MUL_HI_I32: lambda a, b, c: ((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff,
VOP3Op.V_LDEXP_F32: lambda a, b, c: i32(math.ldexp(f32(a), sext(b, 32))),
VOP3Op.V_DIV_FIXUP_F32: lambda a, b, c: i32(math.copysign(float('inf'), f32(c)) if f32(b) == 0.0 else f32(c) / f32(b)),
VOP3Op.V_PACK_B32_F16: lambda a, b, c: (a & 0xffff) | ((b & 0xffff) << 16),
VOP3Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
VOP3Op.V_LSHLREV_B16: lambda a, b, c: ((b & 0xffff) << (a & 0xf)) & 0xffff,
VOP3Op.V_LSHRREV_B16: lambda a, b, c: (b & 0xffff) >> (a & 0xf),
VOP3Op.V_ASHRREV_I16: lambda a, b, c: (sext(b & 0xffff, 16) >> (a & 0xf)) & 0xffff,
VOP3Op.V_ADD_NC_U16: lambda a, b, c: ((a & 0xffff) + (b & 0xffff)) & 0xffff,
VOP3Op.V_SUB_NC_U16: lambda a, b, c: ((a & 0xffff) - (b & 0xffff)) & 0xffff,
VOP3Op.V_MUL_LO_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff)) & 0xffff,
VOP3Op.V_MIN_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff),
VOP3Op.V_MAX_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff),
VOP3Op.V_MIN_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) < sext(b & 0xffff, 16) else (b & 0xffff),
VOP3Op.V_MAX_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) > sext(b & 0xffff, 16) else (b & 0xffff),
VOP3Op.V_MAD_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff) + (c & 0xffff)) & 0xffff,
VOP3Op.V_MAD_I16: lambda a, b, c: (sext(a & 0xffff, 16) * sext(b & 0xffff, 16) + sext(c & 0xffff, 16)) & 0xffff,
VOP3Op.V_FMA_F16: lambda a, b, c: i16(f16(a) * f16(b) + f16(c)),
VOP3Op.V_MIN3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[0] & 0xffffffff,
VOP3Op.V_MAX3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[2] & 0xffffffff,
VOP3Op.V_MED3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[1] & 0xffffffff,
VOP3Op.V_MIN3_F16: lambda a, b, c: i16(min(f16(a), f16(b), f16(c))),
VOP3Op.V_MAX3_F16: lambda a, b, c: i16(max(f16(a), f16(b), f16(c))),
VOP3Op.V_MED3_F16: lambda a, b, c: i16(sorted([f16(a), f16(b), f16(c)])[1]),
VOP3Op.V_MIN3_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff, c & 0xffff),
VOP3Op.V_MAX3_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff, c & 0xffff),
VOP3Op.V_MED3_U16: lambda a, b, c: sorted([a & 0xffff, b & 0xffff, c & 0xffff])[1],
VOP3Op.V_MIN3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[0] & 0xffff,
VOP3Op.V_MAX3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[2] & 0xffff,
VOP3Op.V_MED3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[1] & 0xffff,
}
def _cmp8(a, b): return [False, a < b, a == b, a <= b, a > b, a != b, a >= b, True]
def _cmp6(a, b): return [a < b, a == b, a <= b, a > b, a != b, a >= b]
def vopc(op: int, s0: int, s1: int, s0_hi: int = 0, s1_hi: int = 0) -> int:
base = op & 0x7f
if 16 <= base <= 31: # F32
f0, f1, cmp, nan = f32(s0), f32(s1), base - 16, math.isnan(f32(s0)) or math.isnan(f32(s1))
return int([False, f0<f1, f0==f1, f0<=f1, f0>f1, f0!=f1, f0>=f1, not nan, nan, f0<f1 or nan, f0==f1 or nan, f0<=f1 or nan, f0>f1 or nan, f0!=f1 or nan, f0>=f1 or nan, True][cmp])
if 49 <= base <= 54: return int(_cmp6(sext(s0 & 0xffff, 16), sext(s1 & 0xffff, 16))[base - 49]) # I16
if 57 <= base <= 62: return int(_cmp6(s0 & 0xffff, s1 & 0xffff)[base - 57]) # U16
if 64 <= base <= 79: # I32/U32
cmp = (base - 64) % 8
return int(_cmp8(sext(s0, 32), sext(s1, 32))[cmp] if base < 72 else _cmp8(s0, s1)[cmp])
if 80 <= base <= 95: # I64/U64
s0_64, s1_64 = s0 | (s0_hi << 32), s1 | (s1_hi << 32)
return int(_cmp8(sext(s0_64, 64), sext(s1_64, 64))[(base - 80) % 8] if base < 88 else _cmp8(s0_64, s1_64)[(base - 80) % 8])
if base == 126: # CLASS_F32
f, mask = f32(s0), s1
if math.isnan(f): return int(bool(mask & 0x3))
if math.isinf(f): return int(bool(mask & (0x4 if f < 0 else 0x200)))
if f == 0.0: return int(bool(mask & (0x20 if (s0 >> 31) & 1 else 0x40)))
exp, sign = (s0 >> 23) & 0xff, (s0 >> 31) & 1
return int(bool(mask & ((0x10 if sign else 0x80) if exp == 0 else (0x8 if sign else 0x100))))
raise NotImplementedError(f"VOPC op {op} (base {base})")
-600
View File
@@ -1,600 +0,0 @@
# RDNA3 assembler and disassembler
from __future__ import annotations
import re
from extra.assembly.rdna3.lib import Inst, RawImm, Reg, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
# Decoding helpers
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
SPECIAL_PAIRS = {106: "vcc", 126: "exec"} # Special register pairs (for 64-bit ops)
# GFX11 hwreg names (IDs 16-17 are TBA - not supported, IDs 18-19 are PERF_SNAPSHOT)
HWREG_NAMES = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI',
22: 'HW_REG_XNACK_MASK', 23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
HWREG_IDS = {v.lower(): k for k, v in HWREG_NAMES.items()} # Reverse map for assembler
MSG_NAMES = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
_16BIT_TYPES = ('f16', 'i16', 'u16', 'b16')
def _is_16bit(s: str) -> bool: return any(s.endswith(x) for x in _16BIT_TYPES)
def decode_src(val: int) -> str:
if val <= 105: return f"s{val}"
if val in SPECIAL_DEC: return SPECIAL_DEC[val]
if 108 <= val <= 123: return f"ttmp{val - 108}"
if 128 <= val <= 192: return str(val - 128)
if 193 <= val <= 208: return str(-(val - 192))
if 256 <= val <= 511: return f"v{val - 256}"
return "lit" if val == 255 else f"?{val}"
def _reg(prefix: str, base: int, cnt: int = 1) -> str: return f"{prefix}{base}" if cnt == 1 else f"{prefix}[{base}:{base+cnt-1}]"
def _sreg(base: int, cnt: int = 1) -> str: return _reg("s", base, cnt)
def _vreg(base: int, cnt: int = 1) -> str: return _reg("v", base, cnt)
def _fmt_sdst(v: int, cnt: int = 1) -> str:
"""Format SGPR destination with special register names."""
if v == 124: return "null"
if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
if cnt > 1 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if cnt > 1: return _sreg(v, cnt)
return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
def _fmt_ssrc(v: int, cnt: int = 1) -> str:
"""Format SGPR source with special register names and pairs."""
if cnt == 2:
if v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if v <= 105: return _sreg(v, 2)
if 108 <= v <= 123: return _reg("ttmp", v - 108, 2)
return decode_src(v)
def _fmt_src_n(v: int, cnt: int) -> str:
"""Format source with given register count (1, 2, or 4)."""
if cnt == 1: return decode_src(v)
if v >= 256: return _vreg(v - 256, cnt)
if v <= 105: return _sreg(v, cnt)
if cnt == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
return decode_src(v)
def _fmt_src64(v: int) -> str:
"""Format 64-bit source (VGPR pair, SGPR pair, or special pair)."""
return _fmt_src_n(v, 2)
def _parse_sop_sizes(op_name: str) -> tuple[int, ...]:
"""Parse dst and src sizes from SOP instruction name. Returns (dst_cnt, src0_cnt) or (dst_cnt, src0_cnt, src1_cnt)."""
if op_name in ('s_bitset0_b64', 's_bitset1_b64'): return (2, 1)
if op_name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return (2, 2, 1)
if op_name in ('s_bfm_b64',): return (2, 1, 1)
# SOPC: s_bitcmp0_b64, s_bitcmp1_b64 - 64-bit src0, 32-bit src1 (bit index)
if op_name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return (1, 2, 1)
if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', op_name):
return (2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1)
if m := re.search(r'_(b|i|u)(32|64)$', op_name):
sz = 2 if m.group(2) == '64' else 1
return (sz, sz)
return (1, 1)
# Waitcnt helpers (RDNA3 format: bits 15:10=vmcnt, bits 9:4=lgkmcnt, bits 3:0=expcnt)
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
def decode_waitcnt(val: int) -> tuple[int, int, int]:
return (val >> 10) & 0x3f, val & 0xf, (val >> 4) & 0x3f # vmcnt, expcnt, lgkmcnt
# VOP3SD opcodes (shared encoding with VOP3 but different field layout)
# Note: opcodes 0-255 are VOPC promoted to VOP3 - never treat as VOP3SD
VOP3SD_OPCODES = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
# Disassembler
def disasm(inst: Inst) -> str:
op_val = unwrap(inst._values.get('op', 0))
cls_name = inst.__class__.__name__
# VOP3 and VOP3SD share encoding - check opcode to determine which
is_vop3sd = cls_name == 'VOP3' and op_val in VOP3SD_OPCODES
try:
from extra.assembly.rdna3 import autogen
if is_vop3sd:
op_name = autogen.VOP3SDOp(op_val).name.lower()
else:
op_name = getattr(autogen, f"{cls_name}Op")(op_val).name.lower() if hasattr(autogen, f"{cls_name}Op") else f"op_{op_val}"
except (ValueError, KeyError): op_name = f"op_{op_val}"
def fmt_src(v): return f"0x{inst._literal:x}" if v == 255 and getattr(inst, '_literal', None) else decode_src(v)
# VOP1
if cls_name == 'VOP1':
vdst, src0 = unwrap(inst._values['vdst']), unwrap(inst._values['src0'])
if op_name == 'v_nop': return 'v_nop'
if op_name == 'v_pipeflush': return 'v_pipeflush'
parts = op_name.split('_')
is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name
_F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
is_f64_src = op_name in _F64_OPS or op_name in ('v_cvt_f32_f64', 'v_cvt_i32_f64', 'v_cvt_u32_f64', 'v_frexp_exp_i32_f64')
if op_name == 'v_readfirstlane_b32':
return f"v_readfirstlane_b32 {decode_src(vdst)}, v{src0 - 256 if src0 >= 256 else src0}"
dst_str = _vreg(vdst, 2) if is_f64_dst else f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}" if is_16bit_dst else f"v{vdst}"
src_str = _fmt_src64(src0) if is_f64_src else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit_src and src0 >= 256 else fmt_src(src0)
return f"{op_name}_e32 {dst_str}, {src_str}"
# VOP2
if cls_name == 'VOP2':
vdst, src0_raw, vsrc1 = unwrap(inst._values['vdst']), unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
suffix = "" if op_name == "v_dot2acc_f32_f16" else "_e32"
is_16bit_op = ('_f16' in op_name or '_i16' in op_name or '_u16' in op_name) and '_f32' not in op_name and '_i32' not in op_name and 'pk_' not in op_name
if is_16bit_op:
dst_str = f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}"
src0_str = f"v{(src0_raw - 256) & 0x7f}.{'h' if src0_raw >= 384 else 'l'}" if src0_raw >= 256 else fmt_src(src0_raw)
vsrc1_str = f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}"
else:
dst_str, src0_str, vsrc1_str = f"v{vdst}", fmt_src(src0_raw), f"v{vsrc1}"
return f"{op_name}{suffix} {dst_str}, {src0_str}, {vsrc1_str}" + (", vcc_lo" if op_name == "v_cndmask_b32" else "")
# VOPC
if cls_name == 'VOPC':
src0, vsrc1 = unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
is_64bit = any(x in op_name for x in ('f64', 'i64', 'u64'))
is_64bit_vsrc1 = is_64bit and 'class' not in op_name
is_16bit = any(x in op_name for x in ('_f16', '_i16', '_u16')) and 'f32' not in op_name
is_cmpx = op_name.startswith('v_cmpx') # VOPCX writes to exec, no vcc destination
src0_str = _fmt_src64(src0) if is_64bit else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit and src0 >= 256 else fmt_src(src0)
vsrc1_str = _vreg(vsrc1, 2) if is_64bit_vsrc1 else f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}" if is_16bit else f"v{vsrc1}"
return f"{op_name}_e32 {src0_str}, {vsrc1_str}" if is_cmpx else f"{op_name}_e32 vcc_lo, {src0_str}, {vsrc1_str}"
# SOPP
if cls_name == 'SOPP':
simm16 = unwrap(inst._values.get('simm16', 0))
# No-operand instructions (simm16 is ignored)
no_imm_ops = ('s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_ttracedata_imm',
's_wait_idle', 's_endpgm_saved', 's_code_end', 's_endpgm_ordered_ps_done')
if op_name in no_imm_ops: return op_name
if op_name == 's_waitcnt':
vmcnt, expcnt, lgkmcnt = decode_waitcnt(simm16)
parts = []
if vmcnt != 0x3f: parts.append(f"vmcnt({vmcnt})")
if expcnt != 0x7: parts.append(f"expcnt({expcnt})")
if lgkmcnt != 0x3f: parts.append(f"lgkmcnt({lgkmcnt})")
return f"s_waitcnt {' '.join(parts)}" if parts else "s_waitcnt 0"
if op_name == 's_delay_alu':
dep_names = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
skip_names = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = simm16 & 0xf, (simm16 >> 4) & 0x7, (simm16 >> 7) & 0xf
def dep_name(v): return dep_names[v-1] if 0 < v <= len(dep_names) else str(v)
parts = [f"instid0({dep_name(id0)})"] if id0 else []
if skip: parts.append(f"instskip({skip_names[skip]})")
if id1: parts.append(f"instid1({dep_name(id1)})")
return f"s_delay_alu {' | '.join(p for p in parts if p)}" if parts else "s_delay_alu 0"
if op_name.startswith('s_cbranch') or op_name.startswith('s_branch'):
return f"{op_name} {simm16}"
# Most SOPP ops require immediate (s_nop, s_setkill, s_sethalt, s_sleep, s_setprio, s_sendmsg*, etc.)
return f"{op_name} 0x{simm16:x}"
# SMEM
if cls_name == 'SMEM':
if op_name in ('s_gl1_inv', 's_dcache_inv'): return op_name
sdata, sbase, soffset, offset = unwrap(inst._values['sdata']), unwrap(inst._values['sbase']), unwrap(inst._values['soffset']), unwrap(inst._values.get('offset', 0))
glc, dlc = unwrap(inst._values.get('glc', 0)), unwrap(inst._values.get('dlc', 0))
# Format offset: "soffset offset:X" if both, "0x{offset:x}" if only imm, or decode_src(soffset)
off_str = f"{decode_src(soffset)} offset:0x{offset:x}" if offset and soffset != 124 else f"0x{offset:x}" if offset else decode_src(soffset)
sbase_idx, sbase_cnt = sbase * 2, 4 if (8 <= op_val <= 12 or op_name == 's_atc_probe_buffer') else 2
sbase_str = _fmt_ssrc(sbase_idx, sbase_cnt) if sbase_cnt == 2 else _sreg(sbase_idx, sbase_cnt) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_cnt)
if op_name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{op_name} {sdata}, {sbase_str}, {off_str}"
width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val, 1)
mods = [m for m in ["glc" if glc else "", "dlc" if dlc else ""] if m]
return f"{op_name} {_fmt_sdst(sdata, width)}, {sbase_str}, {off_str}" + (" " + " ".join(mods) if mods else "")
# FLAT
if cls_name == 'FLAT':
vdst, addr, data, saddr, offset, seg = [unwrap(inst._values.get(f, 0)) for f in ['vdst', 'addr', 'data', 'saddr', 'offset', 'seg']]
instr = f"{['flat', 'scratch', 'global'][seg] if seg < 3 else 'flat'}_{op_name.split('_', 1)[1] if '_' in op_name else op_name}"
width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'u8':1, 'i8':1, 'u16':1, 'i16':1}.get(op_name.split('_')[-1], 1)
addr_str = _vreg(addr, 2) if saddr == 0x7F else _vreg(addr)
saddr_str = "" if saddr == 0x7F else f", {_sreg(saddr, 2)}" if saddr < 106 else ", off" if saddr == 124 else f", {decode_src(saddr)}"
off_str = f" offset:{offset}" if offset else ""
vdata_str = _vreg(data if 'store' in op_name else vdst, width)
return f"{instr} {addr_str}, {vdata_str}{saddr_str}{off_str}" if 'store' in op_name else f"{instr} {vdata_str}, {addr_str}{saddr_str}{off_str}"
# VOP3: vector ops with modifiers (can be 1, 2, or 3 sources depending on opcode range)
if cls_name == 'VOP3':
# Handle VOP3SD opcodes (same encoding, different field layout)
if is_vop3sd:
vdst = unwrap(inst._values.get('vdst', 0))
# VOP3SD: sdst is at bits [14:8], but VOP3 decodes opsel at [14:11], abs at [10:8], clmp at [15]
# We need to reconstruct sdst from these fields
opsel_raw = unwrap(inst._values.get('opsel', 0))
abs_raw = unwrap(inst._values.get('abs', 0))
clmp_raw = unwrap(inst._values.get('clmp', 0))
sdst = (clmp_raw << 7) | (opsel_raw << 3) | abs_raw
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg = unwrap(inst._values.get('neg', 0))
omod = unwrap(inst._values.get('omod', 0))
omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
is_f64 = 'f64' in op_name
# v_mad_i64_i32/v_mad_u64_u32: 64-bit dst and src2, 32-bit src0/src1
is_mad64 = 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
def fmt_sd_src(v, neg_bit, is_64bit=False):
s = _fmt_src64(v) if (is_64bit or is_f64) else fmt_src(v)
return f"-{s}" if neg_bit else s
src0_str, src1_str = fmt_sd_src(src0, neg & 1), fmt_sd_src(src1, neg & 2)
src2_str = fmt_sd_src(src2, neg & 4, is_mad64)
dst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}"
sdst_str = _fmt_sdst(sdst, 1)
# v_add_co_u32, v_sub_co_u32, v_subrev_co_u32, v_add_co_ci_u32, etc. only use 2 sources
if op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32', 'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'):
return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}"
# v_div_scale uses 3 sources
return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}, {src2_str}" + omod_str
vdst = unwrap(inst._values.get('vdst', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, abs_, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('abs', 0)), unwrap(inst._values.get('clmp', 0))
opsel = unwrap(inst._values.get('opsel', 0))
# Check if 64-bit op (needs register pairs)
is_f64 = 'f64' in op_name or 'i64' in op_name or 'u64' in op_name or 'b64' in op_name
# v_cmp_class_* has 64-bit src0 but 32-bit src1 (class mask)
is_class = 'class' in op_name
# Shift ops: v_*rev_*64 have 32-bit shift amount (src0), 64-bit value (src1)
is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_')
# v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent)
is_ldexp64 = op_name == 'v_ldexp_f64'
# v_trig_preop_f64: 64-bit dst/src0, 32-bit src1 (exponent/scale)
is_trig_preop = op_name == 'v_trig_preop_f64'
# v_readlane_b32: destination is SGPR (despite vdst field)
is_readlane = op_name == 'v_readlane_b32'
# SAD/QSAD/MQSAD instructions have mixed sizes
# v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1
# v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1
is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk'))
is_mqsad_u32 = 'mqsad_u32' in op_name
# Detect 16-bit and 64-bit operand sizes for various instruction patterns
if 'cvt_pk' in op_name:
is_f16_dst, is_f16_src, is_f16_src2 = False, op_name.endswith('16'), False
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', op_name):
dst_type, src_type = m.group(1), m.group(2)
is_f16_dst, is_f16_src, is_f16_src2 = _is_16bit(dst_type), _is_16bit(src_type), _is_16bit(src_type)
is_f64_dst, is_f64_src, is_f64 = '64' in dst_type, '64' in src_type, False
elif re.match(r'v_mad_[iu]32_[iu]16', op_name):
is_f16_dst, is_f16_src, is_f16_src2 = False, True, False # 32-bit dst, 16-bit src0/src1, 32-bit src2
elif 'pack_b32' in op_name:
is_f16_dst, is_f16_src, is_f16_src2 = False, True, True # 32-bit dst, 16-bit sources
else:
is_16bit_op = any(x in op_name for x in _16BIT_TYPES) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad'))
is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op
def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False):
s = _fmt_src_n(v, reg_cnt) if reg_cnt > 1 else f"v{v - 256}.h" if is_16 and v >= 256 and hi_bit else f"v{v - 256}.l" if is_16 and v >= 256 else fmt_src(v)
if abs_bit: s = f"|{s}|"
return f"-{s}" if neg_bit else s
# Determine register count for each source (check for cvt-specific 64-bit flags first)
is_src0_64 = locals().get('is_f64_src', is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32
is_src1_64 = is_f64 and not is_class and not is_ldexp64 and not is_trig_preop
src0_cnt = 2 if is_src0_64 else 1
src1_cnt = 2 if is_src1_64 else 1
src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1
src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src)
src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src)
src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2)
# Format destination - for 16-bit ops, use .h/.l suffix; readlane uses SGPR dest
is_dst_64 = locals().get('is_f64_dst', is_f64) or is_sad64
dst_cnt = 4 if is_mqsad_u32 else 2 if is_dst_64 else 1
if is_readlane:
dst_str = _fmt_sdst(vdst, 1)
elif dst_cnt > 1:
dst_str = _vreg(vdst, dst_cnt)
elif is_f16_dst:
dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l"
else:
dst_str = f"v{vdst}"
clamp_str = " clamp" if clmp else ""
omod = unwrap(inst._values.get('omod', 0))
omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
# op_sel for non-VGPR sources (when opsel bits are set but source is not a VGPR)
# For 16-bit ops with VGPR sources, opsel is encoded in .h/.l suffix
# For non-VGPR sources or non-16-bit ops, we need explicit op_sel
has_nonvgpr_opsel = (src0 < 256 and (opsel & 1)) or (src1 < 256 and (opsel & 2)) or (src2 < 256 and (opsel & 4))
need_opsel = has_nonvgpr_opsel or (opsel and not is_f16_src)
# Helper to format opsel string based on source count
def fmt_opsel(num_src):
if not need_opsel: return ""
# When dst is .h (for 16-bit ops) and non-VGPR sources have opsel, use all 1s
if is_f16_dst and (opsel & 8): # dst is .h
return f" op_sel:[1,1,1{',1' if num_src == 3 else ''}]"
# Otherwise output actual opsel values
if num_src == 3:
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
# Determine number of sources based on opcode range:
# 0-255: VOPC promoted (comparison, 2 src, sdst)
# 256-383: VOP2 promoted (2 src)
# 384-511: VOP1 promoted (1 src)
# 512+: Native VOP3 (2 or 3 src depending on instruction)
if op_val < 256: # VOPC promoted
# VOPCX (v_cmpx_*) writes to exec, no explicit destination
if op_name.startswith('v_cmpx'):
return f"{op_name}_e64 {src0_str}, {src1_str}"
return f"{op_name}_e64 {_fmt_sdst(vdst, 1)}, {src0_str}, {src1_str}"
elif op_val < 384: # VOP2 promoted
# v_cndmask_b32 in VOP3 format has 3 sources (src2 is mask selector)
if 'cndmask' in op_name:
return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
elif op_val < 512: # VOP1 promoted
if op_name in ('v_nop', 'v_pipeflush'): return f"{op_name}_e64"
return f"{op_name}_e64 {dst_str}, {src0_str}" + fmt_opsel(1) + clamp_str + omod_str
else: # Native VOP3 - determine 2 vs 3 sources based on instruction name
# 3-source ops: fma, mad, min3, max3, med3, div_fixup, div_fmas, sad, msad, qsad, mqsad, lerp, alignbit/byte, cubeid/sc/tc/ma, bfe, bfi, perm_b32, permlane, cndmask
# Note: v_writelane_b32 is 2-src (src0, src1 with vdst as 3rd operand - read-modify-write)
is_3src = any(x in op_name for x in ('fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube',
'bfe', 'bfi', 'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add',
'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit'))
if is_3src:
return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
return f"{op_name} {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
# VOP3SD: 3-source with scalar destination (v_div_scale_*, v_add_co_u32, v_mad_*64_*32, etc.)
if cls_name == 'VOP3SD':
vdst, sdst = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('sdst', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, omod, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('omod', 0)), unwrap(inst._values.get('clmp', 0))
is_f64, is_mad64 = 'f64' in op_name, 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
def fmt_neg(v, neg_bit, is_64=False): return f"-{_fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)}" if neg_bit else _fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)
srcs = [fmt_neg(src0, neg & 1), fmt_neg(src1, neg & 2), fmt_neg(src2, neg & 4, is_mad64)]
dst_str, sdst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}", _fmt_sdst(sdst, 1)
clamp_str, omod_str = " clamp" if clmp else "", {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
is_2src = op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32')
suffix = "_e64" if op_name.startswith('v_') and 'co_' in op_name else ""
return f"{op_name}{suffix} {dst_str}, {sdst_str}, {', '.join(srcs[:2] if is_2src else srcs)}" + clamp_str + omod_str
# VOPD: dual-issue instructions
if cls_name == 'VOPD':
from extra.assembly.rdna3 import autogen
opx, opy, vdstx, vdsty_enc = [unwrap(inst._values.get(f, 0)) for f in ('opx', 'opy', 'vdstx', 'vdsty')]
srcx0, vsrcx1, srcy0, vsrcy1 = [unwrap(inst._values.get(f, 0)) for f in ('srcx0', 'vsrcx1', 'srcy0', 'vsrcy1')]
vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1) # Decode vdsty
def fmt_vopd(op, vdst, src0, vsrc1):
try: name = autogen.VOPDOp(op).name.lower()
except (ValueError, KeyError): name = f"op_{op}"
return f"{name} v{vdst}, {fmt_src(src0)}" if 'mov' in name else f"{name} v{vdst}, {fmt_src(src0)}, v{vsrc1}"
return f"{fmt_vopd(opx, vdstx, srcx0, vsrcx1)} :: {fmt_vopd(opy, vdsty, srcy0, vsrcy1)}"
# VOP3P: packed vector ops
if cls_name == 'VOP3P':
vdst, clmp = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('clmp', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, neg_hi = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('neg_hi', 0))
opsel, opsel_hi, opsel_hi2 = unwrap(inst._values.get('opsel', 0)), unwrap(inst._values.get('opsel_hi', 0)), unwrap(inst._values.get('opsel_hi2', 0))
is_wmma, is_3src = 'wmma' in op_name, any(x in op_name for x in ('fma', 'mad', 'dot', 'wmma'))
def fmt_bits(name, val, n): return f"{name}:[{','.join(str((val >> i) & 1) for i in range(n))}]"
# WMMA: f16/bf16 use 8-reg sources, iu8 uses 4-reg, iu4 uses 2-reg; all have 8-reg dst
if is_wmma:
src_cnt = 2 if 'iu4' in op_name else 4 if 'iu8' in op_name else 8
src0_str, src1_str, src2_str = _fmt_src_n(src0, src_cnt), _fmt_src_n(src1, src_cnt), _fmt_src_n(src2, 8)
dst_str = _vreg(vdst, 8)
else:
src0_str, src1_str, src2_str = _fmt_src_n(src0, 1), _fmt_src_n(src1, 1), _fmt_src_n(src2, 1)
dst_str = f"v{vdst}"
n = 3 if is_3src else 2
full_opsel_hi = opsel_hi | (opsel_hi2 << 2)
mods = [fmt_bits("op_sel", opsel, n)] if opsel else []
if full_opsel_hi != (0b111 if is_3src else 0b11): mods.append(fmt_bits("op_sel_hi", full_opsel_hi, n))
if neg: mods.append(fmt_bits("neg_lo", neg, n))
if neg_hi: mods.append(fmt_bits("neg_hi", neg_hi, n))
if clmp: mods.append("clamp")
mod_str = " " + " ".join(mods) if mods else ""
return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}{mod_str}" if is_3src else f"{op_name} {dst_str}, {src0_str}, {src1_str}{mod_str}"
# VINTERP: interpolation instructions
if cls_name == 'VINTERP':
vdst = unwrap(inst._values.get('vdst', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, waitexp, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('waitexp', 0)), unwrap(inst._values.get('clmp', 0))
def fmt_neg_vi(v, neg_bit): return f"-{v}" if neg_bit else v
srcs = [fmt_neg_vi(f"v{s - 256}" if s >= 256 else fmt_src(s), neg & (1 << i)) for i, s in enumerate([src0, src1, src2])]
mods = [m for m in [f"wait_exp:{waitexp}" if waitexp else "", "clamp" if clmp else ""] if m]
return f"{op_name} v{vdst}, {', '.join(srcs)}" + (" " + " ".join(mods) if mods else "")
# MUBUF/MTBUF helpers
def _buf_vaddr(vaddr, offen, idxen): return _vreg(vaddr, 2) if offen and idxen else f"v{vaddr}" if offen or idxen else "off"
def _buf_srsrc(srsrc): srsrc_base = srsrc * 4; return _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
# MUBUF: buffer load/store
if cls_name == 'MUBUF':
vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
offset, offen, idxen = unwrap(inst._values.get('offset', 0)), unwrap(inst._values.get('offen', 0)), unwrap(inst._values.get('idxen', 0))
glc, dlc, slc, tfe = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc', 'tfe')]
if op_name in ('buffer_gl0_inv', 'buffer_gl1_inv'): return op_name
# Determine data width from op name
if 'd16' in op_name: width = 2 if any(x in op_name for x in ('xyz', 'xyzw')) else 1
elif 'atomic' in op_name:
base_width = 2 if any(x in op_name for x in ('b64', 'u64', 'i64')) else 1
width = base_width * 2 if 'cmpswap' in op_name else base_width
else: width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'b16':1, 'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
if tfe: width += 1
mods = [m for m in ["offen" if offen else "", "idxen" if idxen else "", f"offset:{offset}" if offset else "",
"glc" if glc else "", "dlc" if dlc else "", "slc" if slc else "", "tfe" if tfe else ""] if m]
return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)}" + (" " + " ".join(mods) if mods else "")
# MTBUF: typed buffer load/store
if cls_name == 'MTBUF':
vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
offset, tbuf_fmt, offen, idxen = [unwrap(inst._values.get(f, 0)) for f in ('offset', 'format', 'offen', 'idxen')]
glc, dlc, slc = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc')]
mods = [f"format:{tbuf_fmt}"] + [m for m in ["idxen" if idxen else "", "offen" if offen else "", f"offset:{offset}" if offset else "",
"glc" if glc else "", "dlc" if dlc else "", "slc" if slc else ""] if m]
width = 2 if 'd16' in op_name and any(x in op_name for x in ('xyz', 'xyzw')) else 1 if 'd16' in op_name else {'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)} {' '.join(mods)}"
# SOP1/SOP2/SOPC/SOPK
if cls_name in ('SOP1', 'SOP2', 'SOPC', 'SOPK'):
sizes = _parse_sop_sizes(op_name)
dst_cnt, src0_cnt = sizes[0], sizes[1]
src1_cnt = sizes[2] if len(sizes) > 2 else src0_cnt
if cls_name == 'SOP1':
sdst, ssrc0 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('ssrc0', 0))
if op_name == 's_getpc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}"
if op_name in ('s_setpc_b64', 's_rfe_b64'): return f"{op_name} {_fmt_ssrc(ssrc0, 2)}"
if op_name == 's_swappc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}, {_fmt_ssrc(ssrc0, 2)}"
if op_name in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'):
return f"{op_name} {_fmt_sdst(sdst, 2 if 'b64' in op_name else 1)}, sendmsg({MSG_NAMES.get(ssrc0, str(ssrc0))})"
ssrc0_str = fmt_src(ssrc0) if src0_cnt == 1 else _fmt_ssrc(ssrc0, src0_cnt)
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {ssrc0_str}"
if cls_name == 'SOP2':
sdst, ssrc0, ssrc1 = [unwrap(inst._values.get(f, 0)) for f in ('sdst', 'ssrc0', 'ssrc1')]
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {_fmt_ssrc(ssrc0, src0_cnt)}, {_fmt_ssrc(ssrc1, src1_cnt)}"
if cls_name == 'SOPC':
return f"{op_name} {_fmt_ssrc(unwrap(inst._values.get('ssrc0', 0)), src0_cnt)}, {_fmt_ssrc(unwrap(inst._values.get('ssrc1', 0)), src1_cnt)}"
if cls_name == 'SOPK':
sdst, simm16 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('simm16', 0))
if op_name == 's_version': return f"{op_name} 0x{simm16:x}"
if op_name in ('s_setreg_b32', 's_getreg_b32'):
hwreg_id, hwreg_offset, hwreg_size = simm16 & 0x3f, (simm16 >> 6) & 0x1f, ((simm16 >> 11) & 0x1f) + 1
hwreg_str = f"0x{simm16:x}" if hwreg_id in (16, 17) else f"hwreg({HWREG_NAMES.get(hwreg_id, str(hwreg_id))}, {hwreg_offset}, {hwreg_size})"
return f"{op_name} {hwreg_str}, {_fmt_sdst(sdst, 1)}" if op_name == 's_setreg_b32' else f"{op_name} {_fmt_sdst(sdst, 1)}, {hwreg_str}"
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, 0x{simm16:x}"
# Generic fallback
def fmt_field(n, v):
v = unwrap(v)
if n in SRC_FIELDS: return fmt_src(v) if v != 255 else "0xff"
if n in ('sdst', 'vdst'): return f"{'s' if n == 'sdst' else 'v'}{v}"
return f"v{v}" if n == 'vsrc1' else f"0x{v:x}" if n == 'simm16' else str(v)
ops = [fmt_field(n, inst._values.get(n, 0)) for n in inst._fields if n not in ('encoding', 'op')]
return f"{op_name} {', '.join(ops)}" if ops else op_name
# Assembler
SPECIAL_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125), 'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'scc': RawImm(253)}
FLOAT_CONSTS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
def parse_operand(op: str) -> tuple:
op = op.strip().lower()
neg = op.startswith('-') and not op[1:2].isdigit(); op = op[1:] if neg else op
abs_ = op.startswith('|') and op.endswith('|') or op.startswith('abs(') and op.endswith(')')
op = op[1:-1] if op.startswith('|') else op[4:-1] if op.startswith('abs(') else op
hi_half = op.endswith('.h')
op = re.sub(r'\.[lh]$', '', op)
if op in FLOAT_CONSTS: return (FLOAT_CONSTS[op], neg, abs_, hi_half)
if re.match(r'^-?\d+$', op): return (int(op), neg, abs_, hi_half)
if m := re.match(r'^-?0x([0-9a-f]+)$', op):
v = -int(m.group(1), 16) if op.startswith('-') else int(m.group(1), 16)
return (v, neg, abs_, hi_half)
if op in SPECIAL_REGS: return (SPECIAL_REGS[op], neg, abs_, hi_half)
if op == 'lit': return (RawImm(255), neg, abs_, hi_half) # literal marker (actual value comes from literal word)
if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', op): return (REG_MAP[m.group(1)][int(m.group(2)):int(m.group(3))], neg, abs_, hi_half)
if m := re.match(r'^([svt](?:tmp)?)(\d+)$', op):
reg = REG_MAP[m.group(1)][int(m.group(2))]
reg.hi = hi_half
return (reg, neg, abs_, hi_half)
# hwreg(name, offset, size) or hwreg(name) -> simm16 encoding
if m := re.match(r'^hwreg\((\w+)(?:,\s*(\d+),\s*(\d+))?\)$', op):
name_str = m.group(1).lower()
hwreg_id = HWREG_IDS.get(name_str, int(name_str) if name_str.isdigit() else None)
if hwreg_id is None: raise ValueError(f"unknown hwreg name: {name_str}")
offset, size = int(m.group(2)) if m.group(2) else 0, int(m.group(3)) if m.group(3) else 32
return (((size - 1) << 11) | (offset << 6) | hwreg_id, neg, abs_, hi_half)
raise ValueError(f"cannot parse operand: {op}")
SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
SOP1_SRC_ONLY = {'s_setpc_b64', 's_rfe_b64'}
SOP1_MSG_IMM = {'s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'}
SOPK_IMM_ONLY = {'s_version'}
SOPK_IMM_FIRST = {'s_setreg_b32'}
SOPK_UNSUPPORTED = {'s_setreg_imm32_b32'}
def asm(text: str) -> Inst:
from extra.assembly.rdna3 import autogen
text = text.strip()
clamp = 'clamp' in text.lower()
if clamp: text = re.sub(r'\s+clamp\s*$', '', text, flags=re.I)
modifiers = {}
if m := re.search(r'\s+wait_exp:(\d+)', text, re.I): modifiers['waitexp'] = int(m.group(1)); text = text[:m.start()] + text[m.end():]
parts = text.replace(',', ' ').split()
if not parts: raise ValueError("empty instruction")
mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip()
# Handle s_waitcnt specially before operand parsing
if mnemonic == 's_waitcnt':
vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f
for part in op_str.replace(',', ' ').split():
if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1))
elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1))
elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1))
elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0))
return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt))
# Handle VOPD dual-issue instructions: opx dst, src :: opy dst, src
if '::' in text:
x_part, y_part = text.split('::')
x_parts, y_parts = x_part.strip().replace(',', ' ').split(), y_part.strip().replace(',', ' ').split()
opx_name, opy_name = x_parts[0].upper(), y_parts[0].upper()
opx, opy = autogen.VOPDOp[opx_name], autogen.VOPDOp[opy_name]
x_ops, y_ops = [parse_operand(p)[0] for p in x_parts[1:]], [parse_operand(p)[0] for p in y_parts[1:]]
vdstx, srcx0 = x_ops[0], x_ops[1] if len(x_ops) > 1 else 0
vsrcx1 = x_ops[2] if len(x_ops) > 2 else VGPR(0)
vdsty, srcy0 = y_ops[0], y_ops[1] if len(y_ops) > 1 else 0
vsrcy1 = y_ops[2] if len(y_ops) > 2 else VGPR(0)
# Handle fmaak/fmamk literals (4th operand on x or y side)
lit = None
if 'fmaak' in opx_name.lower() and len(x_ops) > 3: lit = unwrap(x_ops[3])
elif 'fmamk' in opx_name.lower() and len(x_ops) > 3: lit, vsrcx1 = unwrap(x_ops[2]), x_ops[3]
elif 'fmaak' in opy_name.lower() and len(y_ops) > 3: lit = unwrap(y_ops[3])
elif 'fmamk' in opy_name.lower() and len(y_ops) > 3: lit, vsrcy1 = unwrap(y_ops[2]), y_ops[3]
return autogen.VOPD(opx, opy, vdstx=vdstx, vdsty=vdsty, srcx0=srcx0, vsrcx1=vsrcx1, srcy0=srcy0, vsrcy1=vsrcy1, literal=lit)
operands, current, depth, in_pipe = [], "", 0, False
for ch in op_str:
if ch in '[(': depth += 1
elif ch in '])': depth -= 1
elif ch == '|': in_pipe = not in_pipe
if ch == ',' and depth == 0 and not in_pipe: operands.append(current.strip()); current = ""
else: current += ch
if current.strip(): operands.append(current.strip())
parsed = [parse_operand(op) for op in operands]
values = [p[0] for p in parsed]
neg_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[1])
abs_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[2])
opsel_bits = (8 if len(parsed) > 0 and parsed[0][3] else 0) | sum((1 << i) for i, p in enumerate(parsed[1:4]) if p[3])
lit = None
if mnemonic in ('v_fmaak_f32', 'v_fmaak_f16') and len(values) == 4: lit, values = unwrap(values[3]), values[:3]
elif mnemonic in ('v_fmamk_f32', 'v_fmamk_f16') and len(values) == 4: lit, values = unwrap(values[2]), [values[0], values[1], values[3]]
vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32', 'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}
if mnemonic.replace('_e32', '') in vcc_ops and len(values) >= 5: values = [values[0], values[2], values[3]]
if mnemonic.startswith('v_cmp') and len(values) >= 3 and operands[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'):
values = values[1:]
# CMPX instructions with _e64 suffix: prepend implicit EXEC_LO destination (vdst=126)
if 'cmpx' in mnemonic and mnemonic.endswith('_e64') and len(values) == 2:
values = [VGPR(126, 1)] + values
# Recalculate modifiers: parsed[0]=src0, parsed[1]=src1 (no vdst in user input)
neg_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[1])
abs_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[2])
opsel_bits = sum((1 << i) for i, p in enumerate(parsed[:2]) if p[3])
vop3sd_ops = {'v_div_scale_f32', 'v_div_scale_f64'}
if mnemonic in vop3sd_ops and len(parsed) >= 5:
neg_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[1])
abs_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[2])
if mnemonic in SOPK_UNSUPPORTED: raise ValueError(f"unsupported instruction: {mnemonic}")
elif mnemonic in SOP1_SRC_ONLY:
return getattr(autogen, mnemonic)(ssrc0=values[0])
elif mnemonic in SOP1_MSG_IMM:
return getattr(autogen, mnemonic)(sdst=values[0], ssrc0=RawImm(unwrap(values[1])))
elif mnemonic in SOPK_IMM_ONLY:
return getattr(autogen, mnemonic)(simm16=values[0])
elif mnemonic in SOPK_IMM_FIRST:
return getattr(autogen, mnemonic)(simm16=values[0], sdst=values[1])
elif mnemonic in SMEM_OPS and len(operands) >= 3 and re.match(r'^-?[0-9]|^-?0x', operands[2].strip().lower()):
return getattr(autogen, mnemonic)(sdata=values[0], sbase=values[1], offset=values[2], soffset=RawImm(124))
elif mnemonic.startswith('buffer_') and len(operands) >= 2 and operands[1].strip().lower() == 'off':
return getattr(autogen, mnemonic)(vdata=values[0], vaddr=0, srsrc=values[2], soffset=RawImm(unwrap(values[3])) if len(values) > 3 else RawImm(0))
elif (mnemonic.startswith('flat_load') or mnemonic.startswith('global_load') or mnemonic.startswith('scratch_load')) and len(values) >= 3:
offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
return getattr(autogen, mnemonic)(vdst=values[0], addr=values[1], saddr=values[2], offset=offset)
elif (mnemonic.startswith('flat_store') or mnemonic.startswith('global_store') or mnemonic.startswith('scratch_store')) and len(values) >= 3:
offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
return getattr(autogen, mnemonic)(addr=values[0], data=values[1], saddr=values[2], offset=offset)
for suffix in (['_e32', ''] if not (neg_bits or abs_bits or clamp) else ['', '_e32']):
if hasattr(autogen, name := mnemonic.replace('.', '_') + suffix):
use_opsel = 'opsel' in getattr(autogen, name).func._fields
vals = [type(v)(v.idx, v.count, False) if isinstance(v, Reg) and v.hi and use_opsel else v for v in values]
inst = getattr(autogen, name)(*vals, literal=lit, **modifiers)
if neg_bits and 'neg' in inst._fields: inst._values['neg'] = neg_bits
if opsel_bits and use_opsel: inst._values['opsel'] = opsel_bits
if abs_bits and 'abs' in inst._fields: inst._values['abs'] = abs_bits
if clamp and 'clmp' in inst._fields: inst._values['clmp'] = 1
return inst
raise ValueError(f"unknown instruction: {mnemonic}")
-505
View File
@@ -1,505 +0,0 @@
# RDNA3 emulator - pure Python implementation for testing
from __future__ import annotations
import ctypes, struct, math
from typing import Callable
from extra.assembly.rdna3.lib import Inst, Inst32, Inst64, RawImm
Program = dict[int, Inst] # pc (word offset) -> instruction
from extra.assembly.rdna3.autogen import (
SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
)
from extra.assembly.rdna3.alu import (
f32, i32, f16, i16, sext, vopc, SALU, VALU,
SOP1_BASE, SOP2_BASE, SOPC_BASE, SOPK_BASE, VOP1_BASE, VOP2_BASE
)
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
# Pre-computed inline constant table for src operands 128-254 (index = src - 128)
_INLINE_CONSTS = [0] * 127
for _i in range(65): _INLINE_CONSTS[_i] = _i # 128-192 -> 0-64
for _i in range(1, 17): _INLINE_CONSTS[64 + _i] = ((-_i) & 0xffffffff) # 193-208 -> -1 to -16
for _k, _v in {SrcEnum.POS_HALF: 0x3f000000, SrcEnum.NEG_HALF: 0xbf000000, SrcEnum.POS_ONE: 0x3f800000, SrcEnum.NEG_ONE: 0xbf800000,
SrcEnum.POS_TWO: 0x40000000, SrcEnum.NEG_TWO: 0xc0000000, SrcEnum.POS_FOUR: 0x40800000, SrcEnum.NEG_FOUR: 0xc0800000,
SrcEnum.INV_2PI: 0x3e22f983}.items(): _INLINE_CONSTS[_k - 128] = _v
_valid_mem_ranges: list[tuple[int, int]] = []
def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: global _valid_mem_ranges; _valid_mem_ranges = list(ranges)
def _mem_valid(addr: int, size: int) -> bool:
for s, z in _valid_mem_ranges:
if s <= addr and addr + size <= s + z: return True
return not _valid_mem_ranges
def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
def mem_write(addr: int, size: int, val: int) -> None:
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
# Memory op tables - (cnt, sz, sign) for loads, (cnt, sz) for stores
def _mem_ops(ops, suffix_map):
return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
FLAT_LOAD = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP)
FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
DS_LOAD: dict[int, tuple[int,int,int]] = {DSOp.DS_LOAD_B32: (1,4,0), DSOp.DS_LOAD_B64: (2,4,0), DSOp.DS_LOAD_B128: (4,4,0), DSOp.DS_LOAD_U8: (1,1,0), DSOp.DS_LOAD_I8: (1,1,1), DSOp.DS_LOAD_U16: (1,2,0), DSOp.DS_LOAD_I16: (1,2,1)}
DS_STORE: dict[int, tuple[int,int]] = {DSOp.DS_STORE_B32: (1,4), DSOp.DS_STORE_B64: (2,4), DSOp.DS_STORE_B128: (4,4), DSOp.DS_STORE_B8: (1,1), DSOp.DS_STORE_B16: (1,2)}
FLAT_D16_LO = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_U8', (1, 0)), ('LOAD_D16_I8', (1, 1)), ('LOAD_D16_B16', (2, 0))]}
FLAT_D16_HI = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_HI_U8', (1, 0)), ('LOAD_D16_HI_I8', (1, 1)), ('LOAD_D16_HI_B16', (2, 0))]}
FLAT_D16_STORE = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('STORE_D16_HI_B8', 1), ('STORE_D16_HI_B16', 2)]}
SMEM_LOAD: dict[int, int] = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
SOPK_WAIT = {SOPKOp.S_WAITCNT_VSCNT, SOPKOp.S_WAITCNT_VMCNT, SOPKOp.S_WAITCNT_EXPCNT, SOPKOp.S_WAITCNT_LGKMCNT}
class WaveState:
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
def __init__(self):
self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
self.sgpr[EXEC_LO] = 0xffffffff # wave32: all lanes active
self.scc = self.pc = self.literal = 0
self._pend_sgpr = {}
@property
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
@vcc.setter
def vcc(self, v: int) -> None: self.sgpr[VCC_LO] = v & 0xffffffff; self.sgpr[VCC_HI] = (v >> 32) & 0xffffffff
@property
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
@exec_mask.setter
def exec_mask(self, v: int) -> None: self.sgpr[EXEC_LO] = v & 0xffffffff; self.sgpr[EXEC_HI] = (v >> 32) & 0xffffffff
def rsgpr(self, i: int) -> int:
if i == NULL: return 0
if i == SCC: return self.scc
return self.sgpr[i] if i < SGPR_COUNT else 0
def wsgpr(self, i: int, v: int) -> None:
if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & 0xffffffff
def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
def wsgpr64(self, i: int, v: int) -> None: self.wsgpr(i, v & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
def rsrc(self, v: int, lane: int) -> int:
if v < SGPR_COUNT: return self.sgpr[v]
if v == SCC: return self.scc
if v < 255: return _INLINE_CONSTS[v - 128]
if v == 255: return self.literal
return self.vgpr[lane][v - 256] if v <= 511 else 0
def rsrc64(self, v: int, lane: int) -> int:
return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
def pend_sgpr_lane(self, reg: int, lane: int, val: int) -> None:
if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
if val: self._pend_sgpr[reg] |= (1 << lane)
def commit_pends(self) -> None:
for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
self._pend_sgpr.clear()
def decode_format(word: int) -> tuple[type[Inst] | None, bool]:
hi2 = (word >> 30) & 0x3
if hi2 == 0b11:
enc = (word >> 26) & 0xf
if enc == 0b1101: return SMEM, True
if enc == 0b0101:
op = (word >> 16) & 0x3ff
return (VOP3SD, True) if op in (288, 289, 290, 764, 765, 766, 767, 768, 769, 770) else (VOP3, True)
return {0b0011: (VOP3P, True), 0b0110: (DS, True), 0b0111: (FLAT, True), 0b0010: (VOPD, True)}.get(enc, (None, True))
if hi2 == 0b10:
enc = (word >> 23) & 0x7f
return {0b1111101: (SOP1, False), 0b1111110: (SOPC, False), 0b1111111: (SOPP, False)}.get(enc, (SOPK, False) if ((word >> 28) & 0xf) == 0b1011 else (SOP2, False))
enc = (word >> 25) & 0x7f
return (VOPC, False) if enc == 0b0111110 else (VOP1, False) if enc == 0b0111111 else (VOP2, False)
def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
def decode_program(data: bytes) -> Program:
result: Program = {}
i = 0
while i < len(data):
word = int.from_bytes(data[i:i+4], 'little')
inst_class, is_64 = decode_format(word)
if inst_class is None: i += 4; continue
base_size = 8 if is_64 else 4
inst = inst_class.from_bytes(data[i:i+base_size])
for name, val in inst._values.items(): setattr(inst, name, _unwrap(val))
has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
if has_literal: inst._literal = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
inst._words = inst.size() // 4 # cache size for step_wave
result[i // 4] = inst
i += inst._words * 4
return result
# ═══════════════════════════════════════════════════════════════════════════════
# SCALAR EXECUTION
# ═══════════════════════════════════════════════════════════════════════════════
def exec_sop1(st: WaveState, inst: SOP1) -> int:
s0, op = st.rsrc(inst.ssrc0, 0), inst.op
# 64-bit and special ops handled inline
if op == SOP1Op.S_MOV_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0)); return 0
if op == SOP1Op.S_NOT_B64: r = (~st.rsrc64(inst.ssrc0, 0)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP1Op.S_BITSET0_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) & ~(1 << (s0 & 0x1f))); return 0
if op == SOP1Op.S_BITSET1_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) | (1 << (s0 & 0x1f))); return 0
if op == SOP1Op.S_AND_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
if op == SOP1Op.S_OR_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 | old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
if op == SOP1Op.S_AND_NOT1_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & (~old & 0xffffffff); st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
if op == SOP1Op.S_GETPC_B64: return -3
if op == SOP1Op.S_SETPC_B64: return -4
if op == SOP1Op.S_SWAPPC_B64: return -5
if (fn := SALU.get(SOP1_BASE + op)) is None: raise NotImplementedError(f"SOP1 op {op}")
r, scc = fn(s0, 0, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
_SOP2_64: dict[int, Callable[[int, int], int]] = {SOP2Op.S_AND_B64: lambda a, b: a & b, SOP2Op.S_OR_B64: lambda a, b: a | b, SOP2Op.S_XOR_B64: lambda a, b: a ^ b}
def exec_sop2(st: WaveState, inst: SOP2) -> int:
s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
# 64-bit ops handled inline
if op == SOP2Op.S_LSHL_B64: r = (st.rsrc64(inst.ssrc0, 0) << (s1 & 0x3f)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP2Op.S_LSHR_B64: r = st.rsrc64(inst.ssrc0, 0) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP2Op.S_ASHR_I64: r = sext(st.rsrc64(inst.ssrc0, 0), 64) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r & 0xffffffffffffffff); st.scc = int(r != 0); return 0
if (fn := _SOP2_64.get(op)): r = fn(st.rsrc64(inst.ssrc0, 0), st.rsrc64(inst.ssrc1, 0)); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP2Op.S_CSELECT_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0) if st.scc else st.rsrc64(inst.ssrc1, 0)); return 0
if op == SOP2Op.S_FMAC_F32: st.wsgpr(inst.sdst, i32(f32(st.rsgpr(inst.sdst)) + f32(s0) * f32(s1))); return 0
if op == SOP2Op.S_FMAAK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(s1) + f32(inst._literal or 0))); return 0
if op == SOP2Op.S_FMAMK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(inst._literal or 0) + f32(s1))); return 0
if (fn := SALU.get(SOP2_BASE + op)) is None: raise NotImplementedError(f"SOP2 op {op}")
r, scc = fn(s0, s1, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
def exec_sopc(st: WaveState, inst: SOPC) -> int:
s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
if op == SOPCOp.S_CMP_EQ_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) == st.rsrc64(inst.ssrc1, 0)); return 0
if op == SOPCOp.S_CMP_LG_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) != st.rsrc64(inst.ssrc1, 0)); return 0
if (fn := SALU.get(SOPC_BASE + op)) is None: raise NotImplementedError(f"SOPC op {op}")
st.scc = fn(s0, s1, st.scc)[1]; return 0
_SOPK_CMP = frozenset((SOPKOp.S_CMPK_EQ_I32, SOPKOp.S_CMPK_LG_I32, SOPKOp.S_CMPK_GT_I32, SOPKOp.S_CMPK_GE_I32,
SOPKOp.S_CMPK_LT_I32, SOPKOp.S_CMPK_LE_I32, SOPKOp.S_CMPK_EQ_U32, SOPKOp.S_CMPK_LG_U32,
SOPKOp.S_CMPK_GT_U32, SOPKOp.S_CMPK_GE_U32, SOPKOp.S_CMPK_LT_U32, SOPKOp.S_CMPK_LE_U32))
def exec_sopk(st: WaveState, inst: SOPK) -> int:
simm, s0, op = inst.simm16, st.rsgpr(inst.sdst), inst.op
if op in SOPK_WAIT: return 0
if (fn := SALU.get(SOPK_BASE + op)) is None: raise NotImplementedError(f"SOPK op {op}")
r, scc = fn(s0, simm, st.scc)
if op not in _SOPK_CMP: st.wsgpr(inst.sdst, r)
st.scc = scc; return 0
def exec_sopp(st: WaveState, inst: SOPP) -> int:
if inst.op == SOPPOp.S_ENDPGM: return -1
if inst.op == SOPPOp.S_BARRIER: return -2
if inst.op == SOPPOp.S_BRANCH: return sext(inst.simm16, 16)
if inst.op == SOPPOp.S_CBRANCH_SCC0: return sext(inst.simm16, 16) if st.scc == 0 else 0
if inst.op == SOPPOp.S_CBRANCH_SCC1: return sext(inst.simm16, 16) if st.scc == 1 else 0
# In wave32 mode, only VCC_LO is used for lane masks; VCC_HI is a free SGPR
if inst.op == SOPPOp.S_CBRANCH_VCCZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) == 0 else 0
if inst.op == SOPPOp.S_CBRANCH_VCCNZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) != 0 else 0
if inst.op == SOPPOp.S_CBRANCH_EXECZ: return sext(inst.simm16, 16) if st.exec_mask == 0 else 0
if inst.op == SOPPOp.S_CBRANCH_EXECNZ: return sext(inst.simm16, 16) if st.exec_mask != 0 else 0
# Scheduling hints and wait instructions are no-ops in emulation
if inst.op <= 31: return 0 # S_NOP, S_CLAUSE, S_DELAY_ALU, S_WAITCNT, etc.
# S_WAKEUP(52), S_SETPRIO(53), S_SENDMSG(54), S_SENDMSGHALT(55), perf counters, S_ICACHE_INV(60) are no-ops
if inst.op in (52, 53, 54, 55, 56, 57, 60): return 0
raise NotImplementedError(f"SOPP op {inst.op}")
def exec_smem(st: WaveState, inst: SMEM) -> int:
addr = st.rsgpr64(inst.sbase * 2) + sext(inst.offset, 21)
if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & 0xffffffffffffffff, 4))
return 0
# ═══════════════════════════════════════════════════════════════════════════════
# VECTOR EXECUTION
# ═══════════════════════════════════════════════════════════════════════════════
def f64(hi: int, lo: int) -> float: return struct.unpack('<d', struct.pack('<Q', (hi << 32) | lo))[0]
def i64_parts(f: float) -> tuple[int, int]:
if math.isnan(f): val = 0x7ff8000000000000
elif math.isinf(f): val = 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
else: val = struct.unpack('<Q', struct.pack('<d', f))[0]
return val & 0xffffffff, (val >> 32) & 0xffffffff
def exec_vop1(st: WaveState, inst: VOP1, lane: int) -> None:
if inst.op == VOP1Op.V_NOP: return
V, s0 = st.vgpr[lane], st.rsrc(inst.src0, lane)
if inst.op == VOP1Op.V_READFIRSTLANE_B32:
first = (st.exec_mask & -st.exec_mask).bit_length() - 1 if st.exec_mask else 0
st.wsgpr(inst.vdst, st.rsrc(inst.src0, first) if inst.src0 >= 256 else s0); return
# F64 ops handled inline
if inst.op == VOP1Op.V_CVT_F64_F32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(f32(s0))); return
if inst.op == VOP1Op.V_CVT_F64_I32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(sext(s0, 32))); return
if inst.op == VOP1Op.V_CVT_F64_U32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(s0)); return
if inst.op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64):
src = inst.src0 - 256 if inst.src0 >= 256 else inst.src0
lo, hi = (V[src], V[src+1]) if inst.src0 >= 256 else (st.sgpr[src], st.sgpr[src+1])
v = f64(hi, lo)
if inst.op == VOP1Op.V_CVT_F32_F64: V[inst.vdst] = i32(v)
elif inst.op == VOP1Op.V_CVT_I32_F64: V[inst.vdst] = (max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff) if math.isfinite(v) else 0
else: V[inst.vdst] = max(0, min(0xffffffff, int(v))) if math.isfinite(v) and v == v else 0
return
if (fn := VALU.get(VOP1_BASE + inst.op)): V[inst.vdst] = fn(s0, 0, 0); return
raise NotImplementedError(f"VOP1 op {inst.op}")
def exec_vop2(st: WaveState, inst: VOP2, lane: int) -> None:
V, s0, s1, op = st.vgpr[lane], st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1], inst.op
if op == VOP2Op.V_CNDMASK_B32: V[inst.vdst] = s1 if (st.vcc >> lane) & 1 else s0; return
if op == VOP2Op.V_FMAC_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(V[inst.vdst])); return
if op == VOP2Op.V_FMAMK_F32: V[inst.vdst] = i32(f32(s0)*f32(st.literal)+f32(s1)); return
if op == VOP2Op.V_FMAAK_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(st.literal)); return
if op == VOP2Op.V_FMAC_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(V[inst.vdst])); return
if op == VOP2Op.V_FMAMK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(st.literal)+f16(s1)); return
if op == VOP2Op.V_FMAAK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(st.literal)); return
if op == VOP2Op.V_PK_FMAC_F16:
lo = i16(f16(s0 & 0xffff) * f16(s1 & 0xffff) + f16(V[inst.vdst] & 0xffff))
hi = i16(f16((s0 >> 16) & 0xffff) * f16((s1 >> 16) & 0xffff) + f16((V[inst.vdst] >> 16) & 0xffff))
V[inst.vdst] = lo | (hi << 16); return
if op == VOP2Op.V_ADD_CO_CI_U32: r = s0+s1+((st.vcc>>lane)&1); st.pend_sgpr_lane(VCC_LO, lane, r >= 0x100000000); V[inst.vdst] = r & 0xffffffff; return
if op == VOP2Op.V_SUB_CO_CI_U32: b = (st.vcc>>lane)&1; st.pend_sgpr_lane(VCC_LO, lane, s1+b > s0); V[inst.vdst] = (s0-s1-b) & 0xffffffff; return
if (fn := VALU.get(VOP2_BASE + op)): V[inst.vdst] = fn(s0, s1, 0); return
raise NotImplementedError(f"VOP2 op {op}")
def vop3_mod(val: int, neg: int, abs_: int, idx: int) -> int:
if (abs_ >> idx) & 1: val = i32(abs(f32(val)))
if (neg >> idx) & 1: val = i32(-f32(val))
return val
def exec_vop3(st: WaveState, inst: VOP3, lane: int) -> None:
op, src0, src1, src2, vdst, neg, abs_ = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.neg, getattr(inst, 'abs', 0)
V = st.vgpr[lane]
# VOPC encoded in VOP3 (0-255)
if 0 <= op <= 255:
base = op & 0x7f
# For 64-bit comparisons (I64: 80-87, U64: 88-95), read raw 64-bit values (no float modifiers)
if 80 <= base <= 95:
s0_64, s1_64 = st.rsrc64(src0, lane), st.rsrc64(src1, lane)
result = vopc(op, s0_64 & 0xffffffff, s1_64 & 0xffffffff, (s0_64 >> 32) & 0xffffffff, (s1_64 >> 32) & 0xffffffff)
else:
s0, s1 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1)
result = vopc(op, s0, s1)
is_cmpx = op >= 128
st.pend_sgpr_lane(vdst, lane, result)
if is_cmpx: st.pend_sgpr_lane(EXEC_LO, lane, result)
return
s0, s1, s2 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1), vop3_mod(st.rsrc(src2, lane), neg, abs_, 2)
# Special ops
if op == VOP3Op.V_FMAC_F32: V[vdst] = i32(f32(s0)*f32(s1)+f32(V[vdst])); return
if op == VOP3Op.V_READLANE_B32: st.wsgpr(vdst, st.vgpr[s1 & 0x1f][src0 - 256] if src0 >= 256 else s0); return
if op == VOP3Op.V_WRITELANE_B32: st.vgpr[s1 & 0x1f][vdst] = s0; return
if op == VOP3Op.V_CNDMASK_B32:
mask = st.rsgpr(src2) if src2 < 256 else st.vcc
V[vdst] = s1 if (mask >> lane) & 1 else s0; return
if op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64):
v64 = st.rsrc64(src1, lane)
r = ((v64 << (s0 & 0x3f)) & 0xffffffffffffffff if op == VOP3Op.V_LSHLREV_B64 else
v64 >> (s0 & 0x3f) if op == VOP3Op.V_LSHRREV_B64 else sext(v64, 64) >> (s0 & 0x3f))
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff; return
if op in (VOP3Op.V_ADD_F64, VOP3Op.V_MUL_F64, VOP3Op.V_FMA_F64, VOP3Op.V_MAX_F64, VOP3Op.V_MIN_F64):
a, b = f64(st.rsrc(src0+1, lane), s0), f64(st.rsrc(src1+1, lane), s1)
c = f64(st.rsrc(src2+1, lane), s2) if op == VOP3Op.V_FMA_F64 else 0.0
rf = a + b if op == VOP3Op.V_ADD_F64 else a * b if op == VOP3Op.V_MUL_F64 else a * b + c if op == VOP3Op.V_FMA_F64 else max(a, b) if op == VOP3Op.V_MAX_F64 else min(a, b)
V[vdst], V[vdst+1] = i64_parts(rf); return
if (fn := VALU.get(op)): V[vdst] = fn(s0, s1, s2); return
raise NotImplementedError(f"VOP3 op {op}")
def exec_vopc(st: WaveState, inst: VOPC, lane: int) -> None:
result, is_cmpx = vopc(inst.op, st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1]), inst.op >= 128
st.pend_sgpr_lane(EXEC_LO if is_cmpx else VCC_LO, lane, result)
def exec_vop3sd(st: WaveState, inst: VOP3SD, lane: int) -> None:
op, src0, src1, src2, vdst, sdst, neg = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.sdst, inst.neg
s0, s1, s2 = st.rsrc(src0, lane), st.rsrc(src1, lane), st.rsrc(src2, lane)
if (neg >> 0) & 1: s0 = i32(-f32(s0))
if (neg >> 1) & 1: s1 = i32(-f32(s1))
if (neg >> 2) & 1: s2 = i32(-f32(s2))
V = st.vgpr[lane]
if op == VOP3SDOp.V_ADD_CO_U32: r = s0 + s1; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
elif op == VOP3SDOp.V_SUB_CO_U32: V[vdst] = (s0 - s1) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 > s0)
elif op == VOP3SDOp.V_SUBREV_CO_U32: V[vdst] = (s1 - s0) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s0 > s1)
elif op == VOP3SDOp.V_ADD_CO_CI_U32:
cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
r = s0 + s1 + cin; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
elif op == VOP3SDOp.V_SUB_CO_CI_U32:
cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
V[vdst] = (s0 - s1 - cin) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 + cin > s0)
elif op == VOP3SDOp.V_MAD_U64_U32:
s2_64 = s2 | (st.rsrc(src2+1, lane) << 32); r = s0 * s1 + s2_64
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
elif op == VOP3SDOp.V_MAD_I64_I32:
s2_64 = sext(s2 | (st.rsrc(src2+1, lane) << 32), 64)
r = (sext(s0, 32) * sext(s1, 32) + s2_64) & 0xffffffffffffffff
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
elif op == VOP3SDOp.V_DIV_SCALE_F32: V[vdst] = 0; st.pend_sgpr_lane(sdst, lane, False)
elif op == VOP3SDOp.V_DIV_SCALE_F64: V[vdst], V[vdst+1] = s0, st.rsrc(src0+1, lane); st.pend_sgpr_lane(VCC_LO, lane, s0 == s2)
else: raise NotImplementedError(f"VOP3SD op {op}")
def exec_flat(st: WaveState, inst: FLAT, lane: int) -> None:
op, addr_reg, data_reg, vdst, offset, saddr, V = inst.op, inst.addr, inst.data, inst.vdst, sext(inst.offset, 13), inst.saddr, st.vgpr[lane]
addr = V[addr_reg] | (V[addr_reg+1] << 32)
addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
if op in FLAT_LOAD:
cnt, sz, sign = FLAT_LOAD[op]
for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
elif op in FLAT_STORE:
cnt, sz = FLAT_STORE[op]
for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
elif op in FLAT_D16_LO: sz, sign = FLAT_D16_LO[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0xffff0000) | ((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff))
elif op in FLAT_D16_HI: sz, sign = FLAT_D16_HI[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0x0000ffff) | (((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff)) << 16)
elif op in FLAT_D16_STORE: mem_write(addr, FLAT_D16_STORE[op], (V[data_reg] >> 16) & ((1 << (FLAT_D16_STORE[op] * 8)) - 1))
else: raise NotImplementedError(f"FLAT op {op}")
def exec_ds(st: WaveState, inst: DS, lane: int, lds: bytearray) -> None:
op, addr, vdst, V = inst.op, (st.vgpr[lane][inst.addr] + inst.offset0) & 0xffff, inst.vdst, st.vgpr[lane]
if op in DS_LOAD:
cnt, sz, sign = DS_LOAD[op]
for i in range(cnt): val = int.from_bytes(lds[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
elif op in DS_STORE:
cnt, sz = DS_STORE[op]
for i in range(cnt): lds[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
else: raise NotImplementedError(f"DS op {op}")
VOPD_OPS: dict[int, Callable[[int, int, int, int, int], int]] = {
VOPDOp.V_DUAL_MUL_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)), VOPDOp.V_DUAL_ADD_F32: lambda a, b, d, l, lit: i32(f32(a)+f32(b)),
VOPDOp.V_DUAL_SUB_F32: lambda a, b, d, l, lit: i32(f32(a)-f32(b)), VOPDOp.V_DUAL_SUBREV_F32: lambda a, b, d, l, lit: i32(f32(b)-f32(a)),
VOPDOp.V_DUAL_MAX_F32: lambda a, b, d, l, lit: i32(max(f32(a), f32(b))), VOPDOp.V_DUAL_MIN_F32: lambda a, b, d, l, lit: i32(min(f32(a), f32(b))),
VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: lambda a, b, d, l, lit: i32(0.0 if f32(a) == 0.0 or f32(b) == 0.0 else f32(a)*f32(b)),
VOPDOp.V_DUAL_MOV_B32: lambda a, b, d, l, lit: a, VOPDOp.V_DUAL_ADD_NC_U32: lambda a, b, d, l, lit: (a + b) & 0xffffffff,
VOPDOp.V_DUAL_LSHLREV_B32: lambda a, b, d, l, lit: (b << (a & 0x1f)) & 0xffffffff, VOPDOp.V_DUAL_AND_B32: lambda a, b, d, l, lit: a & b,
VOPDOp.V_DUAL_FMAC_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(d)), VOPDOp.V_DUAL_FMAAK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(lit)),
VOPDOp.V_DUAL_FMAMK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(lit)+f32(b)), VOPDOp.V_DUAL_CNDMASK_B32: lambda a, b, d, l, lit: b if l else a,
}
def exec_vopd(st: WaveState, inst: VOPD, lane: int) -> None:
V, vdsty, vcc_lane = st.vgpr[lane], (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), (st.vcc >> lane) & 1
sx0, sx1, sy0, sy1, dstx = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], inst.vdstx
if (fn := VOPD_OPS.get(inst.opx)): V[dstx] = fn(sx0, sx1, V[dstx], vcc_lane, st.literal)
else: raise NotImplementedError(f"VOPD opx {inst.opx}")
if (fn := VOPD_OPS.get(inst.opy)): V[vdsty] = fn(sy0, sy1, V[vdsty], vcc_lane, st.literal)
else: raise NotImplementedError(f"VOPD opy {inst.opy}")
def exec_vop3p(st: WaveState, inst: VOP3P, lane: int) -> None:
op, vdst, V = inst.op, inst.vdst, st.vgpr[lane]
s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
opsel, opsel_hi = [(inst.opsel >> i) & 1 for i in range(3)], [(inst.opsel_hi >> i) & 1 for i in range(2)] + [inst.opsel_hi2]
neg, neg_hi = inst.neg, inst.neg_hi
def get_src(src: int, idx: int, for_mix: bool = False) -> float:
if for_mix:
if not opsel_hi[idx]: return abs(f32(src)) if (neg_hi >> idx) & 1 else f32(src)
return float(f16((src >> 16) & 0xffff) if opsel[idx] else f16(src & 0xffff))
use_hi = opsel[idx]
val = ((src >> 16) & 0xffff) if use_hi else (src & 0xffff)
f = f16(val)
if use_hi and (neg >> idx) & 1: f = -f
elif not use_hi and (neg_hi >> idx) & 1: f = -f
return f
if op == VOP3POp.V_FMA_MIX_F32: V[vdst] = i32(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
elif op == VOP3POp.V_FMA_MIXLO_F16: V[vdst] = (V[vdst] & 0xffff0000) | i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
elif op == VOP3POp.V_FMA_MIXHI_F16: V[vdst] = (V[vdst] & 0x0000ffff) | (i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True)) << 16)
else: raise NotImplementedError(f"VOP3P op {op}")
def exec_wmma_f32_16x16x16_f16(st: WaveState, inst: VOP3P, n_lanes: int) -> None:
src0_base, src1_base, src2_base = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0, (inst.src1 - 256) if inst.src1 >= 256 else inst.src1, (inst.src2 - 256) if inst.src2 >= 256 else inst.src2
src0_is_vgpr, src1_is_vgpr, src2_is_vgpr, vdst = inst.src0 >= 256, inst.src1 >= 256, inst.src2 >= 256, inst.vdst
A, B, C = [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)]
for lane in range(min(n_lanes, 16)):
V = st.vgpr[lane]
for reg in range(8):
val = V[src0_base + reg] if src0_is_vgpr else st.sgpr[src0_base + reg]
A[lane][reg * 2], A[lane][reg * 2 + 1] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
val = V[src1_base + reg] if src1_is_vgpr else st.sgpr[src1_base + reg]
B[reg * 2][lane], B[reg * 2 + 1][lane] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
for row in range(16):
for col in range(16):
idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
if lane_idx < n_lanes:
val = st.vgpr[lane_idx][src2_base + reg] if src2_is_vgpr else st.sgpr[src2_base + reg]
C[row][col] = f32(val)
for row in range(16):
for col in range(16):
for k in range(16): C[row][col] += A[row][k] * B[k][col]
for row in range(16):
for col in range(16):
idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
if lane_idx < n_lanes and (st.exec_mask & (1 << lane_idx)): st.vgpr[lane_idx][vdst + reg] = i32(C[row][col])
# ═══════════════════════════════════════════════════════════════════════════════
# MAIN EXECUTION LOOP
# ═══════════════════════════════════════════════════════════════════════════════
SCALAR: dict[type, Callable[..., int]] = {SOP1: exec_sop1, SOP2: exec_sop2, SOPC: exec_sopc, SOPK: exec_sopk, SOPP: exec_sopp, SMEM: exec_smem}
VECTOR: dict[type, Callable[..., None]] = {VOP1: exec_vop1, VOP2: exec_vop2, VOP3: exec_vop3, VOP3SD: exec_vop3sd, VOPC: exec_vopc, FLAT: exec_flat, DS: exec_ds, VOPD: exec_vopd, VOP3P: exec_vop3p}
_WMMA_OPS = frozenset((VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16,
VOP3POp.V_WMMA_BF16_16X16X16_BF16, VOP3POp.V_WMMA_I32_16X16X16_IU8, VOP3POp.V_WMMA_I32_16X16X16_IU4))
def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
inst = program.get(st.pc)
if inst is None: return 1
inst_words, st.literal, inst_type = inst._words, inst._literal or 0, type(inst)
if (handler := SCALAR.get(inst_type)) is not None:
delta = handler(st, inst)
if delta == -1: return -1
if delta == -2: st.pc += inst_words; return -2
if delta == -3: # S_GETPC_B64
sop1 = inst if isinstance(inst, SOP1) else None
assert sop1 is not None
next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc += inst_words; return 0
if delta == -4: # S_SETPC_B64
sop1 = inst if isinstance(inst, SOP1) else None
assert sop1 is not None
st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
if delta == -5: # S_SWAPPC_B64
sop1 = inst if isinstance(inst, SOP1) else None
assert sop1 is not None
next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
st.pc += inst_words + delta
else:
vec_handler, exec_mask = VECTOR[inst_type], st.exec_mask
if inst_type is DS:
for lane in range(n_lanes):
if exec_mask & (1 << lane): vec_handler(st, inst, lane, lds)
elif inst_type is VOP3P:
vop3p = inst if isinstance(inst, VOP3P) else None
assert vop3p is not None
if vop3p.op in _WMMA_OPS:
exec_wmma_f32_16x16x16_f16(st, vop3p, n_lanes)
else:
for lane in range(n_lanes):
if exec_mask & (1 << lane): vec_handler(st, vop3p, lane)
else:
for lane in range(n_lanes):
if exec_mask & (1 << lane): vec_handler(st, inst, lane)
st.commit_pends(); st.pc += inst_words
return 0
def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int, wg_id: tuple[int,int,int]=(0,0,0), local_size: tuple[int,int,int]=(1,1,1), wave_start: int=0) -> int:
while st.pc in program:
result = step_wave(program, st, lds, n_lanes)
if result == -1: return 0
if result == -2: return -2
return 0
def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, dispatch_dim: int) -> None:
lx, ly, lz = local_size
total_threads, lds = lx * ly * lz, bytearray(65536)
waves: list[tuple[WaveState, int, int]] = []
for wave_start in range(0, total_threads, WAVE_SIZE):
n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
st.exec_mask = (1 << n_lanes) - 1
st.wsgpr64(0, args_ptr)
gx, gy, gz = workgroup_id
if dispatch_dim >= 3: st.sgpr[13], st.sgpr[14], st.sgpr[15] = gx, gy, gz
elif dispatch_dim == 2: st.sgpr[14], st.sgpr[15] = gx, gy
else: st.sgpr[15] = gx
for i in range(n_lanes):
tid = wave_start + i
st.vgpr[i][0] = tid if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
waves.append((st, n_lanes, wave_start))
has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
for _ in range(2 if has_barrier else 1):
for st, n_lanes, wave_start in waves: exec_wave(program, st, lds, n_lanes, workgroup_id, local_size, wave_start)
def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
data = (ctypes.c_char * lib_sz).from_address(lib).raw
program = decode_program(data)
if not program: return -1
dispatch_dim = 3 if gz > 1 else (2 if gy > 1 else 1)
for gidz in range(gz):
for gidy in range(gy):
for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, dispatch_dim)
return 0
-191
View File
@@ -1,191 +0,0 @@
#!/usr/bin/env python3
# generates autogen/__init__.py by parsing the AMD RDNA3.5 ISA PDF
import re, pdfplumber, pathlib
from tinygrad.helpers import fetch
PDF_URL = "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"
FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
FIELD_ORDER = {
'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
'4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
def parse_bits(s: str) -> tuple[int, int] | None:
return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
def parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
fields = []
for row in table[1:]:
if not row or not row[0]: continue
name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
if not (bits := parse_bits(bits_str)): continue
enc_val, hi, lo = None, bits[0], bits[1]
if name == 'ENCODING' and row[2] and (m := re.search(r"'b([01_]+)", row[2])):
enc_bits = m.group(1).replace('_', '')
enc_val = int(enc_bits, 2)
declared_width, actual_width = hi - lo + 1, len(enc_bits)
if actual_width > declared_width: lo = hi - actual_width + 1
ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
fields.append((name, hi, lo, enc_val, ftype))
return fields
def generate(output_path: pathlib.Path|str|None = None) -> dict:
"""Generate RDNA3.5 instruction definitions from the AMD ISA PDF. Returns dict with formats for testing."""
pdf = pdfplumber.open(fetch(PDF_URL))
pages = pdf.pages[150:200]
page_texts = [p.extract_text() or '' for p in pages]
page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
full_text = '\n'.join(page_texts)
# parse SSRC encoding from first page with VCC_LO
src_enum = dict(SRC_EXTRAS)
for text in page_texts[:10]:
if 'SSRC0' in text and 'VCC_LO' in text:
for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
val, name = int(m.group(1)), m.group(2).rstrip('.:')
if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
break
# parse opcode tables
enums: dict[str, dict[int, str]] = {}
for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
enums[m.group(1) + "Op"] = ops
if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
enums["VOPDOp"] = ops
enum_names = set(enums.keys())
def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
def has_header_before_fields(text) -> bool:
return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
# find format headers with their page indices
format_headers = [] # (fmt_name, page_idx)
for i, text in enumerate(page_texts):
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
if m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < len(page_texts):
next_text = page_texts[i + 1].lstrip()
if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
format_headers.append((m.group(1), i, m.start()))
# parse instruction formats
formats: dict[str, list] = {}
for fmt_name, page_idx, header_pos in format_headers:
if fmt_name in formats: continue
text, tables = page_texts[page_idx], page_tables[page_idx]
field_pos = text.find('Field Name', header_pos)
# find fields table with ENCODING (same page or up to 2 pages ahead)
fields = None
for offset in range(3):
if page_idx + offset >= len(pages): break
if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
fields = f
break
if fields: break
# for modifier formats (no ENCODING), accept first fields table on same page
if not fields and field_pos > header_pos:
for t in tables:
if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)):
fields = f
break
if not fields: continue
field_names = {f[0] for f in fields}
# check next pages for continuation fields (tables without ENCODING)
for pg_offset in range(1, 3):
if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
for t in page_tables[page_idx + pg_offset]:
if is_fields_table(t) and (extra := parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
for ef in extra:
if ef[0] not in field_names:
fields.append(ef)
field_names.add(ef[0])
break
formats[fmt_name] = fields
# fix known PDF errors (verified against LLVM test vectors)
# SMEM: PDF says DLC=bit14, GLC=bit16 but actual encoding is DLC=bit13, GLC=bit14
if 'SMEM' in formats:
formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
for n, h, l, e, t in formats['SMEM']]
# generate output
def enum_lines(name, items):
return [f"class {name}(IntEnum):"] + [f" {n} = {v}" for v, n in sorted(items.items())] + [""]
def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
lines = ["# autogenerated from AMD RDNA3.5 ISA PDF by gen.py - do not edit", "from enum import IntEnum",
"from typing import Annotated",
"from extra.assembly.rdna3.lib import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
"import functools", ""]
lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
# Format-specific field defaults (verified against LLVM test vectors)
format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
lines.append("# instruction formats")
for fmt_name, fields in sorted(formats.items()):
base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
order = FIELD_ORDER.get(fmt_name, [])
lines.append(f"class {fmt_name}({base}):")
if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
lines.append(f" encoding = {enc_str}")
if defaults := format_defaults.get(fmt_name):
lines.append(f" _defaults = {defaults}")
for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
# Wrap IntEnum types (ending in Op) with Annotated[BitField, ...] for correct typing
if ftype and ftype.endswith('Op'):
ann = f":Annotated[BitField, {ftype}]"
else:
ann = f":{ftype}" if ftype else ""
lines.append(f" {name.lower()}{ann} = bits[{hi}]" if hi == lo else f" {name.lower()}{ann} = bits[{hi}:{lo}]")
lines.append("")
lines.append("# instruction helpers")
for cls_name, ops in sorted(enums.items()):
fmt = cls_name[:-2]
for op_val, name in sorted(ops.items()):
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=2"}.get(fmt, "")
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
# VOP1/VOP2/VOPC get _e32 suffix, VOP3 promoted ops (< 512) get _e64 suffix
if fmt in ("VOP1", "VOP2", "VOPC"):
suffix = "_e32"
elif fmt == "VOP3" and op_val < 512:
suffix = "_e64"
else:
suffix = ""
lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
# export SrcEnum values, but skip DPP8/DPP16 which conflict with class names
skip_exports = {'DPP8', 'DPP16'}
lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in skip_exports] + ["OFF = NULL\n"]
if output_path is not None: pathlib.Path(output_path).write_text('\n'.join(lines))
return {"formats": formats, "enums": enums, "src_enum": src_enum}
if __name__ == "__main__":
result = generate("extra/assembly/rdna3/autogen/__init__.py")
print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")
-248
View File
@@ -1,248 +0,0 @@
# library for RDNA3 assembly DSL
from __future__ import annotations
from enum import IntEnum
from typing import overload, Annotated, TypeVar, Generic
# Bit field DSL
class BitField:
def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name = hi, lo, name
def __set_name__(self, owner, name): self.name, self._owner = name, owner
def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val) # type: ignore
def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
@property
def marker(self) -> type | None:
# Get marker from Annotated type hint if present
import typing
if hasattr(self, '_owner') and self.name:
hints = typing.get_type_hints(self._owner, include_extras=True)
if self.name in hints:
hint = hints[self.name]
if typing.get_origin(hint) is Annotated:
args = typing.get_args(hint)
return args[1] if len(args) > 1 else None
return None
@overload
def __get__(self, obj: None, objtype: type) -> BitField: ...
@overload
def __get__(self, obj: object, objtype: type | None = None) -> int: ...
def __get__(self, obj, objtype=None):
if obj is None: return self
val = unwrap(obj._values.get(self.name, 0))
# Convert to IntEnum if marker is an IntEnum subclass
if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
try: return self.marker(val)
except ValueError: pass
return val
class _Bits:
def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
bits = _Bits()
# Register types
class Reg:
def __init__(self, idx: int, count: int = 1, hi: bool = False): self.idx, self.count, self.hi = idx, count, hi
def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
T = TypeVar('T', bound=Reg)
class _RegFactory(Generic[T]):
def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
@overload
def __getitem__(self, key: int) -> Reg: ...
@overload
def __getitem__(self, key: slice) -> Reg: ...
def __getitem__(self, key: int | slice) -> Reg:
return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
def __repr__(self): return f"<{self._name} factory>"
class SGPR(Reg): pass
class VGPR(Reg): pass
class TTMP(Reg): pass
s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
# Field type markers (runtime classes for validation)
class _SSrc: pass
class _Src: pass
class _Imm: pass
class _SImm: pass
class _VDSTYEnc: pass # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
class _SGPRField: pass
class _VGPRField: pass
# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
SSrc = Annotated[BitField, _SSrc]
Src = Annotated[BitField, _Src]
Imm = Annotated[BitField, _Imm]
SImm = Annotated[BitField, _SImm]
VDSTYEnc = Annotated[BitField, _VDSTYEnc]
SGPRField = Annotated[BitField, _SGPRField]
VGPRField = Annotated[BitField, _VGPRField]
class RawImm:
def __init__(self, val: int): self.val = val
def __repr__(self): return f"RawImm({self.val})"
def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
def unwrap(val) -> int:
return val.val if isinstance(val, RawImm) else val.value if hasattr(val, 'value') else val.idx if hasattr(val, 'idx') else val
# Encoding helpers
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata'}
def _encode_reg(val) -> int:
if isinstance(val, TTMP): return 108 + val.idx
return val.idx | (0x80 if val.hi else 0)
def encode_src(val) -> int:
if isinstance(val, VGPR): return 256 + _encode_reg(val)
if isinstance(val, Reg): return _encode_reg(val)
if hasattr(val, 'value'): return val.value
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
# Instruction base class
class Inst:
_fields: dict[str, BitField]
_encoding: tuple[BitField, int] | None = None
_defaults: dict[str, int] = {}
_values: dict[str, int | RawImm]
_words: int # size in 32-bit words, set by decode_program
_literal: int | None
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
def __init__(self, *args, literal: int | None = None, **kwargs):
self._values, self._literal = dict(self._defaults), literal
# Map positional args to field names
field_names = [n for n in self._fields if n != 'encoding']
orig_args = dict(zip(field_names, args))
orig_args.update(kwargs)
self._values.update(orig_args)
# Validate register counts for SMEM instructions (before encoding)
if self.__class__.__name__ == 'SMEM':
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
if op_val is not None:
if hasattr(op_val, 'value'): op_val = op_val.value
expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
sdata_val = orig_args.get('sdata')
if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
# Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
if self.__class__.__name__ == 'SOP1':
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
if op_val is not None and hasattr(op_val, 'name'):
expected = 2 if op_val.name.endswith('_B64') else 1
sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
if isinstance(sdst_val, Reg) and sdst_val.count != expected:
raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
# Type check and encode values
for name, val in list(self._values.items()):
if name == 'encoding': continue
# For RawImm, only process RAW_FIELDS to unwrap to int
if isinstance(val, RawImm):
if name in RAW_FIELDS: self._values[name] = val.val
continue
field = self._fields.get(name)
marker = field.marker if field else None
# Type validation
if marker is _SGPRField:
if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
if not isinstance(val, (SGPR, TTMP, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
if marker is _VGPRField:
if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
# Encode source fields as RawImm for consistent disassembly
if name in SRC_FIELDS:
encoded = encode_src(val)
self._values[name] = RawImm(encoded)
# Track literal value if needed (encoded as 255)
if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum):
self._literal = val
elif encoded == 255 and self._literal is None and isinstance(val, float):
import struct
self._literal = struct.unpack('<I', struct.pack('<f', val))[0]
# Encode raw register fields for consistent repr
elif name in RAW_FIELDS:
if isinstance(val, Reg): self._values[name] = _encode_reg(val)
elif hasattr(val, 'value'): self._values[name] = val.value # IntEnum like SrcEnum.NULL
# Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
elif name == 'sbase' and isinstance(val, Reg):
self._values[name] = val.idx // 2
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
self._values[name] = val.idx // 4
# VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
elif marker is _VDSTYEnc and isinstance(val, VGPR):
self._values[name] = val.idx >> 1
def _encode_field(self, name: str, val) -> int:
if isinstance(val, RawImm): return val.val
if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
return val.value if hasattr(val, 'value') else val
def to_int(self) -> int:
word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
for n, bf in self._fields.items():
if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
return word
def _get_literal(self) -> int | None:
for n in SRC_FIELDS:
if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
return None
def to_bytes(self) -> bytes:
result = self.to_int().to_bytes(self._size(), 'little')
return result + (lit & 0xffffffff).to_bytes(4, 'little') if (lit := self._get_literal() or getattr(self, '_literal', None)) else result
@classmethod
def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
def size(self) -> int: return self._size() + (4 if self._literal is not None else 0)
@classmethod
def from_int(cls, word: int):
inst = object.__new__(cls)
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
inst._literal = None
return inst
@classmethod
def from_bytes(cls, data: bytes):
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
op_val = inst._values.get('op', 0)
has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
for n in SRC_FIELDS:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
if has_literal and len(data) >= cls._size() + 4: inst._literal = int.from_bytes(data[cls._size():cls._size()+4], 'little')
return inst
def __repr__(self):
# Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
and not (is_zero(self._values[k]) and k not in {'op'})]
lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
def __eq__(self, other):
if not isinstance(other, Inst): return NotImplemented
return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
def disasm(self) -> str:
from extra.assembly.rdna3.asm import disasm
return disasm(self)
class Inst32(Inst): pass
class Inst64(Inst): pass
-845
View File
@@ -1,845 +0,0 @@
# Unit tests for RDNA3 Python emulator
import unittest
import ctypes
import struct
import math
from extra.assembly.rdna3.emu import (
WaveState, decode_program, exec_wave, exec_workgroup, run_asm,
i32, f32, sext, WAVE_SIZE, set_valid_mem_ranges
)
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import RawImm
def run_kernel(kernel: bytes, n_threads: int = 1, n_outputs: int = 1) -> list[int]:
"""Helper to run a kernel and return output values."""
output = (ctypes.c_uint32 * (n_threads * n_outputs))(*[0xdead] * (n_threads * n_outputs))
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
# Register valid memory ranges for bounds checking
set_valid_mem_ranges({
(output_ptr, ctypes.sizeof(output)),
(args_ptr, ctypes.sizeof(args)),
(kernel_ptr, len(kernel)),
})
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
assert result == 0, f"run_asm failed with {result}"
return [output[i] for i in range(n_threads * n_outputs)]
def make_store_kernel(setup_instrs: list, store_vreg: int = 1) -> bytes:
"""Create a kernel that runs setup instructions then stores v[store_vreg] to output[tid]."""
kernel = b''
# Load output pointer
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
# Run setup instructions
for instr in setup_instrs:
kernel += instr.to_bytes()
# Compute offset: v3 = tid * 4
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
# Store result
kernel += global_store_b32(addr=v[3], data=v[store_vreg], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
return kernel
class TestScalarOps(unittest.TestCase):
def test_s_mov_b32(self):
state = WaveState()
kernel = s_mov_b32(s[5], 42).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[5], 42)
def test_s_add_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 100, 50
kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 150)
self.assertEqual(state.scc, 0) # no carry
def test_s_add_u32_carry(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 0xffffffff, 1
kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 0)
self.assertEqual(state.scc, 1) # carry
def test_s_sub_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 100, 30
kernel = s_sub_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 70)
self.assertEqual(state.scc, 0) # no borrow
def test_s_and_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 0xff00, 0x0ff0
kernel = s_and_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 0x0f00)
def test_s_or_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 0xff00, 0x00ff
kernel = s_or_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 0xffff)
def test_s_lshl_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 1, 4
kernel = s_lshl_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 16)
def test_s_lshr_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 256, 4
kernel = s_lshr_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 16)
def test_s_mul_i32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 7, 6
kernel = s_mul_i32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 42)
def test_s_cmp_eq_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 42, 42
kernel = s_cmp_eq_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.scc, 1)
def test_s_cmp_lg_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 42, 43
kernel = s_cmp_lg_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.scc, 1)
class TestVectorOps(unittest.TestCase):
def test_v_mov_b32(self):
kernel = make_store_kernel([v_mov_b32_e32(v[1], 42)])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_add_nc_u32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 10),
v_mov_b32_e32(v[2], 32),
v_add_nc_u32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_sub_nc_u32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 50),
v_mov_b32_e32(v[2], 8),
v_sub_nc_u32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_mul_lo_u32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 6),
v_mov_b32_e32(v[2], 7),
v_mul_lo_u32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_and_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 0xff0f),
v_mov_b32_e32(v[2], 0x0fff),
v_and_b32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [0x0f0f])
def test_v_or_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 0xf000),
v_mov_b32_e32(v[2], 0x000f),
v_or_b32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [0xf00f])
def test_v_lshlrev_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 1),
v_lshlrev_b32_e32(v[1], 5, v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [32])
def test_v_lshrrev_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 128),
v_lshrrev_b32_e32(v[1], 3, v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [16])
def test_v_add_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(1.5)),
v_mov_b32_e32(v[2], i32(2.5)),
v_add_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 4.0)
def test_v_mul_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.0)),
v_mov_b32_e32(v[2], i32(4.0)),
v_mul_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 12.0)
def test_v_max_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.0)),
v_mov_b32_e32(v[2], i32(5.0)),
v_max_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 5.0)
def test_v_min_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.0)),
v_mov_b32_e32(v[2], i32(5.0)),
v_min_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 3.0)
class TestThreading(unittest.TestCase):
def test_thread_id(self):
"""Each thread should get its own thread ID in v0."""
kernel = make_store_kernel([v_mov_b32_e32(v[1], v[0])], store_vreg=1)
out = run_kernel(kernel, n_threads=4)
self.assertEqual(out, [0, 1, 2, 3])
def test_thread_local_ops(self):
"""Each thread computes tid * 10."""
kernel = make_store_kernel([
v_mov_b32_e32(v[2], 10),
v_mul_lo_u32(v[1], v[0], v[2]),
])
out = run_kernel(kernel, n_threads=4)
self.assertEqual(out, [0, 10, 20, 30])
def test_exec_mask(self):
"""Test that exec mask controls which lanes execute."""
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], 0, soffset=NULL).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_mov_b32_e32(v[1], 100).to_bytes() # default value
kernel += s_mov_b32(EXEC_LO, 0b0101).to_bytes() # only lanes 0 and 2
kernel += v_mov_b32_e32(v[1], 42).to_bytes() # only for active lanes
kernel += s_mov_b32(EXEC_LO, 0xf).to_bytes() # restore all lanes
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
out = run_kernel(kernel, n_threads=4)
self.assertEqual(out, [42, 100, 42, 100])
class TestBranching(unittest.TestCase):
def test_s_branch(self):
"""Test unconditional branch."""
state = WaveState()
kernel = b''
kernel += s_mov_b32(s[0], 1).to_bytes()
kernel += s_branch(1).to_bytes() # skip next instruction
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
kernel += s_mov_b32(s[1], 3).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[0], 1) # not overwritten
self.assertEqual(state.sgpr[1], 3)
def test_s_cbranch_scc0(self):
"""Test conditional branch on SCC=0."""
state = WaveState()
state.scc = 0
kernel = b''
kernel += s_mov_b32(s[0], 1).to_bytes()
kernel += s_cbranch_scc0(1).to_bytes() # branch if scc=0
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[0], 1)
def test_s_cbranch_scc1(self):
"""Test conditional branch on SCC=1."""
state = WaveState()
state.scc = 1
kernel = b''
kernel += s_mov_b32(s[0], 1).to_bytes()
kernel += s_cbranch_scc1(1).to_bytes() # branch if scc=1
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[0], 1)
def test_unknown_sopp_opcode(self):
"""Regression test: unknown SOPP opcodes should be ignored, not crash."""
state = WaveState()
# Create a raw SOPP instruction with opcode 8 (undefined in our enum)
# SOPP format: bits[31:23] = 0b101111111, bits[22:16] = op, bits[15:0] = simm16
unknown_sopp = (0b101111111 << 23) | (8 << 16) | 0 # op=8, simm16=0
kernel = unknown_sopp.to_bytes(4, 'little') + s_endpgm().to_bytes()
prog = decode_program(kernel)
# Should not raise an exception
exec_wave(prog, state, bytearray(65536), 1)
class TestMemory(unittest.TestCase):
def test_global_load_store(self):
"""Test global load followed by store."""
# Create input buffer
input_buf = (ctypes.c_uint32 * 4)(10, 20, 30, 40)
input_ptr = ctypes.addressof(input_buf)
output_buf = (ctypes.c_uint32 * 4)(*[0]*4)
output_ptr = ctypes.addressof(output_buf)
args = (ctypes.c_uint64 * 2)(output_ptr, input_ptr)
args_ptr = ctypes.addressof(args)
# Kernel: load from input[tid], add 1, store to output[tid]
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes() # output ptr
kernel += s_load_b64(s[4:5], s[0:1], soffset=NULL, offset=8).to_bytes() # input ptr
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_lshlrev_b32_e32(v[2], 2, v[0]).to_bytes() # offset = tid * 4
kernel += global_load_b32(vdst=v[1], addr=v[2], saddr=s[4]).to_bytes()
kernel += s_waitcnt(vmcnt=0).to_bytes()
kernel += v_add_nc_u32_e32(v[1], 1, v[1]).to_bytes() # add 1
kernel += global_store_b32(addr=v[2], data=v[1], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({
(input_ptr, ctypes.sizeof(input_buf)),
(output_ptr, ctypes.sizeof(output_buf)),
(args_ptr, ctypes.sizeof(args)),
(kernel_ptr, len(kernel)),
})
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, 4, 1, 1, args_ptr)
self.assertEqual(result, 0)
self.assertEqual([output_buf[i] for i in range(4)], [11, 21, 31, 41])
class TestFloatOps(unittest.TestCase):
def test_v_rcp_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(4.0)),
v_rcp_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertAlmostEqual(f32(out[0]), 0.25, places=5)
def test_v_sqrt_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(16.0)),
v_sqrt_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertAlmostEqual(f32(out[0]), 4.0, places=5)
def test_v_floor_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.7)),
v_floor_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 3.0)
def test_v_ceil_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.2)),
v_ceil_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 4.0)
def test_v_cvt_f32_i32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 42),
v_cvt_f32_i32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 42.0)
def test_v_cvt_i32_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(42.9)),
v_cvt_i32_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 42)
class TestVOP3(unittest.TestCase):
def test_v_fma_f32(self):
"""Test fused multiply-add: a*b + c"""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(2.0)),
v_mov_b32_e32(v[2], i32(3.0)),
v_mov_b32_e32(v[4], i32(4.0)),
v_fma_f32(v[1], v[1], v[2], v[4]), # 2*3+4 = 10
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 10.0)
def test_v_add3_u32(self):
"""Test 3-operand add."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 10),
v_mov_b32_e32(v[2], 20),
v_mov_b32_e32(v[4], 12),
v_add3_u32(v[1], v[1], v[2], v[4]), # 10+20+12 = 42
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 42)
def test_v_neg_modifier(self):
"""Test VOP3 negation modifier."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(5.0)),
v_mov_b32_e32(v[2], i32(3.0)),
# v_add_f32 with neg on src1: 5 + (-3) = 2
v_add_f32_e64(v[1], v[1], v[2], neg=0b010),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 2.0)
def test_v_ldexp_f32(self):
"""Regression test: V_LDEXP_F32 used by exp()."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(1.5)),
v_mov_b32_e32(v[2], 3), # exponent
v_ldexp_f32(v[1], v[1], v[2]), # 1.5 * 2^3 = 12.0
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 12.0)
def test_v_xad_u32(self):
"""Regression test: V_XAD_U32 (xor-add) used by random number generation."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 3),
v_mov_b32_e32(v[2], 4),
v_mov_b32_e32(v[4], 5),
v_xad_u32(v[1], v[1], v[2], v[4]), # (3^4)+5 = 7+5 = 12
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 12)
def test_v_lshl_or_b32(self):
"""Regression test: V_LSHL_OR_B32 operand order is (s0 << s1) | s2, not (s0 << s2) | s1."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 5), # s0 = value to shift
v_mov_b32_e32(v[2], 2), # s1 = shift amount
v_mov_b32_e32(v[4], 3), # s2 = value to OR
v_lshl_or_b32(v[1], v[1], v[2], v[4]), # (5 << 2) | 3 = 20 | 3 = 23
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 23)
def test_v_sqrt_f32_negative(self):
"""Regression test: V_SQRT_F32 should return NaN for negative inputs, not 0."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(-1.0)),
v_sqrt_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertTrue(math.isnan(f32(out[0])))
def test_v_rsq_f32_negative(self):
"""Regression test: V_RSQ_F32 should return NaN for negative inputs, not inf."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(-1.0)),
v_rsq_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertTrue(math.isnan(f32(out[0])))
class TestVOPD(unittest.TestCase):
def test_vopd_add_nc_u32(self):
"""Test VOPD V_DUAL_ADD_NC_U32."""
state = WaveState()
state.vgpr[0][1] = 100
state.vgpr[0][2] = 50
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][3], 100) # MOV result
self.assertEqual(state.vgpr[0][4], 150) # 100 + 50
def test_vopd_lshlrev(self):
"""Test VOPD V_DUAL_LSHLREV_B32."""
state = WaveState()
state.vgpr[0][1] = 0x10
state.vgpr[0][2] = 0
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
opy=VOPDOp.V_DUAL_LSHLREV_B32, srcy0=4, vsrcy1=VGPR(1), vdsty=VGPR(4)).to_bytes() # V4 = V1 << 4
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][3], 0x10) # MOV result
self.assertEqual(state.vgpr[0][4], 0x100) # 0x10 << 4 = 0x100
def test_vopd_and(self):
"""Test VOPD V_DUAL_AND_B32."""
state = WaveState()
state.vgpr[0][1] = 0xff
state.vgpr[0][2] = 0x0f
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
opy=VOPDOp.V_DUAL_AND_B32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][3], 0xff)
self.assertEqual(state.vgpr[0][4], 0x0f) # 0xff & 0x0f = 0x0f
def test_vopd_parallel_read(self):
"""Regression: VOPD must read all inputs before writing - Y op reads register that X op writes."""
state = WaveState()
state.vgpr[0][4] = 0
state.vgpr[0][7] = 5 # Y op reads v7 as vsrcy1, X op writes to v7
# X: MOV v7, v0 (v0=0, so v7 becomes 0)
# Y: ADD v6, v4, v7 (should use original v7=5, not the overwritten 0)
# vdsty_enc=3 with vdstx=7 (odd) -> vdsty = (3 << 1) | (7&1)^1 = 6 | 0 = 6
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[0], vsrcx1=VGPR(0), vdstx=VGPR(7),
opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[4], vsrcy1=VGPR(7), vdsty=VGPR(6)).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][7], 0) # X op: v7 = v0 = 0
self.assertEqual(state.vgpr[0][6], 5) # Y op: v6 = v4 + v7 = 0 + 5 (original v7)
class TestDecoder(unittest.TestCase):
def test_vopd_literal_handling(self):
"""Regression test: VOPD srcx0/srcy0 with literal (255) wasn't consuming the literal dword."""
state = WaveState()
# Create VOPD with srcx0=255 (literal), followed by literal value 0x12345678
vopd_bytes = VOPD(opx=8, srcx0=RawImm(255), vsrcx1=VGPR(0), vdstx=VGPR(1), # MOV: V1 = literal
opy=8, srcy0=RawImm(128), vsrcy1=VGPR(0), vdsty=VGPR(2)).to_bytes() # MOV: V2 = 0
literal_bytes = (0x12345678).to_bytes(4, 'little')
kernel = vopd_bytes + literal_bytes + s_endpgm().to_bytes()
prog = decode_program(kernel)
# Should decode as 3 instructions: VOPD (with literal), then S_ENDPGM
# The literal should NOT be decoded as a separate instruction
self.assertEqual(len(prog), 2) # VOPD + S_ENDPGM
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][1], 0x12345678)
def test_s_endpgm_stops_decode(self):
"""Regression test: decoder should stop at S_ENDPGM, not read past into metadata."""
# Create a kernel followed by garbage that looks like an invalid instruction
kernel = s_mov_b32(s[0], 42).to_bytes() + s_endpgm().to_bytes()
garbage = bytes([0xff] * 16) # garbage after kernel
prog = decode_program(kernel + garbage)
# Should only have 2 instructions (s_mov_b32 and s_endpgm)
self.assertEqual(len(prog), 2)
class TestFloatConversion(unittest.TestCase):
"""Unit tests for i32/i16/f32/f16 float conversion functions."""
def test_i32_preserves_nan_sign(self):
"""NaN sign bit should be preserved when converting float to int bits."""
from extra.assembly.rdna3.emu import i32, f32
# 0 * -inf produces a negative NaN
neg_nan = 0.0 * float('-inf')
bits = i32(neg_nan)
# Should have sign bit set (0xffc00000), not canonical positive NaN (0x7fc00000)
self.assertEqual(bits & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{bits:08x}")
self.assertTrue(math.isnan(f32(bits)))
def test_i32_preserves_positive_nan(self):
"""Positive NaN should remain positive."""
from extra.assembly.rdna3.emu import i32, f32
pos_nan = float('nan')
bits = i32(pos_nan)
# Standard Python NaN is positive (0x7fc00000)
self.assertEqual(bits & 0x80000000, 0, f"Expected positive NaN, got 0x{bits:08x}")
self.assertTrue(math.isnan(f32(bits)))
def test_i32_overflow_to_inf(self):
"""Values too large for f32 should become inf."""
from extra.assembly.rdna3.emu import i32, f32
big = 2.0 ** 200
self.assertEqual(i32(big), 0x7f800000) # +inf
self.assertEqual(i32(-big), 0xff800000) # -inf
def test_i32_inf(self):
"""Infinity should be preserved."""
from extra.assembly.rdna3.emu import i32
self.assertEqual(i32(float('inf')), 0x7f800000)
self.assertEqual(i32(float('-inf')), 0xff800000)
def test_i32_normal_values(self):
"""Normal float values should round-trip correctly (within f32 precision)."""
from extra.assembly.rdna3.emu import i32, f32
# Use values exactly representable in float32
for val in [0.0, 1.0, -1.0, 0.5, -0.5, 100.0, -100.0, 1e10]:
bits = i32(val)
self.assertAlmostEqual(f32(bits), val, places=5)
def test_i16_overflow_to_inf(self):
"""Values too large for f16 should become inf."""
from extra.assembly.rdna3.emu import i16
big = 100000.0 # way larger than f16 max (65504)
self.assertEqual(i16(big), 0x7c00) # +inf
self.assertEqual(i16(-big), 0xfc00) # -inf
def test_i16_inf(self):
"""Infinity should be preserved."""
from extra.assembly.rdna3.emu import i16
self.assertEqual(i16(float('inf')), 0x7c00)
self.assertEqual(i16(float('-inf')), 0xfc00)
def test_fma_nan_sign_preserved(self):
"""FMA producing NaN should preserve the correct sign bit."""
from extra.assembly.rdna3.emu import i32, f32
# 0 * (-inf) + 1.0 = NaN (from 0 * -inf)
a, b, c = 0.0, float('-inf'), 1.0
result = i32(a * b + c)
# The NaN should be negative since 0 * -inf produces negative NaN
self.assertEqual(result & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{result:08x}")
class TestMultiWave(unittest.TestCase):
def test_all_waves_execute(self):
"""Regression test: all waves in a workgroup must execute, not just the first."""
n_threads = 64 # 2 waves of 32 threads each
output = (ctypes.c_uint32 * n_threads)(*[0xdead] * n_threads)
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
# Simple kernel: store tid to output[tid]
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_lshlrev_b32_e32(v[1], 2, v[0]).to_bytes() # offset = tid * 4
kernel += global_store_b32(addr=v[1], data=v[0], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({
(output_ptr, ctypes.sizeof(output)),
(args_ptr, ctypes.sizeof(args)),
(kernel_ptr, len(kernel)),
})
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
self.assertEqual(result, 0)
# All threads should have written their tid
for i in range(n_threads):
self.assertEqual(output[i], i, f"Thread {i} didn't execute")
class TestRegressions(unittest.TestCase):
"""Regression tests for bugs fixed in the emulator."""
def test_v_fmac_f16(self):
"""V_FMAC_F16: fused multiply-add for FP16. Regression for VOP2 op 54."""
from extra.assembly.rdna3.emu import i16, f16
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i16(2.0)), # v1.lo = 2.0 (fp16)
v_mov_b32_e32(v[2], i16(3.0)), # v2.lo = 3.0 (fp16)
# v1 = v1 * v2 + v1 = 2.0 * 3.0 + 2.0 = 8.0
VOP2(VOP2Op.V_FMAC_F16, v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertAlmostEqual(f16(out[0] & 0xffff), 8.0, places=2)
def test_v_cvt_f64_f32(self):
"""V_CVT_F64_F32: convert float32 to float64. Regression for VOP1 op 16."""
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_mov_b32_e32(v[1], i32(3.14159)).to_bytes()
kernel += VOP1(VOP1Op.V_CVT_F64_F32, v[4], v[1]).to_bytes() # v4:v5 = f64(v1)
kernel += v_lshlrev_b32_e32(v[3], 3, v[0]).to_bytes() # offset = tid * 8
kernel += global_store_b64(addr=v[3], data=v[4], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
output = (ctypes.c_double * 1)(0.0)
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
self.assertAlmostEqual(output[0], 3.14159, places=4)
def test_v_add_f64(self):
"""V_ADD_F64: add two float64 values. Regression for VOP3 op 807."""
from extra.assembly.rdna3.emu import i64_parts
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
# Load 1.5 into v1:v2
lo, hi = i64_parts(1.5)
kernel += v_mov_b32_e32(v[1], lo).to_bytes()
kernel += v_mov_b32_e32(v[2], hi).to_bytes()
# Load 2.5 into v3:v4
lo, hi = i64_parts(2.5)
kernel += v_mov_b32_e32(v[3], lo).to_bytes()
kernel += v_mov_b32_e32(v[4], hi).to_bytes()
# v5:v6 = v1:v2 + v3:v4 = 1.5 + 2.5 = 4.0
kernel += VOP3(VOP3Op.V_ADD_F64, v[5], v[1], v[3]).to_bytes()
kernel += v_lshlrev_b32_e32(v[7], 3, v[0]).to_bytes()
kernel += global_store_b64(addr=v[7], data=v[5], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
output = (ctypes.c_double * 1)(0.0)
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
self.assertAlmostEqual(output[0], 4.0, places=10)
def test_flat_load_d16_hi_b16(self):
"""FLAT_LOAD_D16_HI_B16: load 16-bit to high half. Regression for FLAT op 35."""
from extra.assembly.rdna3.emu import i16
# Create a buffer with test data
src_data = (ctypes.c_uint16 * 1)(0x1234)
src_ptr = ctypes.addressof(src_data)
output = (ctypes.c_uint32 * 1)(0xABCD0000) # preset low bits
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 2)(output_ptr, src_ptr)
args_ptr = ctypes.addressof(args)
kernel = b''
kernel += s_load_b128(s[0:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_mov_b32_e32(v[1], 0xDEAD).to_bytes() # initial value with low bits set
kernel += v_mov_b32_e32(v[2], 0).to_bytes() # offset = 0
kernel += FLAT(FLATOp.FLAT_LOAD_D16_HI_B16, v[1], v[2], saddr=s[2], offset=0).to_bytes()
kernel += s_waitcnt(vmcnt=0).to_bytes()
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[0]).to_bytes()
kernel += s_endpgm().to_bytes()
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(output_ptr, 4), (src_ptr, 2), (args_ptr, 16), (kernel_ptr, len(kernel))})
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
# High 16 bits should be 0x1234, low 16 bits preserved as 0xDEAD
self.assertEqual(output[0], 0x1234DEAD)
def test_v_mad_u16(self):
"""V_MAD_U16: multiply-add unsigned 16-bit. Regression for VOP3 op 577."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 10), # a = 10
v_mov_b32_e32(v[2], 20), # b = 20
v_mov_b32_e32(v[4], 5), # c = 5
VOP3(VOP3Op.V_MAD_U16, v[1], v[1], v[2], v[4]), # v1 = 10*20+5 = 205
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0] & 0xffff, 205)
def test_v_lshrrev_b16(self):
"""V_LSHRREV_B16: logical shift right 16-bit. Regression for VOP3 op 825."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 0x8000), # value to shift
v_mov_b32_e32(v[2], 4), # shift amount
VOP3(VOP3Op.V_LSHRREV_B16, v[1], v[2], v[1]), # v1 = 0x8000 >> 4 = 0x0800
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0] & 0xffff, 0x0800)
def test_v_min_u16(self):
"""V_MIN_U16: minimum of two unsigned 16-bit values. Regression for VOP3 op 779."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 100),
v_mov_b32_e32(v[2], 50),
VOP3(VOP3Op.V_MIN_U16, v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0] & 0xffff, 50)
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply Accumulate) instructions."""
def test_wmma_f32_16x16x16_f16_identity(self):
"""V_WMMA_F32_16X16X16_F16 with identity matrix. Regression for VOP3P op 64."""
from extra.assembly.rdna3.emu import i16, f16, exec_wmma_f32_16x16x16_f16, WaveState
# Test using direct emulator call rather than full kernel to simplify
st = WaveState()
st.exec_mask = 0xffffffff # all 32 lanes active
# Set up A as identity matrix: A[i][i] = 1.0, rest = 0.0
# Lane i holds row i of A in 8 regs (2 fp16 per reg)
for lane in range(16):
for reg in range(8):
col0, col1 = reg * 2, reg * 2 + 1
val0 = i16(1.0) if col0 == lane else 0
val1 = i16(1.0) if col1 == lane else 0
st.vgpr[lane][0 + reg] = val0 | (val1 << 16) # src0 = v0:v7
# Set up B as identity matrix: lane i holds column i of B
for lane in range(16):
for reg in range(8):
row0, row1 = reg * 2, reg * 2 + 1
val0 = i16(1.0) if row0 == lane else 0
val1 = i16(1.0) if row1 == lane else 0
st.vgpr[lane][8 + reg] = val0 | (val1 << 16) # src1 = v8:v15
# Set up C as zeros
for lane in range(32):
for reg in range(8):
st.vgpr[lane][16 + reg] = 0 # src2 = v16:v23
# Create a fake VOP3P instruction
inst = VOP3P(VOP3POp.V_WMMA_F32_16X16X16_F16, v[24], src0=VGPR(0), src1=VGPR(8), src2=VGPR(16))
# Execute WMMA
exec_wmma_f32_16x16x16_f16(st, inst, 32)
# Check result: C should be identity (since A @ B where both are identity)
# Output i = row*16+col goes to lane (i%32), reg (i//32)
for row in range(16):
for col in range(16):
idx = row * 16 + col
lane, reg = idx % 32, idx // 32
result = st.vgpr[lane][24 + reg]
expected = 1.0 if row == col else 0.0
self.assertAlmostEqual(f32(result), expected, places=3,
msg=f"C[{row},{col}] = {f32(result)}, expected {expected}")
if __name__ == "__main__":
unittest.main()
-234
View File
@@ -1,234 +0,0 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import Inst
from extra.assembly.rdna3.asm import asm
# Instruction format detection based on encoding bits
def detect_format(data: bytes) -> type[Inst] | None:
"""Detect instruction format from machine code bytes."""
if len(data) < 4: return None
word = int.from_bytes(data[:4], 'little')
enc_9bit = (word >> 23) & 0x1FF # 9-bit encoding for SOP1/SOPC/SOPP
enc_8bit = (word >> 24) & 0xFF
# Check 9-bit encodings first (most specific)
if enc_9bit == 0x17D: return SOP1 # bits 31:23 = 101111101
if enc_9bit == 0x17E: return SOPC # bits 31:23 = 101111110
if enc_9bit == 0x17F: return SOPP # bits 31:23 = 101111111
# SOPK: bits 31:28 = 1011, bits 27:23 = opcode (check after SOP1/SOPC/SOPP)
if enc_8bit in range(0xB0, 0xC0): return SOPK
# SOP2: bits 31:23 in range 0x100-0x17C (0x80-0xBE in bits 31:24, but not SOPK)
if 0x80 <= enc_8bit <= 0x9F: return SOP2
# VOP1: bits 31:25 = 0111111 (0x3F)
if (word >> 25) == 0x3F: return VOP1
# VOPC: bits 31:25 = 0111110 (0x3E)
if (word >> 25) == 0x3E: return VOPC
# VOP2: bits 31:30 = 00
if (word >> 30) == 0: return VOP2
# Check 64-bit formats
if len(data) >= 8:
if enc_8bit in (0xD4, 0xD5, 0xD7): return VOP3
if enc_8bit == 0xD6: return VOP3SD
if enc_8bit == 0xCC: return VOP3P
if enc_8bit == 0xCD: return VINTERP
if enc_8bit in (0xC8, 0xC9): return VOPD
if enc_8bit == 0xF4: return SMEM
if enc_8bit == 0xD8: return DS
if enc_8bit in (0xDC, 0xDD, 0xDE, 0xDF): return FLAT
if enc_8bit in (0xE0, 0xE1, 0xE2, 0xE3): return MUBUF
if enc_8bit in (0xE8, 0xE9, 0xEA, 0xEB): return MTBUF
return None
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
old_stdout = sys.stdout
sys.stdout = io.StringIO()
compiler.disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
results = []
for line in output.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if not instr: continue
comment = line.split('//')[1].strip()
if ':' not in comment: continue
hex_str = comment.split(':')[1].strip().split()[0]
try:
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
results.append((instr, machine_bytes))
except ValueError:
continue
return results
def compile_asm(instr: str, compiler=None) -> bytes | None:
"""Compile a single instruction with llvm-mc and return the machine code bytes."""
import subprocess
try:
result = subprocess.run(
['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=f".text\n{instr}\n", capture_output=True, text=True)
if result.returncode != 0: return None
# Parse encoding: [0x01,0x39,0x0a,0x7e]
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
return bytes.fromhex(hex_vals)
except Exception:
pass
return None
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
def _test_kernel_roundtrip(self, op_fn):
"""Generate kernel from op_fn, test:
1. decode -> reencode matches original bytes
2. asm(disasm()) matches LLVM output
3. our disasm() matches LLVM's disassembly string exactly
"""
from extra.assembly.rdna3.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.compiler_amd import HIPCompiler
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
compiler = HIPCompiler('gfx1100')
decode_passed, decode_failed, decode_skipped = 0, 0, 0
asm_passed, asm_failed, asm_skipped = 0, 0, 0
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
decode_failures, asm_failures, disasm_failures = [], [], []
for ki, kernel in enumerate(kernels):
offset = 0
while offset < len(kernel.code):
remaining = kernel.code[offset:]
fmt = detect_format(remaining)
if fmt is None:
decode_skipped += 1
asm_skipped += 1
disasm_skipped += 1
offset += 4
continue
size = fmt._size()
if len(remaining) < size:
break
orig_bytes = remaining[:size]
# Test 1: decode -> reencode roundtrip
try:
decoded = fmt.from_bytes(orig_bytes)
reencoded = decoded.to_bytes()
if reencoded[:size] == orig_bytes:
decode_passed += 1
else:
decode_failed += 1
decode_failures.append(f"K{ki}@{offset}: {decoded.disasm()}: orig={orig_bytes.hex()} reenc={reencoded[:size].hex()}")
our_disasm = decoded.disasm()
# Test 2: asm(disasm()) matches LLVM output
try:
our_bytes = asm(our_disasm).to_bytes()
llvm_bytes = compile_asm(our_disasm, compiler)
if llvm_bytes is None:
asm_skipped += 1
elif our_bytes[:len(llvm_bytes)] == llvm_bytes:
asm_passed += 1
else:
asm_failed += 1
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
except Exception:
asm_skipped += 1
# Test 3: our disasm() matches LLVM's disassembly string exactly
# Skip if instruction uses op_XX (unknown opcode) or looks malformed (many raw field values)
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm):
disasm_skipped += 1
else:
try:
# Get LLVM's disassembly of our instruction
src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n {our_disasm}\n"
lib = compiler.compile(src)
llvm_instrs = disassemble_lib(lib, compiler)
if llvm_instrs:
llvm_disasm = llvm_instrs[0][0]
if our_disasm == llvm_disasm:
disasm_passed += 1
else:
disasm_failed += 1
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
else:
disasm_skipped += 1
except Exception:
disasm_skipped += 1
except Exception:
decode_skipped += 1
asm_skipped += 1
disasm_skipped += 1
offset += size
print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
self.assertEqual(disasm_failed, 0, f"Disasm failures:\n" + "\n".join(disasm_failures[:20]))
# Basic unary ops
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Reductions
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
# Matmul
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
# Complex ops
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
# Cast ops
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
# Comparison ops
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
if __name__ == "__main__":
unittest.main()
-4
View File
@@ -1,4 +0,0 @@
*.deb
build
src
sniffer/sniff.so
-20
View File
@@ -1,20 +0,0 @@
Built ROCT-Thunk-Interface (hsakmt)
hsakmt-roct-dev_5.4.4.99999-local_amd64.deb
note: installs to /opt/rocm
Built ROCm-Device-Libs
Works with ROCM_PATH=/home/tiny/build/ROCm-Device-Libs/build/dist
rocm-device-libs_1.0.0.99999-local_amd64.deb
Built ROCm-CompilerSupport (amd_comgr)
no deb, sudo make install to /usr/local
Built ROCR-Runtime
hsa-rocr_1.8.0-local_amd64.deb
hsa-rocr-dev_1.8.0-local_amd64.deb
Built ROCm-OpenCL-Runtime
rocm-ocl-icd_2.0.0-local_amd64.deb
ISSUE: these depend on "comgr"
rocm-opencl_2.0.0-local_amd64.deb
rocm-opencl-dev_2.0.0-local_amd64.deb
Did sudo make install
-41
View File
@@ -1,41 +0,0 @@
# run two "rocm-bandwidth-test" in a loop
# amdgpu-6.0.5-1581431.20.04
# fixed in kernel 6.2.14
[ 72.153646] RIP: 0010:pm_send_runlist+0x4a/0x630 [amdgpu]
[ 72.153815] Code: 30 65 48 8b 04 25 28 00 00 00 48 89 45 d0 31 c0 80 fb 01 0f 87 aa 9d 49 00 83 e3 01 0f 85 1c 05 00 00 49 8b 3f b8 01 00 00 00 <48> 8b 97 30 01 00 00 44 8b b7 6c 01 00 00 8b 9f 70 01 00 00 8b 8a
[ 72.153900] RSP: 0018:ffffb48445c03c30 EFLAGS: 00010246
[ 72.153928] RAX: 0000000000000001 RBX: 0000000000000000 RCX: 0000000000000000
[ 72.153962] RDX: 000000000000007b RSI: ffff9395e1562558 RDI: 0000000000000000
[ 72.153996] RBP: ffffb48445c03cb8 R08: 0000000000000000 R09: 0000000000000001
[ 72.154030] R10: ffff9395c900d840 R11: 0000000000000000 R12: 0000000000000000
[ 72.154065] R13: ffff9395c9e00400 R14: 0000000000000001 R15: ffff9395e15624e0
[ 72.154099] FS: 00007f345c6463c0(0000) GS:ffff93a4aee80000(0000) knlGS:0000000000000000
[ 72.154137] CS: 0010 DS: 0000 ES: 0000 CR0: 0000000080050033
[ 72.154165] CR2: 0000000000000130 CR3: 0000000112840000 CR4: 0000000000750ee0
[ 72.154201] PKRU: 55555554
[ 72.154215] Call Trace:
[ 72.154230] <TASK>
[ 72.154244] map_queues_cpsch+0x75/0xc0 [amdgpu]
[ 72.154365] debug_map_and_unlock+0x51/0x90 [amdgpu]
[ 72.154480] debug_refresh_runlist+0x1f/0x30 [amdgpu]
[ 72.154591] kfd_dbg_runtime_disable+0x13c/0x240 [amdgpu]
[ 72.154705] kfd_ioctl_dbg_set_debug_trap+0x69d/0x8b0 [amdgpu]
[ 72.154820] kfd_ioctl+0x24a/0x5b0 [amdgpu]
[ 72.154925] ? kfd_ioctl_create_queue+0x770/0x770 [amdgpu]
[ 72.155035] ? syscall_exit_to_user_mode+0x27/0x50
[ 72.155061] ? exit_to_user_mode_prepare+0x3d/0x1c0
[ 72.155088] __x64_sys_ioctl+0x95/0xd0
[ 72.155109] do_syscall_64+0x5c/0xc0
[ 72.155128] ? syscall_exit_to_user_mode+0x27/0x50
[ 72.155151] ? do_syscall_64+0x69/0xc0
[ 72.155172] entry_SYSCALL_64_after_hwframe+0x61/0xcb
[ 72.155198] RIP: 0033:0x7f345c7f63ab
[ 72.155218] Code: 0f 1e fa 48 8b 05 e5 7a 0d 00 64 c7 00 26 00 00 00 48 c7 c0 ff ff ff ff c3 66 0f 1f 44 00 00 f3 0f 1e fa b8 10 00 00 00 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 8b 0d b5 7a 0d 00 f7 d8 64 89 01 48
[ 72.155301] RSP: 002b:00007ffc97cc89f8 EFLAGS: 00000246 ORIG_RAX: 0000000000000010
[ 72.155339] RAX: ffffffffffffffda RBX: 00007ffc97cc8a30 RCX: 00007f345c7f63ab
[ 72.155375] RDX: 00007ffc97cc8a30 RSI: 00000000c0284b82 RDI: 0000000000000003
[ 72.155411] RBP: 00000000c0284b82 R08: 0000000000000000 R09: 0000000000000000
[ 72.155447] R10: 00007f345cd4ddb0 R11: 0000000000000246 R12: 00007ffc97cc8a30
[ 72.155481] R13: 0000000000000003 R14: 00007ffc97cc8d20 R15: 0000000000000000
[ 72.155517] </TASK>
-41
View File
@@ -1,41 +0,0 @@
# run two tinygrad matrix example in a loop
# amdgpu-6.0.5-1581431.20.04
# NOT fixed in kernel 6.2.14
[ 553.016624] gmc_v11_0_process_interrupt: 30 callbacks suppressed
[ 553.016631] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:9 pasid:32770, for process python3 pid 10001 thread python3 pid 10001)
[ 553.016790] amdgpu 0000:0b:00.0: amdgpu: in page starting at address 0x00007f0000000000 from client 10
[ 553.016892] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00901A30
[ 553.016974] amdgpu 0000:0b:00.0: amdgpu: Faulty UTCL2 client ID: SDMA0 (0xd)
[ 553.017051] amdgpu 0000:0b:00.0: amdgpu: MORE_FAULTS: 0x0
[ 553.017111] amdgpu 0000:0b:00.0: amdgpu: WALKER_ERROR: 0x0
[ 553.017173] amdgpu 0000:0b:00.0: amdgpu: PERMISSION_FAULTS: 0x3
[ 553.017238] amdgpu 0000:0b:00.0: amdgpu: MAPPING_ERROR: 0x0
[ 553.017300] amdgpu 0000:0b:00.0: amdgpu: RW: 0x0
[ 553.123921] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
[ 553.124153] amdgpu: failed to add hardware queue to MES, doorbell=0x1a16
[ 553.124195] amdgpu: MES might be in unrecoverable state, issue a GPU reset
[ 553.124237] amdgpu: Failed to restore queue 2
[ 553.124266] amdgpu: Failed to restore process queues
[ 553.124270] amdgpu: Failed to evict queue 3
[ 553.124297] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
# alternative crash in kernel 6.2.14
[ 151.097948] gmc_v11_0_process_interrupt: 30 callbacks suppressed
[ 151.097953] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:8 pasid:32771, for process python3 pid 7525 thread python3 pid 7525)
[ 151.097993] amdgpu 0000:0b:00.0: amdgpu: in page starting at address 0x00007f0000000000 from client 10
[ 151.098008] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00801A30
[ 151.098020] amdgpu 0000:0b:00.0: amdgpu: Faulty UTCL2 client ID: SDMA0 (0xd)
[ 151.098032] amdgpu 0000:0b:00.0: amdgpu: MORE_FAULTS: 0x0
[ 151.098042] amdgpu 0000:0b:00.0: amdgpu: WALKER_ERROR: 0x0
[ 151.098052] amdgpu 0000:0b:00.0: amdgpu: PERMISSION_FAULTS: 0x3
[ 151.098062] amdgpu 0000:0b:00.0: amdgpu: MAPPING_ERROR: 0x0
[ 151.098071] amdgpu 0000:0b:00.0: amdgpu: RW: 0x0
[ 151.209517] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
[ 151.209724] amdgpu: failed to add hardware queue to MES, doorbell=0x1002
[ 151.209734] amdgpu: MES might be in unrecoverable state, issue a GPU reset
[ 151.209743] amdgpu: Failed to restore queue 1
[ 151.209751] amdgpu: Failed to restore process queues
[ 151.209759] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
[ 151.209858] amdgpu 0000:0b:00.0: amdgpu: GPU reset begin!
-20
View File
@@ -1,20 +0,0 @@
# two tinygrad + two bandwidth test
# RDNA2, driver 6.0.5
# recovered from this!
[ 136.971209] gmc_v10_0_process_interrupt: 39 callbacks suppressed
[ 136.971218] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:11 pasid:32773, for process rocm-bandwidth- pid 20281 thread rocm-bandwidth- pid 20281)
[ 136.971228] amdgpu 0000:0b:00.0: amdgpu: in page starting at address 0x00007f5c2b800000 from client 0x1b (UTCL2)
[ 136.971232] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00B01A31
[ 136.971233] amdgpu 0000:0b:00.0: amdgpu: Faulty UTCL2 client ID: SDMA0 (0xd)
[ 136.971235] amdgpu 0000:0b:00.0: amdgpu: MORE_FAULTS: 0x1
[ 136.971236] amdgpu 0000:0b:00.0: amdgpu: WALKER_ERROR: 0x0
[ 136.971236] amdgpu 0000:0b:00.0: amdgpu: PERMISSION_FAULTS: 0x3
[ 136.971237] amdgpu 0000:0b:00.0: amdgpu: MAPPING_ERROR: 0x0
[ 136.971238] amdgpu 0000:0b:00.0: amdgpu: RW: 0x0
...
[ 136.993979] amdgpu 0000:0b:00.0: amdgpu: IH ring buffer overflow (0x000BE5A0, 0x0003C480, 0x0003E5C0)
[ 138.209072] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004000 flags=0x0000]
[ 138.209078] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004d80 flags=0x0000]
[ 138.209081] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005000 flags=0x0000]
[ 138.209084] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005d80 flags=0x0000]
-33
View File
@@ -1,33 +0,0 @@
# ROCK-Kernel-Driver 0b579de9622f5c93021dcb7927d13926313740a2
# non fatal "crash"
[ 127.418045] ------------[ cut here ]------------
[ 127.418046] User pages unexpectedly invalid
[ 127.418056] WARNING: CPU: 16 PID: 260 at drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gpuvm.c:3000 amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
[ 127.418235] Modules linked in: rfcomm cmac algif_hash algif_skcipher af_alg bnep nls_iso8859_1 iwlmvm mac80211 intel_rapl_msr intel_rapl_common edac_mce_amd snd_hda_codec_realtek snd_hda_codec_generic snd_hda_codec_hdmi kvm_amd binfmt_misc snd_hda_intel snd_intel_dspcfg kvm libarc4 snd_intel_sdw_acpi snd_hda_codec btusb iwlwifi btrtl snd_hda_core btbcm btintel irqbypass btmtk snd_hwdep crct10dif_pclmul snd_pcm polyval_clmulni bluetooth snd_seq_midi snd_seq_midi_event snd_rawmidi snd_seq polyval_generic cfg80211 ghash_clmulni_intel eeepc_wmi snd_seq_device snd_timer aesni_intel asus_wmi ecdh_generic snd platform_profile crypto_simd ledtrig_audio cryptd ecc ccp soundcore sparse_keymap rapl k10temp wmi_bmof mac_hid sch_fq_codel msr parport_pc ppdev lp parport ramoops pstore_blk efi_pstore reed_solomon pstore_zone ip_tables x_tables autofs4 amdgpu hid_generic usbhid hid i2c_algo_bit drm_ttm_helper ttm video iommu_v2 drm_buddy gpu_sched drm_display_helper drm_kms_helper syscopyarea
[ 127.418276] sysfillrect sysimgblt fb_sys_fops drm nvme nvme_core cec r8169 ahci crc32_pclmul rc_core i2c_piix4 xhci_pci libahci nvme_common xhci_pci_renesas realtek wmi
[ 127.418284] CPU: 16 PID: 260 Comm: kworker/16:1 Tainted: G W 6.0.0 #4
[ 127.418286] Hardware name: System manufacturer System Product Name/TUF GAMING X570-PLUS (WI-FI), BIOS 3603 03/20/2021
[ 127.418287] Workqueue: events amdgpu_amdkfd_restore_userptr_worker [amdgpu]
[ 127.418455] RIP: 0010:amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
[ 127.418601] Code: ff e8 2b 8a 96 d1 e9 66 fe ff ff 48 c7 c7 40 4f f5 c0 e8 56 7b 8a d1 0f 0b e9 2e ff ff ff 48 c7 c7 d8 d0 ed c0 e8 43 7b 8a d1 <0f> 0b e9 0a fe ff ff 4c 89 ef e8 f8 89 96 d1 e9 cb fd ff ff e8 ce
[ 127.418603] RSP: 0018:ffffb36740a83dc8 EFLAGS: 00010282
[ 127.418604] RAX: 0000000000000000 RBX: ffff9d159ee9df30 RCX: 0000000000000027
[ 127.418605] RDX: 0000000000000027 RSI: ffffb36740a83c88 RDI: ffff9d242a220568
[ 127.418606] RBP: ffffb36740a83e58 R08: ffff9d242a220560 R09: 0000000000000001
[ 127.418607] R10: 0000000000000001 R11: 0000000000000020 R12: ffff9d159ee9df98
[ 127.418607] R13: ffff9d159ee9df70 R14: ffff9d159ee9dee0 R15: ffff9d159ee9dee0
[ 127.418608] FS: 0000000000000000(0000) GS:ffff9d242a200000(0000) knlGS:0000000000000000
[ 127.418609] CS: 0010 DS: 0000 ES: 0000 CR0: 0000000080050033
[ 127.418610] CR2: 00007fd5d4715000 CR3: 0000000120ffe000 CR4: 0000000000750ee0
[ 127.418611] PKRU: 55555554
[ 127.418611] Call Trace:
[ 127.418612] <TASK>
[ 127.418613] process_one_work+0x21f/0x3f0
[ 127.418615] worker_thread+0x4a/0x3c0
[ 127.418617] ? process_one_work+0x3f0/0x3f0
[ 127.418618] kthread+0xf0/0x120
[ 127.418619] ? kthread_complete_and_exit+0x20/0x20
[ 127.418620] ret_from_fork+0x22/0x30
[ 127.418622] </TASK>
[ 127.418623] ---[ end trace 0000000000000000 ]---
-80
View File
@@ -1,80 +0,0 @@
import numpy as np
import pathlib
from hexdump import hexdump
from tinygrad.helpers import colored
from extra.helpers import enable_early_exec
early_exec = enable_early_exec()
from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
ENABLE_NON_ASM = False
WMMA = True
DUAL_ALU = True
F32 = True
if ENABLE_NON_ASM:
buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
prg_empty = CLProgram("code", "__kernel void code(__global float *a) { a[0] = 1; }")
asm_real = prg_empty.binary()
with open("/tmp/cc.elf", "wb") as f:
f.write(asm_real)
prg_empty([1], [1], buf, wait=True)
print(buf.toCPU())
print(colored("creating CLBuffer", "green"))
buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
code = open(pathlib.Path(__file__).parent / "prog.s", "r").read()
gen = []
FLOPS = 0
MAX_REG = 251
for j in range(1):
if WMMA:
KY, KX = 4, 4
for y in range(KY):
for x in range(KX):
c = (y*KX+x)*8
a = (KY*KX*8) + y*8
b = (KY*KX*8) + (KY*8) + x*8
gen.append(f"v_wmma_f32_16x16x16_f16 v[{c}:{c+7}], v[{a}:{a+7}], v[{b}:{b+7}], v[{c}:{c+7}]")
FLOPS += 16*8*2
else:
for i in range(0, MAX_REG, 6):
if DUAL_ALU:
if F32:
gen.append(f"v_dual_fmac_f32 v{i+0}, v{i+1}, v{i+2} :: v_dual_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
FLOPS += 4
else:
gen.append(f"v_dual_dot2acc_f32_f16 v{i+0}, v{i+1}, v{i+2} :: v_dual_dot2acc_f32_f16 v{i+3}, v{i+4}, v{i+5}")
FLOPS += 8
else:
assert F32
gen.append(f"v_fmac_f32 v{i+0}, v{i+1}, v{i+2}")
gen.append(f"v_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
code = code.replace("// FLOPS", '\n'.join(gen))
print(code)
# fix: COMGR failed to get code object ISA name. set triple to 'amdgcn-amd-amdhsa'
object = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], object))
with open("/tmp/cc2.o", "wb") as f:
f.write(object)
with open("/tmp/cc2.elf", "wb") as f:
f.write(asm)
print(colored("creating CLProgram", "green"))
prg = CLProgram("code", asm)
print(colored("running program", "green"))
G = 512
FLOPS *= 100000*G*G # loop * global_size
for i in range(3):
tm = prg(buf, global_size=[G//256, G, 1], local_size=[256, 1, 1], wait=True)
print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
print(colored("transferring buffer", "green"))
print(buf.toCPU())
-80
View File
@@ -1,80 +0,0 @@
.global _start
_start:
.rodata
.align 0x10
.global code.kd
.type code.kd,STT_OBJECT
# amd_kernel_code_t (must be at 0x440 for kernel_code_entry_byte_offset to be right)
code.kd:
# amd_kernel_..., amd_machine_...
.long 0,0,0,0
# kernel_code_entry_byte_offset, kernel_code_prefetch_byte_offset
.long 0x00000bc0,0x00000000,0x00000000,0x00000000
# kernel_code_prefetch_byte_size, max_scratch_backing_memory_byte_size
.long 0,0,0,0
# compute_pgm_rsrc1, compute_pgm_rsrc2, kernel_code_properties, workitem_private_segment_byte_size
.long 0x60af0000,0x0000009e,0x00000408,0x00000000
# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_32 | AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_16_64
# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_DX10_CLAMP | AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_IEEE_MODE
# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT = 0xF
# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X
# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_KERNARG_SEGMENT_PTR = 1
# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = 1
.text
.global code
.type code,STT_FUNC
code:
# https://llvm.org/docs/AMDGPUUsage.html#initial-kernel-execution-state
# s[0:1] contains the kernarg_address
# TODO: can we use s[2:3] if this was really a wave since we only alloced 2 SGPRs?
s_load_b64 s[2:3], s[0:1], null
s_mov_b32 s8, 0
loop:
s_addk_i32 s8, 1
s_cmp_eq_u32 s8, 100000
// FLOPS
s_cbranch_scc0 loop
# wait for the s_load_b64
s_waitcnt lgkmcnt(0)
v_dual_mov_b32 v0, 4 :: v_dual_mov_b32 v1, 2.0
global_store_b32 v0, v1, s[2:3]
# Deallocate all VGPRs for this wave. Use only when next instruction is S_ENDPGM.
s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
s_endpgm
s_code_end
.amdgpu_metadata
amdhsa.kernels:
- .args:
- .address_space: global
.name: a
.offset: 0
.size: 8
.type_name: 'float*'
.value_kind: global_buffer
.group_segment_fixed_size: 0
.kernarg_segment_align: 8
.kernarg_segment_size: 8
.language: OpenCL C
.language_version:
- 1
- 2
.max_flat_workgroup_size: 256
.name: code
.private_segment_fixed_size: 0
.sgpr_count: 2
.sgpr_spill_count: 0
.symbol: code.kd
.uses_dynamic_stack: false
.vgpr_count: 256
.vgpr_spill_count: 0
.wavefront_size: 32
amdhsa.target: amdgcn-amd-amdhsa--gfx1100
amdhsa.version:
- 1
- 2
.end_amdgpu_metadata
-11
View File
@@ -1,11 +0,0 @@
#!/bin/bash
mkdir -p src
cd src
git clone https://github.com/RadeonOpenCompute/ROCT-Thunk-Interface.git -b rocm-5.5.0
git clone https://github.com/RadeonOpenCompute/ROCm-Device-Libs.git -b rocm-5.5.0
git clone https://github.com/RadeonOpenCompute/llvm-project.git -b rocm-5.5.0 --depth 1
git clone https://github.com/RadeonOpenCompute/ROCR-Runtime.git -b rocm-5.5.0
git clone https://github.com/ROCm-Developer-Tools/ROCclr.git -b rocm-5.5.0
git clone https://github.com/RadeonOpenCompute/ROCm-CompilerSupport.git -b rocm-5.5.0
git clone https://github.com/RadeonOpenCompute/ROCm-OpenCL-Runtime.git -b rocm-5.5.0
cd ../
-69
View File
@@ -1,69 +0,0 @@
#!/bin/bash
mkdir -p build/debs
cd build
# ROCT-Thunk-Interface (hsakmt)
if [ ! -f debs/hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ]
then
mkdir -p ROCT-Thunk-Interface
cd ROCT-Thunk-Interface
cmake ../../src/ROCT-Thunk-Interface
make -j32 package
cp hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ../debs
cd ../
fi
# build custom LLVM
if [ ! -f llvm-project/bin/clang ]
then
mkdir -p llvm-project
cd llvm-project
cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS="llvm;clang;lld" -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" ../../src/llvm-project/llvm
make -j32
cd ..
fi
# use custom LLVM
export PATH="$PWD/llvm-project/bin:$PATH"
# ROCm-Device-Libs
if [ ! -f debs/rocm-device-libs_1.0.0.99999-local_amd64.deb ]
then
mkdir -p ROCm-Device-Libs
cd ROCm-Device-Libs
cmake ../../src/ROCm-Device-Libs
make -j32 package
cp rocm-device-libs_1.0.0.99999-local_amd64.deb ../debs
cd ../
fi
# ROCR-Runtime
if [ ! -f debs/hsa-rocr_1.8.0-local_amd64.deb ]
then
mkdir -p ROCR-Runtime
cd ROCR-Runtime
cmake ../../src/ROCR-Runtime/src
make -j32 package
cp hsa-rocr_1.8.0-local_amd64.deb ../debs
cp hsa-rocr-dev_1.8.0-local_amd64.deb ../debs
cd ../
fi
# ROCm-OpenCL-Runtime (needs ROCclr)
if [ ! -f debs/rocm-opencl_2.0.0-local_amd64.deb ]
then
mkdir -p ROCm-OpenCL-Runtime
cd ROCm-OpenCL-Runtime
cmake ../../src/ROCm-OpenCL-Runtime
make -j32 package
cp rocm-opencl_2.0.0-local_amd64.deb ../debs
cp rocm-opencl-dev_2.0.0-local_amd64.deb ../debs
cp rocm-ocl-icd_2.0.0-local_amd64.deb ../debs
fi
# ROCm-CompilerSupport (broken)
#mkdir -p ROCm-CompilerSupport
#cd ROCm-CompilerSupport
#cmake ../../src/ROCm-CompilerSupport/lib/comgr
#make -j32
-14
View File
@@ -1,14 +0,0 @@
#!/bin/bash
rm amdgpu-install_5.5.50500-1_all.deb
wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/$(lsb_release -cs)/amdgpu-install_5.5.50500-1_all.deb
sudo dpkg -i amdgpu-install_5.5.50500-1_all.deb
sudo apt-get update
# kernel driver
sudo apt-get install amdgpu-dkms
# for opencl
sudo apt-get install rocm-opencl-runtime
# for HIP
sudo apt-get install hip-runtime-amd rocm-device-libs hip-dev
-11
View File
@@ -1,11 +0,0 @@
#!/bin/bash -e
clang sniff.cc -Werror -shared -fPIC -I../src/ -I../src/ROCT-Thunk-Interface/include -I../src/ROCm-Device-Libs/ockl/inc -o sniff.so -lstdc++
#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so /home/tiny/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so $HOME/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
#AMD_LOG_LEVEL=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
DEBUG=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
#AMD_LOG_LEVEL=5 HSAKMT_DEBUG_LEVEL=7 DEBUG=5 LD_PRELOAD=$PWD/sniff.so strace -F python3 ../rdna3/asm.py
#LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so FORWARD_ONLY=1 DEBUG=2 python3 ../../../test/test_ops.py TestOps.test_add
#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 0 -d 1 -m 1
#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 1 -d 2 -m 1
-282
View File
@@ -1,282 +0,0 @@
// template copied from https://github.com/geohot/cuda_ioctl_sniffer/blob/master/sniff.cc
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dlfcn.h>
#include <signal.h>
#include <ucontext.h>
#include <sys/mman.h>
// includes from the ROCm sources
#include <linux/kfd_ioctl.h>
#include <hsa.h>
#include <amd_hsa_kernel_code.h>
#include <ROCR-Runtime/src/core/inc/sdma_registers.h>
using namespace rocr::AMD;
#include <string>
#include <map>
std::map<int, std::string> files;
std::map<uint64_t, uint64_t> ring_base_addresses;
#define D(args...) fprintf(stderr, args)
uint64_t doorbell_offset = -1;
std::map<uint64_t, int> queue_types;
void hexdump(void *d, int l) {
for (int i = 0; i < l; i++) {
if (i%0x10 == 0 && i != 0) printf("\n");
if (i%0x10 == 8) printf(" ");
if (i%0x10 == 0) printf("%8X: ", i);
printf("%2.2X ", ((uint8_t*)d)[i]);
}
printf("\n");
}
extern "C" {
// https://defuse.ca/online-x86-assembler.htm#disassembly2
static void handler(int sig, siginfo_t *si, void *unused) {
ucontext_t *u = (ucontext_t *)unused;
uint8_t *rip = (uint8_t*)u->uc_mcontext.gregs[REG_RIP];
int store_size = 0;
uint64_t value;
if (rip[0] == 0x48 && rip[1] == 0x89 && rip[2] == 0x30) {
// 0: 48 89 30 mov QWORD PTR [rax],rsi
store_size = 8;
value = u->uc_mcontext.gregs[REG_RSI];
u->uc_mcontext.gregs[REG_RIP] += 3;
} else if (rip[0] == 0x4c && rip[1] == 0x89 && rip[2] == 0x28) {
// 0: 4c 89 28 mov QWORD PTR [rax],r13
store_size = 8;
value = u->uc_mcontext.gregs[REG_R13];
u->uc_mcontext.gregs[REG_RIP] += 3;
} else {
D("segfault %02X %02X %02X %02X %02X %02X %02X %02X rip: %p addr: %p\n", rip[0], rip[1], rip[2], rip[3], rip[4], rip[5], rip[6], rip[7], rip, si->si_addr);
D("rax: %llx rcx: %llx rdx: %llx rsi: %llx rbx: %llx\n", u->uc_mcontext.gregs[REG_RAX], u->uc_mcontext.gregs[REG_RCX], u->uc_mcontext.gregs[REG_RDX], u->uc_mcontext.gregs[REG_RSI], u->uc_mcontext.gregs[REG_RBX]);
exit(-1);
}
uint64_t ring_base_address = ring_base_addresses[((uint64_t)si->si_addr)&0xFFF];
int queue_type = queue_types[((uint64_t)si->si_addr)&0xFFF];
D("%16p: \u001b[31mDING DONG\u001b[0m (queue_type %d) store(%d): 0x%8lx -> %p ring_base_address:0x%lx\n", rip, queue_type, store_size, value, si->si_addr, ring_base_address);
if (queue_type == KFD_IOC_QUEUE_TYPE_SDMA) {
uint8_t *sdma_ptr = (uint8_t*)(ring_base_address);
while (sdma_ptr < ((uint8_t*)(ring_base_address)+value)) {
D("0x%3lx: ", sdma_ptr-(uint8_t*)(ring_base_address));
if (sdma_ptr[0] == SDMA_OP_TIMESTAMP) {
D("SDMA_PKT_TIMESTAMP\n");
sdma_ptr += sizeof(SDMA_PKT_TIMESTAMP);
} else if (sdma_ptr[0] == SDMA_OP_GCR) {
D("SDMA_PKT_GCR\n");
sdma_ptr += sizeof(SDMA_PKT_GCR);
} else if (sdma_ptr[0] == SDMA_OP_ATOMIC) {
D("SDMA_PKT_ATOMIC\n");
sdma_ptr += sizeof(SDMA_PKT_ATOMIC);
} else if (sdma_ptr[0] == SDMA_OP_FENCE) {
D("SDMA_PKT_FENCE\n");
sdma_ptr += sizeof(SDMA_PKT_FENCE);
} else if (sdma_ptr[0] == SDMA_OP_TRAP) {
D("SDMA_PKT_TRAP\n");
sdma_ptr += sizeof(SDMA_PKT_TRAP);
} else if (sdma_ptr[0] == SDMA_OP_COPY && sdma_ptr[1] == SDMA_SUBOP_COPY_LINEAR) {
SDMA_PKT_COPY_LINEAR *pkt = (SDMA_PKT_COPY_LINEAR *)sdma_ptr;
D("SDMA_PKT_COPY_LINEAR: count:0x%x src:0x%lx dst:0x%lx\n", pkt->COUNT_UNION.count+1,
(uint64_t)pkt->SRC_ADDR_LO_UNION.src_addr_31_0 | ((uint64_t)pkt->SRC_ADDR_HI_UNION.src_addr_63_32 << 32),
(uint64_t)pkt->DST_ADDR_LO_UNION.dst_addr_31_0 | ((uint64_t)pkt->DST_ADDR_HI_UNION.dst_addr_63_32 << 32)
);
sdma_ptr += sizeof(SDMA_PKT_COPY_LINEAR);
} else {
D("unhandled packet type %d %d, exiting\n", sdma_ptr[0], sdma_ptr[1]);
break;
}
}
//hexdump((void*)(ring_base_address), 0x100);
} else if (queue_type == KFD_IOC_QUEUE_TYPE_COMPUTE_AQL) {
hsa_kernel_dispatch_packet_t *pkt = (hsa_kernel_dispatch_packet_t *)(ring_base_address+value*0x40);
if ((pkt->header&0xFF) == HSA_PACKET_TYPE_KERNEL_DISPATCH) {
D("HSA_PACKET_TYPE_KERNEL_DISPATCH -- setup:%d workgroup[%d, %d, %d] grid[%d, %d, %d] kernel_object:0x%lx kernarg_address:%p\n", pkt->setup, pkt->workgroup_size_x, pkt->workgroup_size_y, pkt->workgroup_size_z, pkt->grid_size_x, pkt->grid_size_y, pkt->grid_size_z, pkt->kernel_object, pkt->kernarg_address);
amd_kernel_code_t *code = (amd_kernel_code_t *)pkt->kernel_object;
D("kernel_code_entry_byte_offset:%lx\n", code->kernel_code_entry_byte_offset);
uint32_t *kernel_code = (uint32_t*)(pkt->kernel_object + code->kernel_code_entry_byte_offset);
int code_len = 0;
while (kernel_code[code_len] != 0xbf9f0000 && kernel_code[code_len] != 0) code_len++;
hexdump(kernel_code, code_len*4);
/*FILE *f = fopen("/tmp/kernel_code", "wb");
fwrite(kernel_code, 4, code_len, f);
fclose(f);
system("python -c 'print(\" \".join([(\"0x%02X\"%x) for x in open(\"/tmp/kernel_code\", \"rb\").read()]))' | ../build/llvm-project/bin/llvm-mc --disassemble --arch=amdgcn --mcpu=gfx1100 --show-encoding");*/
D("kernargs (kernarg_segment_byte_size:0x%lx)\n", code->kernarg_segment_byte_size);
// get length
int i;
for (i = 0; i < 0x400; i+=0x10) {
if (memcmp((void*)((uint64_t)pkt->kernarg_address+i), "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", 0x10) == 0) break;
}
hexdump((void*)pkt->kernarg_address, i+0x10);
} else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_BARRIER_AND) {
hsa_barrier_and_packet_t *pkt_and = (hsa_barrier_and_packet_t *)(ring_base_address+value*0x40);
D("HSA_PACKET_TYPE_BARRIER_AND completion_signal:0x%lx\n", pkt_and->completion_signal.handle);
//hexdump((void*)(ring_base_address+value*0x40), 0x40);
} else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_VENDOR_SPECIFIC) {
D("HSA_PACKET_TYPE_VENDOR_SPECIFIC\n");
hexdump((void*)(ring_base_address+value*0x40), 0x40);
} else {
hexdump((void*)(ring_base_address+value*0x40), 0x40);
}
}
mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_READ | PROT_WRITE);
if (store_size == 8) {
*(volatile uint64_t*)(si->si_addr) = value;
} else if (store_size == 4) {
*(volatile uint32_t*)(si->si_addr) = value;
} else if (store_size == 2) {
*(volatile uint16_t*)(si->si_addr) = value;
} else {
D("store size not supported\n");
exit(-1);
}
mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_NONE);
}
void register_sigsegv_handler() {
struct sigaction sa = {0};
sa.sa_flags = SA_SIGINFO;
sigemptyset(&sa.sa_mask);
sa.sa_sigaction = handler;
if (sigaction(SIGSEGV, &sa, NULL) == -1) {
D("ERROR: failed to register sigsegv handler");
exit(-1);
}
// NOTE: python (or ocl runtime?) blocks the SIGSEGV signal
sigset_t x;
sigemptyset(&x);
sigaddset(&x, SIGSEGV);
sigprocmask(SIG_UNBLOCK, &x, NULL);
}
int (*my_open)(const char *pathname, int flags, mode_t mode);
#undef open
int open(const char *pathname, int flags, mode_t mode) {
if (my_open == NULL) my_open = reinterpret_cast<decltype(my_open)>(dlsym(RTLD_NEXT, "open"));
int ret = my_open(pathname, flags, mode);
//D("open %s (0o%o) = %d\n", pathname, flags, ret);
files[ret] = pathname;
return ret;
}
int (*my_open64)(const char *pathname, int flags, mode_t mode);
#undef open
int open64(const char *pathname, int flags, mode_t mode) {
if (my_open64 == NULL) my_open64 = reinterpret_cast<decltype(my_open64)>(dlsym(RTLD_NEXT, "open64"));
int ret = my_open64(pathname, flags, mode);
//D("open %s (0o%o) = %d\n", pathname, flags, ret);
files[ret] = pathname;
return ret;
}
void *(*my_mmap)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
#undef mmap
void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
if (my_mmap == NULL) my_mmap = reinterpret_cast<decltype(my_mmap)>(dlsym(RTLD_NEXT, "mmap"));
void *ret = my_mmap(addr, length, prot, flags, fd, offset);
if (doorbell_offset != -1 && offset == doorbell_offset) {
D("HIDDEN DOORBELL %p, handled by %p\n", addr, handler);
register_sigsegv_handler();
mprotect(addr, length, PROT_NONE);
}
if (fd != -1) D("mmapped %p (target %p) with flags 0x%x length 0x%zx fd %d %s offset 0x%lx\n", ret, addr, flags, length, fd, files[fd].c_str(), offset);
return ret;
}
void *(*my_mmap64)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
#undef mmap64
void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { return mmap(addr, length, prot, flags, fd, offset); }
int ioctl_num = 1;
int (*my_ioctl)(int filedes, unsigned long request, void *argp) = NULL;
#undef ioctl
int ioctl(int filedes, unsigned long request, void *argp) {
if (my_ioctl == NULL) my_ioctl = reinterpret_cast<decltype(my_ioctl)>(dlsym(RTLD_NEXT, "ioctl"));
int ret = 0;
ret = my_ioctl(filedes, request, argp);
if (!files.count(filedes)) return ret;
uint8_t type = (request >> 8) & 0xFF;
uint8_t nr = (request >> 0) & 0xFF;
uint16_t size = (request >> 16) & 0xFFF;
D("%3d: %d = %3d(%20s) 0x%3x ", ioctl_num, ret, filedes, files[filedes].c_str(), size);
if (request == AMDKFD_IOC_SET_EVENT) {
kfd_ioctl_set_event_args *args = (kfd_ioctl_set_event_args *)argp;
D("AMDKFD_IOC_SET_EVENT event_id:%d", args->event_id);
} else if (request == AMDKFD_IOC_ALLOC_MEMORY_OF_GPU) {
kfd_ioctl_alloc_memory_of_gpu_args *args = (kfd_ioctl_alloc_memory_of_gpu_args *)argp;
D("AMDKFD_IOC_ALLOC_MEMORY_OF_GPU va_addr:0x%llx size:0x%llx handle:%llX gpu_id:0x%x", args->va_addr, args->size, args->handle, args->gpu_id);
} else if (request == AMDKFD_IOC_MAP_MEMORY_TO_GPU) {
kfd_ioctl_map_memory_to_gpu_args *args = (kfd_ioctl_map_memory_to_gpu_args *)argp;
D("AMDKFD_IOC_MAP_MEMORY_TO_GPU handle:%llX", args->handle);
} else if (request == AMDKFD_IOC_CREATE_EVENT) {
kfd_ioctl_create_event_args *args = (kfd_ioctl_create_event_args *)argp;
D("AMDKFD_IOC_CREATE_EVENT event_page_offset:0x%llx event_type:%d event_id:%d", args->event_page_offset, args->event_type, args->event_id);
} else if (request == AMDKFD_IOC_WAIT_EVENTS) {
D("AMDKFD_IOC_WAIT_EVENTS");
} else if (request == AMDKFD_IOC_SET_XNACK_MODE) {
D("AMDKFD_IOC_SET_XNACK_MODE");
} else if (request == AMDKFD_IOC_SVM || (type == 0x4b && nr == 0x20)) {
// NOTE: this one is variable length
kfd_ioctl_svm_args *args = (kfd_ioctl_svm_args *)argp;
D("AMDKFD_IOC_SVM start_addr:0x%llx size:0x%llx op:%d", args->start_addr, args->size, args->op);
} else if (request == AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU) {
kfd_ioctl_unmap_memory_from_gpu_args *args = (kfd_ioctl_unmap_memory_from_gpu_args *)argp;
D("AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU handle:%llX", args->handle);
} else if (request == AMDKFD_IOC_FREE_MEMORY_OF_GPU) {
D("AMDKFD_IOC_FREE_MEMORY_OF_GPU");
} else if (request == AMDKFD_IOC_SET_SCRATCH_BACKING_VA) {
D("AMDKFD_IOC_SET_SCRATCH_BACKING_VA");
} else if (request == AMDKFD_IOC_GET_TILE_CONFIG) {
D("AMDKFD_IOC_GET_TILE_CONFIG");
} else if (request == AMDKFD_IOC_SET_TRAP_HANDLER) {
D("AMDKFD_IOC_SET_TRAP_HANDLER");
} else if (request == AMDKFD_IOC_GET_VERSION) {
kfd_ioctl_get_version_args *args = (kfd_ioctl_get_version_args *)argp;
D("AMDKFD_IOC_GET_VERSION major_version:%d minor_version:%d", args->major_version, args->minor_version);
} else if (request == AMDKFD_IOC_GET_PROCESS_APERTURES_NEW) {
D("AMDKFD_IOC_GET_PROCESS_APERTURES_NEW");
} else if (request == AMDKFD_IOC_ACQUIRE_VM) {
D("AMDKFD_IOC_ACQUIRE_VM");
} else if (request == AMDKFD_IOC_SET_MEMORY_POLICY) {
D("AMDKFD_IOC_SET_MEMORY_POLICY");
} else if (request == AMDKFD_IOC_GET_CLOCK_COUNTERS) {
D("AMDKFD_IOC_GET_CLOCK_COUNTERS");
} else if (request == AMDKFD_IOC_CREATE_QUEUE) {
kfd_ioctl_create_queue_args *args = (kfd_ioctl_create_queue_args *)argp;
D("AMDKFD_IOC_CREATE_QUEUE\n");
D("queue_type:%d ring_base_address:0x%llx\n", args->queue_type, args->ring_base_address);
D("eop_buffer_address:0x%llx ctx_save_restore_address:0x%llx\n", args->eop_buffer_address, args->ctx_save_restore_address);
D("ring_size:0x%x queue_priority:%d\n", args->ring_size, args->queue_priority);
D("RETURNS write_pointer_address:0x%llx read_pointer_address:0x%llx doorbell_offset:0x%llx queue_id:%d\n", args->write_pointer_address, args->read_pointer_address, args->doorbell_offset, args->queue_id);
//D("RETURNS *write_pointer_address:0x%llx *read_pointer_address:0x%llx\n", *(uint64_t*)args->write_pointer_address, *(uint64_t*)args->read_pointer_address);
ring_base_addresses[args->doorbell_offset&0xFFF] = args->ring_base_address;
queue_types[args->doorbell_offset&0xFFF] = args->queue_type;
doorbell_offset = args->doorbell_offset&~0xFFF;
} else {
D("type:0x%x nr:0x%x size:0x%x", type, nr, size);
}
D("\n");
ioctl_num++;
return ret;
}
}
+38 -13755
View File
File diff suppressed because it is too large Load Diff
+83
View File
@@ -0,0 +1,83 @@
.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 32
# register usage (RSRC1)
.amdhsa_next_free_vgpr 504
.amdhsa_next_free_sgpr 96
# workgroup / workitem IDs (RSRC2)
.amdhsa_system_sgpr_workgroup_id_x 1
.amdhsa_system_sgpr_workgroup_id_y 1
.amdhsa_system_sgpr_workgroup_id_z 1
# user 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
- .name: num_wg
.offset: 28
.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: 32
.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
+23 -15
View File
@@ -1,17 +1,22 @@
# Run assembly on the AMD runtime and check correctness
# VIZ=2 to profile
import pathlib
from tinygrad import Tensor, Device, dtypes
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.engine.realize import ExecItem, CompiledRunner
from tinygrad.renderer import ProgramSpec
from tinygrad.uop.ops import track_rewrites, UOp
from tinygrad.helpers import TracingKey
from tinygrad.helpers import TracingKey, 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
N = 8192
scale = 10.0
import torch
@@ -34,23 +39,26 @@ C_asm.uop.buffer.allocate()
# ** run gemms
@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name,), ret=ret))
def get_asm_gemm(ast:UOp, fp:pathlib.Path) -> ProgramSpec:
src = fp.read_text()
lib = Device[Device.DEFAULT].compiler.compile(src)
return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[1024, 1, 1], local_size=[256, 1, 1], globals=[0, 1, 2])
# baseline tinygrad
sched = C_tiny.schedule()
assert len(sched) == 1
eis:list[ExecItem] = [sched[-1].lower()]
ast = eis[0].ast
prg = get_asm_gemm(ast, fp)
eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], prg=CompiledRunner(prg)))
ast = sched[-1].ast
# assembly gemm
@track_rewrites(name=lambda ret: TracingKey(ret.name, (ret.function_name,), ret))
def get_asm_prg() -> ProgramSpec:
src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
lib = Device[Device.DEFAULT].compiler.compile(src)
return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[NUM_WG, 1, 1], local_size=[THREADS_PER_WG, 1, 1],
globals=[0, 1, 2], vars=[UOp.variable("SZ", 256, 8192), UOp.variable("NUM_WG", 1, 1024)])
eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], fixedvars={"SZ":N, "NUM_WG":NUM_WG},
prg=CompiledRunner(get_asm_prg())))
for ei in eis:
et = ei.run(wait=True)
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
with Context(DEBUG=2):
for ei in eis:
et = ei.run(wait=True)
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
# ** correctness
+2 -1
View File
@@ -4,7 +4,8 @@ import os, pathlib
os.environ["AMD_AQL"] = "1"
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.runtime.ops_amd import AMDProgram
NUM_WORKGROUPS = 96
WAVE_SIZE = 32
+4 -12
View File
@@ -11,8 +11,8 @@ from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
from tinygrad.renderer import ProgramSpec
from tinygrad.engine.realize import CompiledRunner
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.asm import waitcnt
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.asm import waitcnt
from test.testextra.test_cfg_viz import template
def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
@@ -48,14 +48,12 @@ class TestHW(unittest.TestCase):
])
self.assertEqual(out, [2])
# assembler err
@unittest.expectedFailure
def test_simple_s_mov(self):
out = get_output([
s_mov_b32(s[7], 0x7fffffff),
v_mov_b32_e32(v[1], s[7]),
])
self.assertEqual(out, [2])
self.assertEqual(out, [0x7fffffff])
def test_exec_mov(self):
out = get_output([
@@ -102,8 +100,6 @@ class TestHW(unittest.TestCase):
self.assertEqual(run_fmac(a, -b), f16_to_bits(-10.0))
self.assertEqual(run_fmac(-a, -b), f16_to_bits(14.0))
# assembler err
@unittest.expectedFailure
def test_s_abs_i32(self):
def check(x, y, dst=s[10], scc=0):
for reg,val in [(dst, y), (SCC, scc)]:
@@ -121,8 +117,6 @@ class TestHW(unittest.TestCase):
check(0xffffffff, 0x00000001, scc=1)
check(0, 0, scc=0)
# how do I negate a VGPR operand?
@unittest.expectedFailure
def test_v_rcp_f32_neg_vop3(self):
def v_neg_rcp_f32(x:float, y:float):
out = get_output([
@@ -138,14 +132,12 @@ class TestHW(unittest.TestCase):
v_neg_rcp_f32(-2.0, 0.5)
v_neg_rcp_f32(2.0, -0.5)
# how do I negate a VGPR operand?
@unittest.expectedFailure
def test_v_cndmask_b32_neg(self):
def v_neg(x:float, y:float):
out = get_output([
v_mov_b32_e32(v[1], f32_to_bits(x)),
s_mov_b32(s[10], 1),
v_cndmask_b32_e32(v[1], v[1], -v[1], s[10]),
v_cndmask_b32_e64(v[1], v[1], -v[1], s[10]),
])[0]
assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
Regular → Executable
+42 -18
View File
@@ -1,11 +1,9 @@
import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
#!/usr/bin/env python3
import ctypes, pathlib, argparse, pickle, dataclasses, threading
from typing import Generator
from tinygrad.helpers import temp, unwrap, DEBUG
from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
from tinygrad.runtime.autogen import llvm, rocprof
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.viz.serve import llvm_disasm
from tinygrad.runtime.ops_amd import ProfileSQTTEvent
from tinygrad.runtime.autogen import rocprof
@dataclasses.dataclass(frozen=True)
class InstExec:
@@ -117,26 +115,52 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[st
def worker():
try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
except AttributeError as e:
raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
(t:=threading.Thread(target=worker, daemon=True)).start()
t.join()
return ROCParseCtx
def print_pmc(events:list[ProfilePMCEvent]) -> None:
from tinygrad.viz.serve import unpack_pmc
def print_data(data:dict) -> None:
from tabulate import tabulate
for e in events:
print("**", e.kern)
data = unpack_pmc(e)
print(tabulate([r[:-1] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
# plaintext
if "src" in data: print(data["src"])
# table format
elif "cols" in data:
print(tabulate([r[:len(data["cols"])] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
def main() -> None:
import tinygrad.viz.serve as viz
viz.ctxs = []
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
default=pathlib.Path(temp("profile.pkl", append_user=True)))
parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Kernel to focus on (optional name, default: all kernels)')
parser.add_argument('-n', type=int, default=3, metavar="NUM", help='Max traces to print (optional number, default: 3 traces)')
args = parser.parse_args()
with args.profile.open("rb") as f: profile = pickle.load(f)
#rctx = decode(profile, disasm)
#print('SQTT:', rctx.inst_execs.keys())
print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])
viz.get_profile(profile)
# List all kernels
if args.kernel is None:
for c in viz.ctxs:
print(c["name"])
for s in c["steps"]: print(" "+s["name"])
return None
# Find kernel trace
trace = next((c for c in viz.ctxs if c["name"] == f"Exec {args.kernel}"), None)
if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
n = 0
for s in trace["steps"]:
print(s["name"])
data = viz.get_render(s["query"])
print_data(data)
n += 1
if n > args.n: break
if __name__ == "__main__":
main()
+2 -3
View File
@@ -49,7 +49,7 @@ arm = ["unicorn"]
triton = ["triton-nightly>=2.1.0.dev20231014192330"]
linting = [
"pylint",
"mypy==1.18.1",
"mypy==1.19.1",
"typing-extensions",
"pre-commit",
"ruff",
@@ -61,7 +61,7 @@ linting = [
# ]
testing_minimal = [
"numpy",
"torch==2.9.0",
"torch==2.9.1",
"pytest",
"pytest-xdist",
"pytest-timeout",
@@ -188,7 +188,6 @@ select = [
"E72",
"E112", # no-indented-block
"E113", # unexpected-indentation
# "E124",
"E203", # whitespace-before-punctuation
"E272", # multiple-spaces-before-keyword
"E275", # missing-whitespace-after-keyword
+1 -1
View File
@@ -76,7 +76,7 @@ class TestHCQ(unittest.TestCase):
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
@unittest.skipIf(MOCKGPU or Device.DEFAULT in {"CPU"}, "Can't handle async update on MOCKGPU for now")
@unittest.skipIf(Device.DEFAULT in {"CPU"}, "Can't handle async update on CPU device")
def test_wait_late_set(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
-74
View File
@@ -1,74 +0,0 @@
#!/usr/bin/env python
import unittest
from tinygrad.device import Device, BufferSpec
from tinygrad.dtype import dtypes
@unittest.skipUnless(Device.DEFAULT == "QCOM", "QCOM device required to run")
class TestQcom(unittest.TestCase):
def test_image_pitch(self):
dev = Device["QCOM"]
def __validate(imgdt, expected_pitch):
img = dev.allocator.alloc(imgdt.shape[0] * imgdt.shape[1] * 16, options:=BufferSpec(image=imgdt))
pitch = img.texture_info.pitch
assert pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{pitch:X}, expected 0x{expected_pitch:X}"
dev.allocator.free(img, imgdt.shape[0] * imgdt.shape[1] * 16, options)
# Match opencl pitches for perf
__validate(dtypes.imageh((1, 201)), 0x680)
__validate(dtypes.imageh((16, 216)), 0x700)
__validate(dtypes.imageh((16, 9)), 0x80)
__validate(dtypes.imageh((48, 64)), 0x200)
__validate(dtypes.imageh((32, 128)), 0x400)
__validate(dtypes.imageh((96, 128)), 0x400)
__validate(dtypes.imageh((64, 256)), 0x840)
__validate(dtypes.imageh((64, 9)), 0x80)
__validate(dtypes.imageh((192, 256)), 0x840)
__validate(dtypes.imageh((64, 768)), 0x1840)
__validate(dtypes.imageh((256, 49)), 0x1C0)
__validate(dtypes.imageh((128, 9)), 0x80)
__validate(dtypes.imageh((16, 1024)), 0x2080)
__validate(dtypes.imageh((64, 512)), 0x1040)
__validate(dtypes.imageh((16, 512)), 0x1080)
__validate(dtypes.imageh((132, 64)), 0x200)
__validate(dtypes.imageh((4, 512)), 0x1200)
__validate(dtypes.imageh((8, 512)), 0x1100)
__validate(dtypes.imageh((128, 128)), 0x400)
__validate(dtypes.imageh((32, 512)), 0x1040)
__validate(dtypes.imageh((26, 64)), 0x200)
__validate(dtypes.imageh((32, 516)), 0x1040)
__validate(dtypes.imageh((32, 1024)), 0x2040)
__validate(dtypes.imageh((16, 2048)), 0x4080)
__validate(dtypes.imageh((8, 2048)), 0x4100)
__validate(dtypes.imageh((4, 4096)), 0x8200)
__validate(dtypes.imagef((16, 49)), 0x380)
__validate(dtypes.imagef((16, 1024)), 0x4080)
__validate(dtypes.imagef((256, 64)), 0x400)
__validate(dtypes.imagef((64, 512)), 0x2040)
__validate(dtypes.imagef((16, 512)), 0x2080)
__validate(dtypes.imagef((132, 64)), 0x400)
__validate(dtypes.imagef((4, 512)), 0x2200)
__validate(dtypes.imagef((4, 16)), 0x200)
__validate(dtypes.imagef((2, 16)), 0x400)
__validate(dtypes.imagef((8, 512)), 0x2100)
__validate(dtypes.imagef((12, 64)), 0x400)
__validate(dtypes.imagef((3, 32)), 0x400)
__validate(dtypes.imagef((128, 128)), 0x840)
__validate(dtypes.imagef((32, 512)), 0x2040)
__validate(dtypes.imagef((8, 3072)), 0xC100)
__validate(dtypes.imagef((4, 2048)), 0x8200)
__validate(dtypes.imagef((4, 1024)), 0x4200)
__validate(dtypes.imagef((4, 4096)), 0x10200)
__validate(dtypes.imagef((10, 384)), 0x1900)
__validate(dtypes.imagef((24, 64)), 0x400)
__validate(dtypes.imagef((128, 12)), 0xC0)
__validate(dtypes.imagef((10, 24)), 0x200)
__validate(dtypes.imagef((1, 129)), 0x840)
__validate(dtypes.imagef((1, 32)), 0x200)
__validate(dtypes.imagef((1, 64)), 0x400)
__validate(dtypes.imagef((1, 1239)), 0x4D80)
__validate(dtypes.imagef((1, 1)), 0x40)
if __name__ == "__main__":
unittest.main()
+16 -22
View File
@@ -1,7 +1,7 @@
from tinygrad import Tensor, Device, GlobalCounters, TinyJit, dtypes
from tinygrad.helpers import getenv, Context, RING, DEBUG
from tinygrad.helpers import getenv, Context, DEBUG
def test(devs: list[str], N: int, iters:int = 10):
def test(devs: list[str], N: int, iters:int = 10, name:str = "allreduce"):
@TinyJit
def f(t: Tensor) -> Tensor: t.sum(0).realize()
@@ -17,39 +17,33 @@ def test(devs: list[str], N: int, iters:int = 10):
i_secs = GlobalCounters.time_sum_s
i_gflops = GlobalCounters.global_ops/i_secs/10**9
i_gbs = (N*4)/i_secs/10**9
print(f"{'ring_allreduce' if RING >= 2 else 'naive_allreduce'} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
print(f"{name} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
secs += i_secs
gflops += i_gflops
gbs += i_gbs
return (gflops/iters, gbs/iters, secs/iters)
def run(sz, n_gpus=6, iters=10, use_ring=False):
def run(sz, n_gpus=6, iters=10, ring=0, all2all=0):
devs = tuple([f"{Device.DEFAULT}:{x}" for x in range(n_gpus)])
N = sz // dtypes.float32.itemsize
with Context(RING=(2 if use_ring else 0), DEBUG=max(DEBUG.value, 2)): return test(devs, N, iters=iters)
name = "all2all" if all2all else ("ring" if ring else "naive")
with Context(RING=(2 if ring else 0), ALL2ALL=(2 if all2all else 0), JIT_BATCH_SIZE=0, DEBUG=max(DEBUG.value, 2)):
return test(devs, N, iters=iters, name=name)
def main():
ONLY_RING = getenv("ONLY_RING", 0)
n_gpus = getenv("GPUS", 6)
iters = getenv("ITERS", 10)
sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
if getenv("BENCHMARK_SPLIT"):
l, r = 0, 512
while r - l > 1:
m = (l + r) // 2
(ring_gflops, ring_gbs, ring_secs) = run(m * 1024 * 4, n_gpus=n_gpus, iters=100, use_ring=True)
(naive_gflops, naive_gbs, naive_secs) = run(m * 1024 * 4, n_gpus=n_gpus, iters=100, use_ring=False)
if ring_secs > naive_secs: l = m
else: r = m
print("Better split", r * 1024, "elements")
else:
sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
(ring_gflops, ring_gbs, ring_secs) = run(sz, use_ring=True, n_gpus=n_gpus, iters=iters)
if not ONLY_RING: (naive_gflops, naive_gbs, naive_secs) = run(sz, use_ring=False, n_gpus=n_gpus, iters=iters)
print(f"Ring:\n {ring_secs:.6f} seconds/iter\n {ring_gflops:.2f} GFLOP/s\n {ring_gbs:.2f} GB/s")
if not ONLY_RING: print(f"Naive:\n {naive_secs:.6f} seconds/iter\n {naive_gflops:.2f} GFLOP/s\n {naive_gbs:.2f} GB/s")
results = {}
for name, kwargs in [("naive", {}), ("ring", {"ring": 2}), ("all2all", {"all2all": 2})]:
results[name] = run(sz, n_gpus=n_gpus, iters=iters, **kwargs)
print("\n=== RESULTS ===")
for name, (gflops, gbs, secs) in results.items():
print(f"{name.upper()}:\n {secs:.6f} seconds/iter\n {gflops:.2f} GFLOP/s\n {gbs:.2f} GB/s")
if __name__ == "__main__":
main()
+15 -6
View File
@@ -52,6 +52,7 @@ class AMDDriver(VirtDriver):
self.doorbells = {}
self.next_doorbell = collections.defaultdict(int)
self.mmu_event_ids = []
self._executing = False # re-entrancy guard for _emulate_execute
for i in range(gpus): self._prepare_gpu(i+1)
@@ -125,6 +126,9 @@ class AMDDriver(VirtDriver):
if struct.gpu_id not in self.gpus: return -1
struct.handle = self._alloc_handle()
self.object_by_handle[struct.handle] = copy.deepcopy(struct) # save memory struct to know what mem it is
# Track signal memory (uncached + coherent) - progress queues when written to
if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED:
self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute())
elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
self.object_by_handle.pop(struct.handle)
elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
@@ -173,9 +177,14 @@ class AMDDriver(VirtDriver):
return 0
def _emulate_execute(self):
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.executing: any_progress |= q.execute() > 0
if self._executing: return # prevent re-entrancy
self._executing = True
try:
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.executing: any_progress |= q.execute() > 0
finally:
self._executing = False
+4 -1
View File
@@ -7,6 +7,7 @@ import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.
SDMA_MAX_COPY_SIZE = 0x400000
regCOMPUTE_PGM_LO = 0x1bac + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_PGM_RSRC2 = 0x1bb3 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_USER_DATA_0 = 0x1be0 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_NUM_THREAD_X = 0x1ba7 + amd_gpu.GC_BASE__INST0_SEG0
regGRBM_GFX_INDEX = 0x2200 + amd_gpu.GC_BASE__INST0_SEG1
@@ -179,14 +180,16 @@ class PM4Executor(AMDQueue):
prg_addr = (self.gpu.regs[regCOMPUTE_PGM_LO] + (self.gpu.regs[regCOMPUTE_PGM_LO + 1] << 32)) << 8
args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32)
lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)]
rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2]
prg_sz = 0
for st,sz in self.gpu.mapped_ranges:
if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
# Pass valid memory ranges to Python emulator for bounds checking
# Pass valid memory ranges and rsrc2 to Python emulator for bounds checking and SGPR/VGPR layout
if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
if hasattr(remu, 'rsrc2'): remu.rsrc2 = rsrc2
err = remu.run_asm(prg_addr, prg_sz, *gl, *lc, args_addr)
if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
+3 -2
View File
@@ -18,12 +18,13 @@ def _try_dlopen_gpuocelot():
class PythonRemu:
"""Python RDNA3 emulator wrapper that matches the libremu.so interface."""
valid_mem_ranges: set[tuple[int, int]] = set()
rsrc2: int = 0x19c # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
from extra.assembly.rdna3.emu import run_asm, set_valid_mem_ranges
from extra.assembly.amd.emu import run_asm, set_valid_mem_ranges
# Pad ranges to handle GPU loads that may read past small buffers (e.g. s_load_b128 on 12-byte buffer)
set_valid_mem_ranges({(start, size + 4096) for start, size in self.valid_mem_ranges})
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr)
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2)
def _try_dlopen_remu():
# Use Python emulator only if PYTHON_REMU=1
+33 -12
View File
@@ -15,7 +15,7 @@ libc.munmap.restype = ctypes.c_int
NVSubDevice = collections.namedtuple('NVSubDevice', ['device'])
NVUserMode = collections.namedtuple('NVUserMode', ['subdevice'])
NVVASpace = collections.namedtuple('NVVASpace', ['device'])
NVAllocation = collections.namedtuple('NVAllocation', ['device', 'size'])
NVAllocation = collections.namedtuple('NVAllocation', ['device', 'size', 'is_signal'])
NVChannelGroup = collections.namedtuple('NVChannelGroup', ['device'])
NVContextShare = collections.namedtuple('NVContextShare', ['channel_group'])
NVGPFIFO = collections.namedtuple('NVGPFIFO', ['device', 'token'])
@@ -41,12 +41,14 @@ class NVDevFileDesc(VirtFileDesc):
super().__init__(fd)
self.driver, self.gpu = driver, gpu
self._mapping_userland = False
self._mapping_signal = False
def ioctl(self, fd, request, argp): return self.driver.dev_ioctl(self.gpu, request, argp)
def mmap(self, start, sz, prot, flags, fd, offset):
start = libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
if self._mapping_userland:
if self._mapping_userland or self._mapping_signal:
self.driver.track_address(start, start+sz, lambda mv,off: None, lambda mv, off: self.driver._gpu_mmio_write(mv, off, self.gpu))
self._mapping_signal = False
return start
class NVDriver(VirtDriver):
@@ -65,6 +67,7 @@ class NVDriver(VirtDriver):
self.object_by_handle = {}
self.opened_fds = {}
self.next_doorbell = collections.defaultdict(int)
self._executing = False # re-entrancy guard for _gpu_mmio_write
for i in range(gpus): self._prepare_gpu(i)
@@ -115,7 +118,8 @@ class NVDriver(VirtDriver):
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
params = nv_gpu.NV_MEMORY_ALLOCATION_PARAMS.from_address(params_ptr)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVAllocation(self.object_by_handle[struct.hObjectParent], params.size)
is_signal = struct.hClass == nv_gpu.NV1_MEMORY_SYSTEM # signal memory uses NV1_MEMORY_SYSTEM (uncached)
self.object_by_handle[struct.hObjectNew] = NVAllocation(self.object_by_handle[struct.hObjectParent], params.size, is_signal)
elif struct.hClass == nv_gpu.KEPLER_CHANNEL_GROUP_A:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVGPU)
struct.hObjectNew = self._alloc_handle()
@@ -206,7 +210,6 @@ class NVDriver(VirtDriver):
def ctl_ioctl(self, req, argp):
nr = req & 0xff
if nr == nv_gpu.NV_ESC_RM_ALLOC: return self.rm_alloc(argp)
elif nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY: pass
elif nr == nv_gpu.NV_ESC_RM_CONTROL: return self.rm_control(argp)
elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
st:Any = nv_gpu.nv_ioctl_nvos33_parameters_with_fd.from_address(argp)
@@ -215,6 +218,10 @@ class NVDriver(VirtDriver):
file = self.opened_fds[st.fd]
assert isinstance(file, NVDevFileDesc)
file._mapping_userland = True
elif isinstance(obj, NVAllocation) and obj.is_signal:
file = self.opened_fds[st.fd]
assert isinstance(file, NVDevFileDesc)
file._mapping_signal = True
elif nr == nv_gpu.NV_ESC_RM_FREE:
st = nv_gpu.NVOS00_PARAMETERS.from_address(argp)
self.object_by_handle.pop(st.hObjectOld)
@@ -256,12 +263,26 @@ class NVDriver(VirtDriver):
else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
return 0
def dev_ioctl(self, dev, req, argp): return 0
def dev_ioctl(self, dev, req, argp):
nr = req & 0xff
# Handle NV_ESC_RM_ALLOC_MEMORY for host/signal memory
if nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
st:Any = nv_gpu.nv_ioctl_nvos02_parameters_with_fd.from_address(argp)
# Track host memory (signal memory) - progress queues when written to
if st.params.hClass == nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR:
self.track_address(st.params.pMemory, st.params.pMemory + st.params.limit + 1,
lambda mv,off: None, lambda mv, off: self._gpu_mmio_write(mv, off, None))
return 0
def _gpu_mmio_write(self, mv, off, gpu):
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.ctrl.GPGet != q.ctrl.GPPut:
any_progress |= q.execute()
if self._executing: return # prevent re-entrancy
self._executing = True
try:
any_progress = True
while any_progress:
any_progress = False
for gpu in self.gpus.values():
for q in gpu.queues:
if q.ctrl.GPGet != q.ctrl.GPPut:
any_progress |= q.execute()
finally:
self._executing = False
+2
View File
@@ -1,4 +1,6 @@
#!/usr/bin/env python
import os
os.environ['USE_TF'] = '0' # prevent transformers from importing tensorflow
import unittest
from tinygrad import Tensor
import numpy as np
+1 -1
View File
@@ -68,7 +68,7 @@ class TestEfficientNet(unittest.TestCase):
self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
def test_chicken_car(self):
labels = _infer(self.model, np.concat([chicken_img, car_img], axis=0))
labels = _infer(self.model, np.concatenate([chicken_img, car_img], axis=0))
self.assertEqual(_LABELS[labels[0]], "hen")
self.assertEqual(_LABELS[labels[1]], "sports car, sport car")
+10
View File
@@ -193,6 +193,16 @@ class TestCustomKernel(unittest.TestCase):
err = (tst - (a@b)).square().max()
self.assertLess(err.item(), 1e-6)
def test_gemm_multi(self):
devs = ("CPU:0", "CPU:1")
N = 16
a = Tensor.randn(N, N).shard_(devs, axis=0)
b = Tensor.randn(N, N).to(devs)
c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
err = (tst - (a@b)).square().max()
self.assertLess(err.item(), 1e-6)
def test_gemm_backward_custom(self): self.test_gemm_backward(True)
# NOTE: grad_fxn doesn't work with pyrender
def test_gemm_backward(self, custom_backward_gemm=False):
+12 -19
View File
@@ -1,7 +1,7 @@
import unittest, operator, math
from tinygrad import Tensor, dtypes, Device
from tinygrad.dtype import DType, truncate
from tinygrad.helpers import CI, getenv, CPU_LLVM
from tinygrad.helpers import CI, getenv
from tinygrad.tensor import _to_np_dtype
from tinygrad.device import is_dtype_supported
from tinygrad.runtime.ops_python import from_storage_scalar
@@ -9,7 +9,7 @@ from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
import numpy as np
import pytest
from hypothesis import assume, given, strategies as strat, settings, HealthCheck
from hypothesis import assume, given, strategies as strat, settings
pytestmark = pytest.mark.filterwarnings("ignore")
@@ -48,7 +48,7 @@ class ht:
int32 = strat.integers(-2147483648, 2147483647)
int64 = strat.integers(-9223372036854775808, 9223372036854775807)
bool = strat.booleans()
ht.bfloat16 = ht.uint16
ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0) # filter subnormal bfloat16
ht.fp8e4m3 = ht.uint8
ht.fp8e5m2 = ht.uint8
@@ -138,7 +138,6 @@ class TestDTypeALU(unittest.TestCase):
def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
@unittest.skipIf(CPU_LLVM, "bfloat16 precision issues with CPU_LLVM")
@given(ht.bfloat16, strat.sampled_from(unary_operations))
def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
@@ -206,29 +205,23 @@ class TestDTypeALU(unittest.TestCase):
@given(ht.int32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype)
@settings(suppress_health_check=[HealthCheck.filter_too_much])
@given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
@given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
float_strat = float_strat.filter(lambda x: 0 < x < dtypes.max(unsigned_dtype))
universal_test_cast(a.draw(float_strat), float_dtype, unsigned_dtype)
universal_test_cast(a, float_dtype, unsigned_dtype)
@settings(suppress_health_check=[HealthCheck.filter_too_much])
@given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
@given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
overflow_strat = float_strat.filter(lambda x: x > dtypes.max(unsigned_dtype) and x <= dtypes.max(dtypes.int32))
universal_test_cast(a.draw(overflow_strat), float_dtype, unsigned_dtype)
universal_test_cast(a, float_dtype, unsigned_dtype)
@settings(suppress_health_check=[HealthCheck.filter_too_much])
@given(strat.data(), strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
@given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
float_strat = {dtypes.float16: ht.float16, dtypes.float32: ht.float32, dtypes.float64: ht.float64}[float_dtype]
underflow_strat = float_strat.filter(lambda x: x < 0 and x >= dtypes.min(dtypes.int32))
universal_test_cast(a.draw(underflow_strat), float_dtype, unsigned_dtype)
universal_test_cast(a, float_dtype, unsigned_dtype)
@unittest.expectedFailure
def test_unsafe_cast_float_to_int_failure(self):
+60
View File
@@ -44,6 +44,66 @@ class TestImageCopy(unittest.TestCase):
@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageDType(unittest.TestCase):
def test_image_pitch(self):
def __validate(imgdt, expected_pitch):
assert imgdt.pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{imgdt.pitch:X}, expected 0x{expected_pitch:X}"
# Match opencl pitches for perf
__validate(dtypes.imageh((1, 201)), 0x680)
__validate(dtypes.imageh((16, 216)), 0x700)
__validate(dtypes.imageh((16, 9)), 0x80)
__validate(dtypes.imageh((48, 64)), 0x200)
__validate(dtypes.imageh((32, 128)), 0x400)
__validate(dtypes.imageh((96, 128)), 0x400)
__validate(dtypes.imageh((64, 256)), 0x840)
__validate(dtypes.imageh((64, 9)), 0x80)
__validate(dtypes.imageh((192, 256)), 0x840)
__validate(dtypes.imageh((64, 768)), 0x1840)
__validate(dtypes.imageh((256, 49)), 0x1C0)
__validate(dtypes.imageh((128, 9)), 0x80)
__validate(dtypes.imageh((16, 1024)), 0x2080)
__validate(dtypes.imageh((64, 512)), 0x1040)
__validate(dtypes.imageh((16, 512)), 0x1080)
__validate(dtypes.imageh((132, 64)), 0x200)
__validate(dtypes.imageh((4, 512)), 0x1200)
__validate(dtypes.imageh((8, 512)), 0x1100)
__validate(dtypes.imageh((128, 128)), 0x400)
__validate(dtypes.imageh((32, 512)), 0x1040)
__validate(dtypes.imageh((26, 64)), 0x200)
__validate(dtypes.imageh((32, 516)), 0x1040)
__validate(dtypes.imageh((32, 1024)), 0x2040)
__validate(dtypes.imageh((16, 2048)), 0x4080)
__validate(dtypes.imageh((8, 2048)), 0x4100)
__validate(dtypes.imageh((4, 4096)), 0x8200)
__validate(dtypes.imagef((16, 49)), 0x380)
__validate(dtypes.imagef((16, 1024)), 0x4080)
__validate(dtypes.imagef((256, 64)), 0x400)
__validate(dtypes.imagef((64, 512)), 0x2040)
__validate(dtypes.imagef((16, 512)), 0x2080)
__validate(dtypes.imagef((132, 64)), 0x400)
__validate(dtypes.imagef((4, 512)), 0x2200)
__validate(dtypes.imagef((4, 16)), 0x200)
__validate(dtypes.imagef((2, 16)), 0x400)
__validate(dtypes.imagef((8, 512)), 0x2100)
__validate(dtypes.imagef((12, 64)), 0x400)
__validate(dtypes.imagef((3, 32)), 0x400)
__validate(dtypes.imagef((128, 128)), 0x840)
__validate(dtypes.imagef((32, 512)), 0x2040)
__validate(dtypes.imagef((8, 3072)), 0xC100)
__validate(dtypes.imagef((4, 2048)), 0x8200)
__validate(dtypes.imagef((4, 1024)), 0x4200)
__validate(dtypes.imagef((4, 4096)), 0x10200)
__validate(dtypes.imagef((10, 384)), 0x1900)
__validate(dtypes.imagef((24, 64)), 0x400)
__validate(dtypes.imagef((128, 12)), 0xC0)
__validate(dtypes.imagef((10, 24)), 0x200)
__validate(dtypes.imagef((1, 129)), 0x840)
__validate(dtypes.imagef((1, 32)), 0x200)
__validate(dtypes.imagef((1, 64)), 0x400)
__validate(dtypes.imagef((1, 1239)), 0x4D80)
__validate(dtypes.imagef((1, 1)), 0x40)
def test_image_and_back(self):
data = Tensor.randn(9*27*4).realize()
tst = data.numpy()
+5
View File
@@ -256,6 +256,11 @@ class TestMultiTensor(unittest.TestCase):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_allreduce_all2all(self):
with Context(ALL2ALL=2):
a,b = _test_allreduce(Tensor.rand(256, 256))
np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
def test_copy_jit(self):
@TinyJit
def copy_tensor(x:Tensor): return (x.to(f"{x.device.split(':')[0]}:1") + 1)
+15 -15
View File
@@ -848,7 +848,7 @@ class TestOps(unittest.TestCase):
helper_test_op([(45,65)], lambda x: x.cos())
helper_test_op([()], lambda x: x.cos())
if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
helper_test_op(None, lambda x: x.cos(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
@@ -859,8 +859,8 @@ class TestOps(unittest.TestCase):
helper_test_op([(45,65)], lambda x: x.tan(), low=-5, high=5)
helper_test_op([()], lambda x: x.tan())
if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
helper_test_op(None, lambda x: x.tan(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
helper_test_op(None, lambda x: x.tan(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
def test_asin(self):
@@ -1655,7 +1655,7 @@ class TestOps(unittest.TestCase):
def test_broadcast_full(self):
for torch_op, tinygrad_op in [(torch.add, Tensor.add), (torch.sub, Tensor.sub), (torch.mul, Tensor.mul),
(torch.div, Tensor.div), (torch.pow, Tensor.pow)]:
for shapes in [((5,13,24,16), (5,1,24,1)), ((1,3,1,7,1), (2,1,5,1,8))]:
for shapes in [((5,3,14,16), (5,1,14,1)), ((1,3,1,7,1), (2,1,5,1,8))]:
with self.subTest(op=torch_op.__name__, shapes=shapes):
if tinygrad_op != Tensor.pow:
helper_test_op(shapes, torch_op, tinygrad_op)
@@ -2078,7 +2078,7 @@ class TestOps(unittest.TestCase):
lambda x,w: Tensor.conv2d(x,w,padding=[1,1,1,1,1,1]), grad_rtol=1e-5)
def test_simple_conv2d_m4(self):
helper_test_op([(1,16,18,18), (16,16,3,3)],
helper_test_op([(1,16,9,9), (16,16,3,3)],
lambda x,w: torch.nn.functional.conv2d(x,w),
lambda x,w: Tensor.conv2d(x,w), atol=1e-05, grad_rtol=1e-5)
@@ -2535,7 +2535,7 @@ class TestOps(unittest.TestCase):
@slow_test
def test_avg_pool2d(self):
shape = (32,2,111,28)
shape = (32,2,11,28)
for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
with self.subTest(kernel_size=ksz):
helper_test_op([shape],
@@ -2549,7 +2549,7 @@ class TestOps(unittest.TestCase):
@slow_test
def test_avg_pool2d_padding(self):
shape = (32,2,111,28)
shape = (32,2,11,28)
for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
for p in [1, (1,0), (0,1)]:
with self.subTest(kernel_size=ksz, padding=p):
@@ -2557,10 +2557,10 @@ class TestOps(unittest.TestCase):
lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, padding=p),
lambda x: Tensor.avg_pool2d(x, kernel_size=ksz, padding=p), rtol=1e-5)
with self.assertRaises(ValueError):
Tensor.avg_pool2d(Tensor.randn((32,2,111,28)), kernel_size=(2,2), padding=(1,1,1))
Tensor.avg_pool2d(Tensor.randn((32,2,11,28)), kernel_size=(2,2), padding=(1,1,1))
def test_avg_pool2d_asymmetric_padding(self):
shape = (32,2,111,28)
shape = (32,2,11,28)
for p in [(0,1,0,1), (2,1,2,1), (2,0,2,1)]:
with self.subTest(padding=p):
helper_test_op([shape],
@@ -2571,7 +2571,7 @@ class TestOps(unittest.TestCase):
@slow_test
def test_avg_pool2d_padding_not_counted(self):
shape = (32,2,111,28)
shape = (32,2,11,28)
for ksz in [(2,2), (3,3), 2, 3, (3,2)]:
with self.subTest(kernel_size=ksz):
helper_test_op([shape],
@@ -2607,9 +2607,9 @@ class TestOps(unittest.TestCase):
lambda x: Tensor.avg_pool2d(x, kernel_size=(3,3), stride=3, padding=1, ceil_mode=True, count_include_pad=True))
def test_global_avg_pool2d(self):
helper_test_op([(32,2,111,28)],
lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(111,28)),
lambda x: Tensor.avg_pool2d(x, kernel_size=(111,28)), rtol=1e-5)
helper_test_op([(32,2,11,28)],
lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=(11,28)),
lambda x: Tensor.avg_pool2d(x, kernel_size=(11,28)), rtol=1e-5)
def test_avg_pool3d(self):
# TODO: AMD_LLVM has larger atol
@@ -3142,10 +3142,10 @@ class TestOps(unittest.TestCase):
lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), Tensor(weight), reduction=r))
def test_nll_loss_3d_weight(self):
target = np.random.randint(0, 10, (32,3,3,3), dtype=np.int32).tolist()
target = np.random.randint(0, 10, (16,3,3,3), dtype=np.int32).tolist()
weight = np.random.normal(0, 1, (10,)).astype(np.float32).tolist()
for r in ("mean", "sum", "none"):
helper_test_op([(32,10,3,3,3)],
helper_test_op([(16,10,3,3,3)],
lambda x: torch.nn.functional.nll_loss(torch.nn.functional.log_softmax(x, dim=1), torch.tensor(target), torch.tensor(weight), reduction=r),
lambda x: x.log_softmax(axis=1).nll_loss(Tensor(target), Tensor(weight), reduction=r))
-137
View File
@@ -478,143 +478,6 @@ class TestUOpGraph(unittest.TestCase):
for u in uops:
self.assertNotEqual(u.dtype, dtypes.long)
def test_in_out_of_bounds_access(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 0), ptr=True),))
to_uops_list([ld0])
ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 15), ptr=True),))
to_uops_list([ld1])
ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 7), ptr=True),))
to_uops_list([ld1])
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 42), ptr=True),))
with self.assertRaises(RuntimeError): to_uops_list([ld0])
def test_in_out_of_bounds_access_symbolic(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 1, 10), ptr=True),))
to_uops_list([ld0])
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 15), ptr=True),))
to_uops_list([ld0])
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(Variable("i", 0, 20), ptr=True),))
with self.assertRaises(RuntimeError): to_uops_list([ld0])
def test_in_out_of_bounds_access_gated_store(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
v = Variable("v", 0, 20)
st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v.valid(v<16)), UOp.const(dtypes.int, 0)))
to_uops_list([st0])
st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v.valid(v<20)), v))
with self.assertRaises(RuntimeError): to_uops_list([st1])
@unittest.skip("if not allowed in graph")
def test_in_bounds_access_gated_local(self):
with Context(IGNORE_OOB=0):
# Define buffers
gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")
# Define indices, valids and barrier
gidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 416),), "gidx0")
lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "lidx0")
gate = (gidx<400) & (lidx<8)
local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
# Load from local memory (after the IF/barrier)
local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx, ptr=True), if_barrier))
# Store to global memory
global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
to_uops_list([global_store])
def test_load_with_float_in_index(self):
with Context(IGNORE_OOB=0):
ridx = UOp.range(20, 0)
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
i = (ridx.cast(dtypes.float)*0.68).trunc().cast(dtypes.int)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid((0<=i)&(i<16)), ptr=True),))
to_uops_list([ld0])
glblfloat = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(20), (), 0)
ldfloat = UOp(Ops.LOAD, dtypes.float, (glblfloat.index(ridx),))
i = (ldfloat+3.14).cast(dtypes.int)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i, ((0<=i)&(i<16)), ptr=True),))
def test_load_cast_to_bool(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
ridx = UOp.range(20, 0)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(ridx.valid(ridx.cast(dtypes.bool).logical_not()), ptr=True),))
to_uops_list([ld0])
@unittest.skip("Bool load is not supported yet")
def test_load_mask(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
ridx = UOp.range(20, 0)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
to_uops_list([ld0])
def test_out_of_bounds_off_by_one_access(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(dtypes.int, 16), ptr=True),))
with self.assertRaises(RuntimeError): to_uops_list([ld0])
def test_in_out_bounds_access_with_mask(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid((5<gidx0)&(gidx0<16)), ptr=True),))
ld1 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<16), ptr=True),))
to_uops_list([ld0, ld1])
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<17), ptr=True),))
with self.assertRaises(RuntimeError): to_uops_list([ld0])
def test_in_out_of_bounds_access_symbolic_mask(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
i = Variable("i", 1, 80)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<10), ptr=True),))
to_uops_list([ld0])
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<15), ptr=True),))
to_uops_list([ld0])
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(i.valid(i<20), ptr=True),))
with self.assertRaises(RuntimeError): to_uops_list([ld0])
def test_in_out_of_bounds_access_index_load(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
gidx0 = UOp.range(42, 0, AxisType.GLOBAL)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(gidx0.valid(gidx0<8), ptr=True),)).cast(dtypes.index)
ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index((ld0*2).valid((ld0>=0)&(ld0<32)), ptr=True),))
to_uops_list([ld1])
ld1 = UOp(Ops.LOAD, dtypes.int, (glbl1.index((ld0*2).valid((ld0>=0)&(ld0<64)), ptr=True),))
with self.assertRaises(RuntimeError): to_uops_list([ld1])
def test_bounds_with_loaded_bool(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 0)
gidx0 = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
ld0 = glbl0.index(gidx0, ptr=True).load()
ld1 = glbl1.index(gidx0.valid(ld0), ptr=True).load()
with self.assertRaises(RuntimeError): to_uops_list([ld1])
def test_fold_gated_load(self):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0)
glbl1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 1)
+1 -1
View File
@@ -10,7 +10,7 @@ from tinygrad.renderer import ProgramSpec
from tinygrad.helpers import TracingKey, getenv
from tinygrad.engine.realize import ExecItem, CompiledRunner
from extra.assembly.rdna3.autogen import *
from extra.assembly.amd.autogen.rdna3.ins import *
# TODO: use the RDNA3 renderer when it's in master
template = """.text
+2 -2
View File
@@ -47,10 +47,10 @@ class TestKeccak(unittest.TestCase):
ha_ref, hb_ref = hasher(a), hasher(b)
tres = Tensor.stack(*(Tensor(d) for d in (a, b))).keccak(name)
ha, hb = tres[0].data(), tres[1].data()
ha, hb = bytes(tres[0].data()), bytes(tres[1].data())
self.assertEqual(ha_ref, ha)
self.assertEqual(ha_ref, Tensor(a).keccak(name).data())
self.assertEqual(ha_ref, bytes(Tensor(a).keccak(name).data()))
self.assertEqual(hb_ref, hb)
def test_referenced(self):
+38
View File
@@ -0,0 +1,38 @@
import unittest
from tinygrad import UOp, dtypes
class TestUOpRepr(unittest.TestCase):
def test_simple_const(self):
a = UOp.const(dtypes.int, 42)
self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())")
def test_different_consts(self):
a, b = UOp.const(dtypes.int, 42), UOp.const(dtypes.int, 3)
expected = (
"UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
" UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
" UOp(Ops.CONST, dtypes.int, arg=3, src=()),))"
)
self.assertEqual(repr(a+b), expected)
def test_walrus_operator_indentation(self):
# The reference should have the same indentation as the definition
a = UOp.const(dtypes.int, 42)
expected = (
"UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
" x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
" x0,))"
)
self.assertEqual(repr(a+a), expected)
def test_nested_walrus_indentation(self):
# Ensure indentation is consistent at multiple levels
b = (a:=UOp.const(dtypes.int, 1)) + a
expected = (
"UOp(Ops.MUL, dtypes.int, arg=None, src=(\n" +
" x0:=UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
" x1:=UOp(Ops.CONST, dtypes.int, arg=1, src=()),\n" +
" x1,)),\n" +
" x0,))"
)
self.assertEqual(repr(b*b), expected)
if __name__ == '__main__':
unittest.main()
+179
View File
@@ -0,0 +1,179 @@
import unittest
from tinygrad import dtypes, Variable
from tinygrad.dtype import AddrSpace
from tinygrad.helpers import Context
from tinygrad.uop.ops import Ops, UOp, AxisType
from test.test_uops import to_uops_list
class TestValidateOOB(unittest.TestCase):
"""Test z3 validation of index bounds for different ALU ops and patterns."""
# basic index patterns
def test_const_index(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
to_uops_list([buf.index(UOp.const(dtypes.int, 0), ptr=True).load(dtype=dtypes.int)]) # valid
to_uops_list([buf.index(UOp.const(dtypes.int, 15), ptr=True).load(dtype=dtypes.int)]) # valid (last element)
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(UOp.const(dtypes.int, 16), ptr=True).load(dtype=dtypes.int)]) # off by one
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(UOp.const(dtypes.int, 42), ptr=True).load(dtype=dtypes.int)]) # way out
def test_variable_index(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
to_uops_list([buf.index(Variable("i", 0, 15), ptr=True).load(dtype=dtypes.int)]) # valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(Variable("i", 0, 20), ptr=True).load(dtype=dtypes.int)]) # oob
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(Variable("i", -5, 10), ptr=True).load(dtype=dtypes.int)]) # negative
def test_range_with_mask(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
r = UOp.range(42, 0, AxisType.GLOBAL)
to_uops_list([buf.index(r.valid(r < 16), ptr=True).load(dtype=dtypes.int)]) # valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(r.valid(r < 17), ptr=True).load(dtype=dtypes.int)]) # oob
def test_variable_with_mask(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
v = Variable("v", -5, 80)
to_uops_list([buf.index(v.valid((v >= 0) & (v < 16)), ptr=True).load(dtype=dtypes.int)]) # valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(v.valid(v < 20), ptr=True).load(dtype=dtypes.int)]) # negative not masked
def test_gated_store(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
v = Variable("v", 0, 20)
to_uops_list([buf.index(v.valid(v < 16)).store(0)]) # valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(v.valid(v < 20)).store(0)]) # oob
# ALU ops in index
def test_idiv(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2, ptr=True).load(dtype=dtypes.int)]) # 0..16 oob
def test_mod(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
r = UOp.range(100, 0, AxisType.GLOBAL)
to_uops_list([buf.index(r % 16, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(r % 20, ptr=True).load(dtype=dtypes.int)]) # 0..19 oob
def test_shr(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2, ptr=True).load(dtype=dtypes.int)]) # 0..31 oob
def test_shl(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 0)
r = UOp.range(8, 0, AxisType.GLOBAL)
to_uops_list([buf.index(r << 2, ptr=True).load(dtype=dtypes.int)]) # 0..28 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(r << 4, ptr=True).load(dtype=dtypes.int)]) # 0..112 oob
def test_and(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
r = UOp.range(100, 0, AxisType.GLOBAL)
to_uops_list([buf.index(r & 15, ptr=True).load(dtype=dtypes.int)]) # 0..15 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(r & 31, ptr=True).load(dtype=dtypes.int)]) # 0..31 oob
def test_max(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
to_uops_list([buf.index(Variable("v", -10, 15).maximum(0), ptr=True).load(dtype=dtypes.int)]) # 0..15 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0), ptr=True).load(dtype=dtypes.int)]) # 0..20 oob
def test_xor_in_mask(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
r = UOp.range(32, 0, AxisType.GLOBAL)
to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16))), ptr=True).load(dtype=dtypes.int)]) # 0..15 valid
with self.assertRaises(RuntimeError):
to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20)), ptr=True).load(dtype=dtypes.int)]) # 0..9,20..31 oob
# cast patterns
def test_float_cast_in_index(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
r = UOp.range(20, 0)
i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
to_uops_list([buf.index(i.valid((i >= 0) & (i < 16)), ptr=True).load(dtype=dtypes.int)])
def test_bool_cast_in_mask(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), (), 0)
r = UOp.range(20, 0)
to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not()), ptr=True).load(dtype=dtypes.int)]) # only r=0 valid
# load result as index/mask
def test_load_as_index(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
buf1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(64), (), 1)
r = UOp.range(42, 0, AxisType.GLOBAL)
ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.index)
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)]) # valid
with self.assertRaises(RuntimeError):
to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64)), ptr=True).load(dtype=dtypes.int)]) # oob
def test_load_bool_as_mask(self):
with Context(IGNORE_OOB=0, SPEC=2):
buf_bool = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
buf_int = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(8), (), 1)
gidx = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
ld_bool = buf_bool.index(gidx, ptr=True).load()
with self.assertRaises(RuntimeError):
to_uops_list([buf_int.index(gidx.valid(ld_bool), ptr=True).load()]) # gidx 0..15, buf_int size 8
# skipped tests (moved from test_uop_graph.py)
@unittest.skip("if not allowed in graph")
def test_in_bounds_access_gated_local(self):
with Context(IGNORE_OOB=0):
# Define buffers
gbuf = UOp(Ops.DEFINE_GLOBAL, dtypes.uint.ptr(400), (), 0)
sbuf = UOp(Ops.DEFINE_LOCAL, dtypes.uint.ptr(8, addrspace=AddrSpace.LOCAL), (), "temp0")
# Define indices, valids and barrier
gidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 416),), "gidx0")
lidx = UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 10),), "lidx0")
gate = (gidx<400) & (lidx<8)
local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
# Load from local memory (after the IF/barrier)
local_load = UOp(Ops.LOAD, dtypes.uint, (sbuf.index(lidx, ptr=True), if_barrier))
# Store to global memory
global_store = UOp(Ops.STORE, dtypes.void, (gbuf.index(gidx), local_load))
to_uops_list([global_store])
@unittest.skip("Bool load is not supported yet")
def test_load_mask(self):
with Context(IGNORE_OOB=0):
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), (), 0)
mask = UOp(Ops.DEFINE_GLOBAL, dtypes.bool.ptr(16), (), 0)
ridx = UOp.range(20, 0)
ld0 = UOp(Ops.LOAD, dtypes.int, (glbl0.index(UOp.const(ridx, ridx<16&mask), ptr=True)))
to_uops_list([ld0])
if __name__ == "__main__":
unittest.main()
+10 -1
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
from typing import Final, ClassVar, Callable, Literal
import math, struct, ctypes, functools
from dataclasses import dataclass, fields
from tinygrad.helpers import getenv, prod
from tinygrad.helpers import getenv, prod, round_up, next_power2
from enum import Enum, auto
class InvalidTypeMetaClass(type):
@@ -101,6 +101,15 @@ class ImageDType(PtrDType):
assert addrspace == AddrSpace.GLOBAL, "images can't be local"
return self
def __repr__(self): return f"dtypes.{self.name}({self.shape})" + (f'.vec({self.v})' if self.v != 1 else '')
@property
def pitch(self):
imgw, imgh, itemsize_log = self.shape[1], self.shape[0], int(math.log2(self.itemsize))
pitchalign = max(6, 11 - int(math.log2(imgh))) if imgh > 1 else 6
align_up = max(1, (8 // itemsize_log + 1) - imgh // 32) if pitchalign == 6 else (2 ** (pitchalign - itemsize_log - 2))
granularity = 128 if self.itemsize == 4 else 256
pitch_add = (1 << pitchalign) if min(next_power2(imgw), round_up(imgw, granularity)) - align_up + 1 <= imgw and imgw > granularity//2 else 0
return round_up(imgw * 4 * self.itemsize, 1 << pitchalign) + pitch_add
class dtypes:
@staticmethod
+4 -3
View File
@@ -50,6 +50,7 @@ def strip_parens(fst:str) -> str: return fst[1:-1] if fst and fst[0]=='(' and fs
def ceildiv(num, amt): return int(ret) if isinstance((ret:=-(num//-amt)), float) else ret
def round_up(num:int, amt:int) -> int: return (num+amt-1)//amt * amt
def round_down(num:int, amt:int) -> int: return -round_up(-num, amt)
def next_power2(x): return 1 if x == 0 else 1 << (x - 1).bit_length()
# cstyle div and mod
def cdiv(x:int, y:int) -> int: return abs(x)//abs(y)*(1,-1)[x*y<0] if y != 0 else 0
def cmod(x:int, y:int) -> int: return x-cdiv(x,y)*y
@@ -181,8 +182,8 @@ JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVa
WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1)
LRU = ContextVar("LRU", 1)
SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
RING, ALL2ALL = ContextVar("RING", 1), ContextVar("ALL2ALL", 0)
CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
@@ -197,7 +198,7 @@ AMD_CC, CPU_CC, NV_CC, CUDA_CC = ContextVar("AMD_CC", ""), ContextVar("CPU_CC",
QCOM_CC = ContextVar("QCOM_CC", "")
# VIZ implies PROFILE, but you can run PROFILE without VIZ
VIZ = ContextVar("VIZ", 0)
PROFILE = ContextVar("PROFILE", VIZ.value)
PROFILE = ContextVar("PROFILE", abs(VIZ.value))
SPEC = ContextVar("SPEC", 1)
# TODO: disable by default due to speed
IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
+13 -19
View File
@@ -1,6 +1,6 @@
from typing import Callable, cast, Any
from tinygrad.dtype import AddrSpace, DType, PtrDType, ImageDType, dtypes
from tinygrad.helpers import DEBUG, OSX, unwrap, charptr
from tinygrad.helpers import DEBUG, OSX, unwrap, charptr, fromimport
from tinygrad.renderer import Renderer
from tinygrad.renderer.cstyle import CUDARenderer
from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
@@ -115,7 +115,8 @@ def nidx(b:mesa.nir_builder, buf, off, dtype, gate=None) -> mesa.nir_def:
return if_phi(b, gate, f, lambda: buf) if gate is not None else f()
class NIRRenderer(Renderer):
suffix = "NAK"
suffix = "NIR"
nir_options: bytes
global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
@@ -158,13 +159,17 @@ class NIRRenderer(Renderer):
(UPat(Ops.ENDIF, name="x"), lambda ctx,x: (lambda _: mesa.nir_def())(mesa.nir_pop_if(ctx.b, ctx.r[x.src[0]])))
])
def __init__(self): mesa.glsl_type_singleton_init_or_ref()
def __reduce__(self): return self.__class__, self.args
def __init__(self, *args):
self.compiler = fromimport("tinygrad.runtime.support.compiler_mesa", self.__class__.__name__.replace("Renderer", "Compiler"))(*args)
self.args = args
if hasattr(self.compiler, "nir_options"): self.nir_options = self.compiler.nir_options
mesa.glsl_type_singleton_init_or_ref()
def __del__(self):
with contextlib.suppress(AttributeError): mesa.glsl_type_singleton_decref()
@property
def nir_options(self): raise NotImplementedError("needs nir_options")
def param(self, b:mesa.nir_builder, x, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
def prerender(self, uops:list[UOp]):
self.b = mesa.nir_builder_init_simple_shader(mesa.MESA_SHADER_COMPUTE, mesa.nir_shader_compiler_options.from_buffer_copy(self.nir_options), None)
@@ -216,20 +221,9 @@ class NIRRenderer(Renderer):
return ret
class NIRRendererWithOpts(NIRRenderer):
def __init__(self, dev=None, nir_options=None):
self.dev, self._nir_options = dev, nir_options
super().__init__()
def __reduce__(self): return self.__class__, (None, self.nir_options)
@property
def nir_options(self):
if self._nir_options is None: self._nir_options = self.dev.compiler.nir_options
return self._nir_options
class NAKRenderer(NIRRendererWithOpts):
class NAKRenderer(NIRRenderer):
device = "NV"
param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz),
intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])(
lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
@@ -261,7 +255,7 @@ _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2
nc=4, bs=32, num_components=4, srcs=lambda b,img,coord:[nsrc(x) for x in [img, tovec(b, coord), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
lambda b,img,coord,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
class IR3Renderer(NIRRendererWithOpts):
class IR3Renderer(NIRRenderer):
device = "QCOM"
def nload_img(ctx,img,coord):
+1 -1
View File
@@ -12,7 +12,7 @@ llvm_lib = (r"'C:\\Program Files\\LLVM\\bin\\LLVM-C.dll' if WIN else '/opt/homeb
repr(['LLVM'] + [f'LLVM-{i}' for i in reversed(range(14, 21+1))]))
webgpu_lib = "os.path.join(sysconfig.get_paths()['purelib'], 'pydawn', 'lib', 'libwebgpu_dawn.dll') if WIN else 'webgpu_dawn'"
nv_lib_path = "f'/usr/local/cuda/targets/{sysconfig.get_config_var(\"MULTIARCH\").rsplit(\"-\", 1)[0]}/lib'"
nv_lib_path = "f'/usr/local/cuda/targets/{sysconfig.get_config_vars().get(\"MULTIARCH\", \"\").rsplit(\"-\", 1)[0]}/lib'"
def load(name, dll, files, **kwargs):
if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
+1 -1
View File
@@ -2,7 +2,7 @@
import ctypes
from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
import sysconfig
dll = DLL('nvjitlink', 'nvJitLink', f'/usr/local/cuda/targets/{sysconfig.get_config_var("MULTIARCH").rsplit("-", 1)[0]}/lib')
dll = DLL('nvjitlink', 'nvJitLink', f'/usr/local/cuda/targets/{sysconfig.get_config_vars().get("MULTIARCH", "").rsplit("-", 1)[0]}/lib')
nvJitLinkResult = CEnum(ctypes.c_uint32)
NVJITLINK_SUCCESS = nvJitLinkResult.define('NVJITLINK_SUCCESS', 0)
NVJITLINK_ERROR_UNRECOGNIZED_OPTION = nvJitLinkResult.define('NVJITLINK_ERROR_UNRECOGNIZED_OPTION', 1)
+1 -1
View File
@@ -2,7 +2,7 @@
import ctypes
from tinygrad.runtime.support.c import DLL, Struct, CEnum, _IO, _IOW, _IOR, _IOWR
import sysconfig
dll = DLL('nvrtc', 'nvrtc', f'/usr/local/cuda/targets/{sysconfig.get_config_var("MULTIARCH").rsplit("-", 1)[0]}/lib')
dll = DLL('nvrtc', 'nvrtc', f'/usr/local/cuda/targets/{sysconfig.get_config_vars().get("MULTIARCH", "").rsplit("-", 1)[0]}/lib')
nvrtcResult = CEnum(ctypes.c_uint32)
NVRTC_SUCCESS = nvrtcResult.define('NVRTC_SUCCESS', 0)
NVRTC_ERROR_OUT_OF_MEMORY = nvrtcResult.define('NVRTC_ERROR_OUT_OF_MEMORY', 1)
+14 -9
View File
@@ -1,6 +1,6 @@
import collections, time
import collections, itertools, time
from typing import Any, cast
from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv, dedup, suppress_finalizing
from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, dedup, suppress_finalizing
from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator, MMIOInterface
from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
from tinygrad.dtype import dtypes
@@ -22,7 +22,7 @@ class HCQGraph(MultiGraphRunner):
for (j,i), input_idx in self.input_replace.items():
x = self.input_replace_to_var.setdefault((j,i), UOp.variable(f"input_{input_idx}", 0, 0xffffffffffffffff, dtype=dtypes.uint64))
self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size, texture_info=self.hcq_bufs[j][i].texture_info) # Create fake buffer with variable
self.hcq_bufs[j][i] = HCQBuffer(x, self.hcq_bufs[j][i].size, image=self.hcq_bufs[j][i].image) # Create fake buffer with variable
# Allocate kernel args.
kernargs_size: dict[Compiled, int] = collections.defaultdict(int)
@@ -49,7 +49,9 @@ class HCQGraph(MultiGraphRunner):
self.ji_schedule: dict[int, tuple[HCQCompiled, HWQueue, list, list, HCQSignal, int|None]] = {}
self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation
self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", 2 if ALL2ALL >= 1 else 1)
self.copy_queue_cnt: collections.defaultdict[HCQCompiled, itertools.count] = collections.defaultdict(itertools.count)
self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if not dev._is_cpu()},
**{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev._is_cpu()}}
@@ -85,7 +87,8 @@ class HCQGraph(MultiGraphRunner):
enqueue_queue = self.comp_queues[enqueue_dev]
else:
assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"
enqueue_queue = self.copy_queues.setdefault(enqueue_dev, enqueue_dev.hw_copy_queue_t())
queue_idx = next(self.copy_queue_cnt[enqueue_dev]) % self.num_copy_queues
enqueue_queue = self.copy_queues.setdefault((enqueue_dev, queue_idx), enqueue_dev.hw_copy_queue_t(queue_idx=queue_idx))
out_signal = self.signals.setdefault(enqueue_queue, self.devices[0].new_signal(value=0))
@@ -175,14 +178,17 @@ class HCQGraph(MultiGraphRunner):
for dev in self.devices:
for dep_dev in list(self.copy_to_devs[dev]) + [dev]:
if dep_dev in self.copy_queues: self.comp_queues[dev].wait(self.signals[(copy_q:=self.copy_queues[dep_dev])], cast(int, last_j[copy_q]) + 1)
for copy_q in self._dev_copy_queues(dep_dev):
if copy_q in self.signals: self.comp_queues[dev].wait(self.signals[copy_q], cast(int, last_j[copy_q]) + 1)
self.comp_queues[dev].signal(self.virt_timeline_signals[dev], self.virt_timeline_vals[dev] + 1).bind(dev)
if dev in self.copy_queues: self.copy_queues[dev].bind(dev)
for copy_q in self._dev_copy_queues(dev): copy_q.bind(dev)
self.last_timeline: dict[HCQCompiled, tuple[HCQSignal, int]] = {dev: (dev.timeline_signal, 0) for dev in self.devices}
self.queue_signals_to_reset = [self.signals[q] for q in list(self.comp_queues.values()) + list(self.copy_queues.values()) if q in self.signals]
def _dev_copy_queues(self, dev): return [q for (d, _), q in self.copy_queues.items() if d == dev]
def __call__(self, input_rawbuffers: list[Buffer], var_vals: dict[str, int], wait=False) -> float|None:
# Wait and restore signals
self.kickoff_value += 1
@@ -205,8 +211,7 @@ class HCQGraph(MultiGraphRunner):
for dev in self.devices:
self.comp_queues[dev].submit(dev, hcq_var_vals_local:=hcq_var_vals|self.fixedvars.get(dev, {}))
if (copy_queue:=self.copy_queues.get(dev, None)) is not None: copy_queue.submit(dev, hcq_var_vals_local)
for copy_queue in self._dev_copy_queues(dev): copy_queue.submit(dev, hcq_var_vals_local)
self.last_timeline[dev] = (dev.timeline_signal, dev.next_timeline())
if wait:
+41 -34
View File
@@ -446,8 +446,8 @@ class AMDComputeAQLQueue(AMDComputeQueue):
dev.compute_queue.signal_doorbell(dev, doorbell_value=dev.compute_queue.put_value-1)
class AMDCopyQueue(HWQueue):
def __init__(self, dev, max_copy_size=0x40000000):
self.dev, self.sdma, self.internal_cmd_sizes, self.max_copy_size = dev, dev.sdma, [], max_copy_size
def __init__(self, dev, max_copy_size=0x40000000, queue_idx=0):
self.dev, self.sdma, self.internal_cmd_sizes, self.max_copy_size, self.queue_idx = dev, dev.sdma, [], max_copy_size, queue_idx
super().__init__()
def q(self, *arr):
@@ -501,41 +501,42 @@ class AMDCopyQueue(HWQueue):
self._q, self.cmd_sizes = hw_view, [len(self.indirect_cmd)]
def _submit(self, dev:AMDDevice):
sdma_queue = dev.sdma_queue(self.queue_idx)
if self.binded_device == dev:
# An IB packet must end on a 8 DW boundary.
add = (8 - (((dev.sdma_queue.put_value % 32) // 4) + len(self.indirect_cmd) % 8)) % 8
add = (8 - (((sdma_queue.put_value % 32) // 4) + len(self.indirect_cmd) % 8)) % 8
cmds, cmd_sizes = ([0] * add) + self.indirect_cmd, [len(self.indirect_cmd) + add]
if len(cmds) * 4 >= (dev.sdma_queue.ring.nbytes - dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes):
if len(cmds) * 4 >= (sdma_queue.ring.nbytes - sdma_queue.put_value % sdma_queue.ring.nbytes):
cmds, cmd_sizes = [0, 0] + self.indirect_cmd, [8]
else: cmds, cmd_sizes = self._q, self.internal_cmd_sizes
tail_blit_dword = 0
for cmdsz in cmd_sizes:
if (tail_blit_dword + cmdsz) * 4 >= dev.sdma_queue.ring.nbytes - dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes: break
if (tail_blit_dword + cmdsz) * 4 >= sdma_queue.ring.nbytes - sdma_queue.put_value % sdma_queue.ring.nbytes: break
tail_blit_dword += cmdsz
# Force align of submits to hit our usb layer write cache.
if (rem_packet_cnt := len(cmds) - tail_blit_dword) > 0 and dev.is_usb(): tail_blit_dword = 0
# USB devices run in single-step mode, so they can't overrun the queue.
total_bytes = (tail_blit_dword * 4 if rem_packet_cnt == 0 else -dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes) + rem_packet_cnt * 4
assert total_bytes < dev.sdma_queue.ring.nbytes, "SDMA queue overrun"
while not dev.is_usb() and dev.sdma_queue.put_value + total_bytes - dev.sdma_queue.read_ptr > dev.sdma_queue.ring.nbytes: pass
total_bytes = (tail_blit_dword * 4 if rem_packet_cnt == 0 else -sdma_queue.put_value % sdma_queue.ring.nbytes) + rem_packet_cnt * 4
assert total_bytes < sdma_queue.ring.nbytes, "SDMA queue overrun"
while not dev.is_usb() and sdma_queue.put_value + total_bytes - sdma_queue.read_ptr > sdma_queue.ring.nbytes: pass
start_idx = (dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes) // 4
dev.sdma_queue.ring[start_idx : start_idx + tail_blit_dword] = array.array('I', cmds[:tail_blit_dword])
dev.sdma_queue.put_value += tail_blit_dword * 4
start_idx = (sdma_queue.put_value % sdma_queue.ring.nbytes) // 4
sdma_queue.ring[start_idx : start_idx + tail_blit_dword] = array.array('I', cmds[:tail_blit_dword])
sdma_queue.put_value += tail_blit_dword * 4
if (rem_packet_cnt := len(cmds) - tail_blit_dword) > 0:
zero_fill = dev.sdma_queue.ring.nbytes - dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes
dev.sdma_queue.ring.view(dev.sdma_queue.put_value % dev.sdma_queue.ring.nbytes, zero_fill, fmt='B')[:] = bytes(zero_fill)
dev.sdma_queue.put_value += zero_fill
zero_fill = sdma_queue.ring.nbytes - sdma_queue.put_value % sdma_queue.ring.nbytes
sdma_queue.ring.view(sdma_queue.put_value % sdma_queue.ring.nbytes, zero_fill, fmt='B')[:] = bytes(zero_fill)
sdma_queue.put_value += zero_fill
dev.sdma_queue.ring[0:rem_packet_cnt] = array.array('I', cmds[tail_blit_dword:])
dev.sdma_queue.put_value += rem_packet_cnt * 4
sdma_queue.ring[0:rem_packet_cnt] = array.array('I', cmds[tail_blit_dword:])
sdma_queue.put_value += rem_packet_cnt * 4
dev.sdma_queue.signal_doorbell(dev)
sdma_queue.signal_doorbell(dev)
class AMDProgram(HCQProgram):
def __init__(self, dev:AMDDevice, name:str, lib:bytes):
@@ -641,14 +642,14 @@ class AMDQueueDesc:
def read_ptr(self): return min(p[0] for p in self.read_ptrs)
def signal_doorbell(self, dev, doorbell_value:int|None=None):
for write_ptr in self.write_ptrs: write_ptr[0] = self.put_value
# Ensure all prior writes are visible to the GPU.
System.memory_barrier()
# Flush hdp if queue is in dev mem.
if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
try:
for write_ptr in self.write_ptrs: write_ptr[0] = self.put_value
# Ensure all prior writes are visible to the GPU.
System.memory_barrier()
# Flush hdp if queue is in dev mem.
if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
except Exception as e:
dev.error_state = e
@@ -756,7 +757,8 @@ class KFDIface:
stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
assert stm.n_success == 1
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
xcc_id=0, idx=0):
queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
@@ -826,12 +828,14 @@ class PCIIface(PCIIfaceBase):
'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
xcc_id=0, idx=0):
assert cwsr_buffer is None, "no cwsr buffer for am"
if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
assert idx <= 3, "only 4 SDMA queues supported in am"
pv = self.dev_impl.sdma.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0), pipe=0, queue=0)
doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0 + idx * 0xA * 4), pipe=0, queue=idx)
else:
pv = self.dev_impl.gfx.setup_ring(ring_addr=ring.va_addr, ring_size=ring.size, rptr_addr=gart.va_addr+rptr, wptr_addr=gart.va_addr+wptr,
eop_addr=eop_buffer.va_addr, eop_size=eop_buffer.size, doorbell=(doorbell_index:=am.AMDGPU_NAVI10_DOORBELL_MEC_RING0), pipe=0,
@@ -875,9 +879,10 @@ class USBIface(PCIIface):
barview = self.pci_dev.map_bar(bar=0, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
return HCQBuffer(mapping.va_addr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=False), view=barview, owner=self.dev)
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0, xcc_id=0):
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
xcc_id=0, idx=0):
if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.pci_dev.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id)
return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id, idx)
def sleep(self, timeout): pass
@@ -931,8 +936,7 @@ class AMDDevice(HCQCompiled):
0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size, debug_memory_size=debug_memory_size)
max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
@@ -940,7 +944,7 @@ class AMDDevice(HCQCompiled):
super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
functools.partial(AMDCopyQueue, self, max_copy_size=max_copy_size),
functools.partial(AMDCopyQueue, self, max_copy_size=self.max_copy_size),
kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000)
# Scratch setup
@@ -976,7 +980,7 @@ class AMDDevice(HCQCompiled):
self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
self.sqtt_next_cmd_id = itertools.count(0)
def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0):
def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
gart = self.iface.alloc(0x100, uncached=True, cpu_access=True)
@@ -993,7 +997,10 @@ class AMDDevice(HCQCompiled):
return (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size))
ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
@functools.lru_cache(None)
def sdma_queue(self, idx:int=0): return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
def _ensure_has_local_memory(self, private_segment_size):
if self.max_private_segment_size >= private_segment_size: return
+2 -3
View File
@@ -9,7 +9,6 @@ from tinygrad.renderer.cstyle import ClangJITRenderer
from tinygrad.renderer.llvmir import LLVMRenderer
from tinygrad.renderer.nir import LVPRenderer
from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
from tinygrad.runtime.support.compiler_mesa import LVPCompiler
from tinygrad.runtime.support.elf import jit_loader
from tinygrad.uop.ops import sint
@@ -72,7 +71,7 @@ class CPUProgram(HCQProgram):
except OSError: pass
def __init__(self, dev, name:str, lib:bytes):
LVP = isinstance(dev.compiler, LVPCompiler)
LVP = isinstance(dev.renderer, LVPRenderer)
if sys.platform == "win32": # mypy doesn't understand when WIN is used here
PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
@@ -137,5 +136,5 @@ class CPUDevice(HCQCompiled):
self.tasks:queue.Queue = queue.Queue()
CPUWorker(self, self.tasks, thread_id=0).start()
compilers = CompilerSet([CompilerPair(ClangJITRenderer, None), CompilerPair(LLVMRenderer, CPULLVMCompiler, ctrl_var=CPU_LLVM),
CompilerPair(LVPRenderer, LVPCompiler, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
CompilerPair(LVPRenderer, None, ctrl_var=CPU_LVP)], ctrl_var=CPU_CC)
super().__init__(device, CPUAllocator(self), compilers, functools.partial(CPUProgram, self), CPUSignal, CPUComputeQueue)
+2 -4
View File
@@ -6,7 +6,6 @@ from tinygrad.renderer.llvmir import AMDLLVMRenderer
from tinygrad.uop.ops import Ops
from tinygrad.helpers import cpu_profile, EMULATE, NULL_IR3, NULL_NAK
from tinygrad.renderer.nir import IR3Renderer, NAKRenderer
from tinygrad.runtime.support.compiler_mesa import IR3Compiler, NAKCompiler
class NullRenderer(CStyleLanguage):
device = "NULL"
@@ -39,7 +38,6 @@ class NullDevice(Compiled):
case "AMD_RDNA4": renderer = functools.partial(AMDLLVMRenderer, "gfx1201")
case "": renderer = NullRenderer
case _: raise RuntimeError(f"can't EMULATE device: {EMULATE.value}")
compilers = CompilerSet([CompilerPair(renderer, Compiler),
CompilerPair(functools.partial(IR3Renderer, self), functools.partial(IR3Compiler, 0x6030001), NULL_IR3), # adreno 630
CompilerPair(functools.partial(NAKRenderer, self), functools.partial(NAKCompiler, "sm_120", 48), NULL_NAK)]) # 5090
compilers = CompilerSet([CompilerPair(renderer, Compiler), CompilerPair(functools.partial(IR3Renderer, 0x6030001), None, NULL_IR3), # adreno 630
CompilerPair(functools.partial(NAKRenderer, "sm_120", 48), None, NULL_NAK)]) # 5090
super().__init__(device, NullAllocator(self), compilers, functools.partial(NullProgram, device), NullGraph)
+6 -3
View File
@@ -11,7 +11,6 @@ from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, pr
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import NVRenderer
from tinygrad.runtime.support.compiler_cuda import CUDACompiler, PTXCompiler, NVPTXCompiler, NVCompiler
from tinygrad.runtime.support.compiler_mesa import NAKCompiler
from tinygrad.runtime.autogen import nv_570, nv_580, pci, mesa
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager
@@ -166,6 +165,10 @@ class NVComputeQueue(NVCommandQueue):
def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.compute_gpfifo)
class NVCopyQueue(NVCommandQueue):
def __init__(self, queue_idx=0):
self.queue_idx = queue_idx
super().__init__()
def copy(self, dest:sint, src:sint, copy_size:int):
for off in range(0, copy_size, step:=(1 << 31)):
self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src+off), *data64(dest+off))
@@ -216,7 +219,7 @@ class NVProgram(HCQProgram):
self.dev, self.name, self.lib = dev, name, lib
self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[va_addr, size]]
if (NAK:=isinstance(dev.compiler, NAKCompiler)):
if (NAK:=isinstance(dev.renderer, NAKRenderer)):
image, self.cbuf_0 = memoryview(bytearray(lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(lib)):])), []
self.regs_usage, self.shmem_usage, self.lcmem_usage = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
elif MOCKGPU: image, sections, relocs = memoryview(bytearray(lib) + b'\x00' * (4 - len(lib)%4)).cast("I"), [], [] # type: ignore
@@ -586,7 +589,7 @@ class NVDevice(HCQCompiled[HCQSignal]):
cucc, ptxcc = (CUDACompiler, PTXCompiler) if MOCKGPU else (NVCompiler, NVPTXCompiler)
compilers = CompilerSet(ctrl_var=NV_CC, cset=[CompilerPair(functools.partial(NVRenderer, self.arch),functools.partial(cucc, self.arch)),
CompilerPair(functools.partial(PTXRenderer, self.arch, device="NV"), functools.partial(ptxcc, self.arch), NV_PTX),
CompilerPair(functools.partial(NAKRenderer, dev=self), functools.partial(NAKCompiler, self.arch, self.max_warps_per_sm), NV_NAK)])
CompilerPair(functools.partial(NAKRenderer, self.arch, self.max_warps_per_sm), None, NV_NAK)])
super().__init__(device, NVAllocator(self), compilers, functools.partial(NVProgram, self), HCQSignal, NVComputeQueue, NVCopyQueue)
self._setup_gpfifos()
+42 -72
View File
@@ -1,7 +1,6 @@
from __future__ import annotations
import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib
assert sys.platform != 'win32'
from types import SimpleNamespace
from typing import Any, cast
from tinygrad.device import BufferSpec, CompilerSet, CompilerPair
from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator
@@ -10,9 +9,8 @@ from tinygrad.runtime.autogen import kgsl, mesa
from tinygrad.runtime.ops_cl import CLCompiler, CLDevice
from tinygrad.renderer.cstyle import QCOMRenderer
from tinygrad.renderer.nir import IR3Renderer
from tinygrad.runtime.support.compiler_mesa import IR3Compiler
from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE, suppress_finalizing
from tinygrad.helpers import flatten, QCOM_IR3, QCOM_CC
from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC
from tinygrad.runtime.support.system import System
if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import
@@ -27,7 +25,7 @@ def _qreg_exec(__reg, __val=0, **kwargs):
return __val
qreg: Any = type("QREG", (object,), {name[4:].lower(): functools.partial(_qreg_exec, name) for name in mesa.__dict__.keys() if name[:4] == 'REG_'})
def next_power2(x): return 1 if x == 0 else 1 << (x - 1).bit_length()
def ctz(v): return (v & -v).bit_length() - 1
def parity(val: int):
for i in range(4,1,-1): val ^= val >> (1 << i)
@@ -192,45 +190,37 @@ class QCOMComputeQueue(HWQueue):
class QCOMArgsState(HCQArgsState):
def __init__(self, buf:HCQBuffer, prg:QCOMProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
super().__init__(buf, prg, bufs, vals=vals)
if len(bufs) + len(vals) != len(prg.buf_info): raise RuntimeError(f'incorrect args size given={len(bufs)+len(vals)} != want={len(prg.buf_info)}')
self.buf_info, self.args_info = prg.buf_info[:len(bufs)], prg.buf_info[len(bufs):]
ctypes.memset(cast(int, self.buf.va_addr), 0, prg.kernargs_alloc_size)
ubos, uavs = [b for b in bufs if b.image is None], [b for b in bufs if b.image is not None]
ibos, texs = (uavs, []) if prg.tex_cnt == 0 else (uavs[:-prg.tex_cnt], uavs[-prg.tex_cnt:])
for cnst_val,cnst_off,cnst_sz in prg.consts_info: to_mv(self.buf.va_addr + cnst_off, cnst_sz)[:] = cnst_val.to_bytes(cnst_sz, byteorder='little')
if prg.samp_cnt > 0: to_mv(self.buf.va_addr + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers)
for i, b in enumerate(bufs):
if prg.buf_info[i].type in {BUFTYPE_TEX, BUFTYPE_IBO}:
obj = b.texture_info.desc if prg.buf_info[i].type is BUFTYPE_TEX else b.texture_info.ibo
to_mv(self.buf.va_addr + prg.buf_info[i].offset, len(obj) * 4).cast('I')[:] = array.array('I', obj)
self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=self.buf_info[i].offset+(0 if self.buf_info[i].type is BUFTYPE_BUF else 16))
if prg.NIR:
self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off)
self.bind_sints_to_buf(*vals, buf=self.buf, fmt='I', offset=prg.buf_off + len(ubos) * 8)
else:
for i, b in enumerate(ubos): self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=prg.buf_offs[i])
for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=prg.buf_offs[i+len(ubos)])
for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=self.args_info[i].offset)
def _tex(b, ibo=False):
fmt = mesa.FMT6_32_32_32_32_FLOAT if b.image.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
return [qreg.a6xx_tex_const_0(fmt=fmt) if ibo else qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=fmt),
qreg.a6xx_tex_const_1(width=b.image.shape[1], height=b.image.shape[0]),
qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=b.image.pitch, pitchalign=ctz(b.image.pitch)-6), 0, *data64_le(b.va_addr),
qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13), 0, 0, 0, 0, 0, 0, 0, 0]
class IR3ArgsState(HCQArgsState):
def __init__(self, buf:HCQBuffer, prg:QCOMProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
super().__init__(buf, prg, bufs, vals=vals)
ctypes.memset(cast(int, self.buf.va_addr), 0, prg.kernargs_alloc_size)
to_mv(self.buf.va_addr + prg.imm_off, len(prg.imm_vals))[:] = prg.imm_vals
ubos, uavs = [b for b in bufs if b.texture_info is None], [b for b in bufs if b.texture_info is not None]
ibos, texs = (uavs, []) if prg.tex_cnt == 0 else (uavs[:-prg.tex_cnt], uavs[-prg.tex_cnt:]) # textures are at the end
if prg.samp_cnt > 0: to_mv(self.buf.va_addr + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers)
self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off)
self.bind_sints_to_buf(*vals, buf=self.buf, fmt='I', offset=prg.buf_off + len(ubos) * 8)
self.bind_sints_to_buf(*flatten([b.texture_info.desc + ([0] * 8) for b in texs]), buf=self.buf, fmt='I', offset=prg.tex_off)
self.bind_sints_to_buf(*flatten([b.texture_info.ibo + ([0] * 8) for b in ibos]), buf=self.buf, fmt='I', offset=prg.ibo_off)
self.bind_sints_to_buf(*flatten(map(_tex, texs)), buf=self.buf, fmt='I', offset=prg.tex_off)
self.bind_sints_to_buf(*flatten(map(functools.partial(_tex, ibo=True), ibos)), buf=self.buf, fmt='I', offset=prg.ibo_off)
class QCOMProgram(HCQProgram):
def __init__(self, dev: QCOMDevice, name: str, lib: bytes):
self.dev: QCOMDevice = dev
self.name, self.lib, self.NIR = name, lib, isinstance(dev.compiler, IR3Compiler)
self.name, self.lib, self.NIR = name, lib, isinstance(dev.renderer, IR3Renderer)
if self.NIR:
from tinygrad.runtime.autogen import mesa
from tinygrad.runtime.support.compiler_mesa import IR3Compiler
v, cs, self.imm_vals, self.image = IR3Compiler.unpack_lib(lib)
self.prg_offset, self.brnchstck, self.image_size, self.pvtmem, self.shmem = 0, v.branchstack, v.info.size, v.pvtmem_size, v.shared_size
self.wgsz = alloc.offset_vec4 * 4 + 8 if (alloc:=cs.allocs.consts[mesa.IR3_CONST_ALLOC_DRIVER_PARAMS]).size_vec4 else 0xfc
@@ -247,10 +237,11 @@ class QCOMProgram(HCQProgram):
self.tex_off, self.ibo_off, self.samp_off = 2048, 2048 + 0x40 * self.tex_cnt, 2048 + 0x40 * (self.tex_cnt + self.ibo_cnt)
self.fregs, self.hregs = v.info.max_reg + 1, v.info.max_half_reg + 1
self.consts_info:list[tuple] = []
else: self._parse_lib()
self.lib_gpu: HCQBuffer = self.dev.allocator.alloc(self.image_size, buf_spec:=BufferSpec(cpu_access=True, nolru=True))
to_mv(cast(int, self.lib_gpu.va_addr), self.image_size)[:] = self.image
to_mv(self.lib_gpu.va_addr, self.image_size)[:] = self.image
self.pvtmem_size_per_item: int = round_up(self.pvtmem, 512) >> 9
self.pvtmem_size_total: int = self.pvtmem_size_per_item * 128 * 2
@@ -260,7 +251,7 @@ class QCOMProgram(HCQProgram):
dev._ensure_stack_size(self.hw_stack_offset * 4)
kernargs_alloc_size = round_up(2048 + (self.tex_cnt + self.ibo_cnt) * 0x40 + len(self.samplers) * 4, 0x100)
super().__init__(IR3ArgsState if self.NIR else QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
super().__init__(QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
@@ -280,7 +271,7 @@ class QCOMProgram(HCQProgram):
self.pvtmem, self.shmem = _read_lib(self.lib, image_desc_off+0xc8), _read_lib(self.lib, image_desc_off+0xd8)
# Fill up constants and buffers info
self.buf_info, self.consts_info = [], []
self.consts_info = []
# Collect sampler info.
self.samp_cnt = samp_cnt_in_file = _read_lib(self.lib, image_desc_off + 0xdc)
@@ -292,20 +283,17 @@ class QCOMProgram(HCQProgram):
else: self.samplers = []
# Collect kernel arguments (buffers) info.
bdoff = round_up(image_desc_off + 0x158 + len(self.name), 4) + 8 * samp_cnt_in_file
bdoff, binfos = round_up(image_desc_off + 0x158 + len(self.name), 4) + 8 * samp_cnt_in_file, []
while bdoff + 32 <= len(self.lib):
length, _, _, offset_words, _, _, _, typ = struct.unpack("IIIIIIII", self.lib[bdoff:bdoff+32])
length, _, _, offset_words, _, _, _, typ = struct.unpack("8I", self.lib[bdoff:bdoff+32])
if length == 0: break
self.buf_info.append(SimpleNamespace(offset=offset_words * 4, type=typ))
binfos.append((offset_words * 4, typ))
bdoff += length
self.buf_offs = [off for off,typ in binfos if typ not in {BUFTYPE_TEX, BUFTYPE_IBO}]
# Setting correct offsets to textures/ibos.
self.tex_cnt, self.ibo_cnt = sum(x.type is BUFTYPE_TEX for x in self.buf_info), sum(x.type is BUFTYPE_IBO for x in self.buf_info)
self.tex_cnt, self.ibo_cnt = sum(typ is BUFTYPE_TEX for _,typ in binfos), sum(typ is BUFTYPE_IBO for _,typ in binfos)
self.ibo_off, self.tex_off, self.samp_off = 2048, 2048 + 0x40 * self.ibo_cnt, 2048 + 0x40 * self.tex_cnt + 0x40 * self.ibo_cnt
cur_ibo_off, cur_tex_off = self.ibo_off, self.tex_off
for x in self.buf_info:
if x.type is BUFTYPE_IBO: x.offset, cur_ibo_off = cur_ibo_off, cur_ibo_off + 0x40
elif x.type is BUFTYPE_TEX: x.offset, cur_tex_off = cur_tex_off, cur_tex_off + 0x40
if _read_lib(self.lib, 0xb0) != 0: # check if we have constants.
cdoff = _read_lib(self.lib, 0xac)
@@ -323,28 +311,10 @@ class QCOMTextureInfo:
self.pitch, self.real_stride, self.desc, self.ibo = pitch, real_stride, desc, ibo
class QCOMAllocator(HCQAllocatorBase):
def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer:
# Recalculate real size for texture
if options.image is not None:
imgw, imgh, itemsize_log = options.image.shape[1], options.image.shape[0], int(math.log2(options.image.itemsize))
pitchalign = max(6, 11 - int(math.log2(imgh))) if imgh > 1 else 6
align_up = max(1, (8 // itemsize_log + 1) - imgh // 32) if pitchalign == 6 else (2 ** (pitchalign - itemsize_log - 2))
granularity = 128 if options.image.itemsize == 4 else 256
pitch_add = (1 << pitchalign) if min(next_power2(imgw), round_up(imgw, granularity)) - align_up + 1 <= imgw and imgw > granularity//2 else 0
pitch = round_up((real_stride:=imgw * 4 * options.image.itemsize), 1 << pitchalign) + pitch_add
size = pitch * imgh
buf = self.dev._gpu_map(options.external_ptr, size) if options.external_ptr else self.dev._gpu_alloc(size)
if options.image is not None:
tex_fmt = mesa.FMT6_32_32_32_32_FLOAT if options.image.itemsize == 4 else mesa.FMT6_16_16_16_16_FLOAT
desc = [qreg.a6xx_tex_const_0(0x8, swiz_x=0, swiz_y=1, swiz_z=2, swiz_w=3, fmt=tex_fmt), qreg.a6xx_tex_const_1(width=imgw, height=imgh),
qreg.a6xx_tex_const_2(type=mesa.A6XX_TEX_2D, pitch=pitch, pitchalign=pitchalign-6), 0,
*data64_le(buf.va_addr), qreg.a6xx_tex_const_6(plane_pitch=0x400000), qreg.a6xx_tex_const_7(13)]
buf.texture_info = QCOMTextureInfo(pitch, real_stride, desc, [desc[0] & (~0xffff), *desc[1:len(desc)]])
return buf
if opts.image is not None: size = opts.image.pitch* opts.image.shape[0]
return self.dev._gpu_map(opts.external_ptr, size, image=opts.image) if opts.external_ptr else self.dev._gpu_alloc(size, image=opts.image)
def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, prof_text, dest_off=0, src_off=0):
with cpu_profile(prof_text, self.dev.device, is_copy=True):
@@ -353,13 +323,13 @@ class QCOMAllocator(HCQAllocatorBase):
src_off, dest_off = src_off+src_stride, dest_off+dest_stride
def _copyin(self, dest:HCQBuffer, src:memoryview):
stride, pitch = (src.nbytes, src.nbytes) if (ti:=cast(QCOMTextureInfo, dest.texture_info)) is None else (ti.real_stride, ti.pitch)
stride, pitch = (dest.image.shape[1] * 4 * dest.image.itemsize, dest.image.pitch) if dest.image else (src.nbytes, src.nbytes)
self._do_copy(mv_address(src), dest.cpu_view().addr, src.nbytes, stride, stride, pitch, f"TINY -> {self.dev.device}")
def _copyout(self, dest:memoryview, src:HCQBuffer):
self.dev.synchronize()
stride, pitch = (src.size, src.size) if (ti:=cast(QCOMTextureInfo, src.texture_info)) is None else (ti.real_stride, ti.pitch)
stride, pitch = (src.image.shape[1] * 4 * src.image.itemsize, src.image.pitch) if src.image else (src.size, src.size)
self._do_copy(src.cpu_view().addr, mv_address(dest), src.size, stride, pitch, stride, f"{self.dev.device} -> TINY")
def _as_buffer(self, src:HCQBuffer) -> memoryview:
@@ -402,11 +372,11 @@ class QCOMDevice(HCQCompiled):
System.write_sysfs("/sys/class/kgsl/kgsl-3d0/idle_timer", value="4000000000", msg="Failed to disable suspend mode", expected="4294967276")
compilers = CompilerSet(ctrl_var=QCOM_CC, cset=[CompilerPair(QCOMRenderer, functools.partial(QCOMCompiler, device)),
CompilerPair(functools.partial(IR3Renderer, self), functools.partial(IR3Compiler, info.chip_id), QCOM_IR3)])
CompilerPair(functools.partial(IR3Renderer, info.chip_id), None, QCOM_IR3)])
super().__init__(device, QCOMAllocator(self), compilers, functools.partial(QCOMProgram, self), QCOMSignal,
functools.partial(QCOMComputeQueue, self), None)
def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False, **kwargs) -> HCQBuffer:
flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)
@@ -414,15 +384,15 @@ class QCOMDevice(HCQCompiled):
va_addr = self.fd.mmap(0, bosz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, alloc.id * 0x1000)
if fill_zeroes: ctypes.memset(va_addr, 0, size)
return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self)
return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self, **kwargs)
def _gpu_map(self, ptr:int, size:int) -> HCQBuffer:
def _gpu_map(self, ptr:int, size:int, **kwargs) -> HCQBuffer:
ptr_aligned, size_aligned = (ptr & ~0xfff), round_up(size + (ptr & 0xfff), 0x1000)
try:
mapinfo = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
return HCQBuffer(mapinfo.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mapinfo, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
mi = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)
except OSError as e:
if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs)
raise RuntimeError("Failed to map external pointer to GPU memory") from e
def _gpu_free(self, mem:HCQBuffer):
+22 -18
View File
@@ -189,16 +189,18 @@ class AM_SMU(AM_IP):
return table_t.from_buffer(bytearray(self.adev.vram.view(self.driver_table_paddr, ctypes.sizeof(table_t))[:]))
def set_clocks(self, level):
if self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12)}: return # TODO
if not hasattr(self, 'clcks'):
clks = [self.smu_mod.PPCLK_UCLK, self.smu_mod.PPCLK_FCLK, self.smu_mod.PPCLK_SOCCLK]
if self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,12)}: clks.append(self.smu_mod.PPCLK_GFXCLK)
self.clcks = {}
for clck in [self.smu_mod.PPCLK_GFXCLK, self.smu_mod.PPCLK_UCLK, self.smu_mod.PPCLK_FCLK, self.smu_mod.PPCLK_SOCCLK]:
for clck in clks:
cnt = self._send_msg(self.smu_mod.PPSMC_MSG_GetDpmFreqByIndex, (clck<<16)|0xff, read_back_arg=True)&0x7fffffff
self.clcks[clck] = [self._send_msg(self.smu_mod.PPSMC_MSG_GetDpmFreqByIndex, (clck<<16)|i, read_back_arg=True)&0x7fffffff for i in range(cnt)]
for clck, vals in self.clcks.items():
self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMinByFreq, clck << 16 | (vals[level]))
if not vals: continue
with contextlib.suppress(TimeoutError): self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMinByFreq, clck << 16 | (vals[level]), timeout=20)
self._send_msg(self.smu_mod.PPSMC_MSG_SetSoftMaxByFreq, clck << 16 | (vals[level]))
def _smu_cmn_send_msg(self, msg:int, param=0, debug=False):
@@ -412,33 +414,34 @@ class AM_IH(AM_IP):
self.adev.regIH_RB_RPTR.write(wptr % self.ring_size)
class AM_SDMA(AM_IP):
def init_sw(self): self.sdma_name = "F32" if self.adev.ip_ver[am.SDMA0_HWIP] < (7,0,0) else "MCU"
def init_sw(self): self.sdma_reginst, self.sdma_name = [], "F32" if self.adev.ip_ver[am.SDMA0_HWIP] < (7,0,0) else "MCU"
def init_hw(self):
for pipe_id in range(1):
pipe = "" if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else str(pipe_id)
pipe, inst = ("", pipe_id) if self.adev.ip_ver[am.SDMA0_HWIP] < (5,0,0) else (str(pipe_id), 0)
if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100) # 10s, 100ms per unit
self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9)
self.adev.reg(f"regSDMA{pipe}_WATCHDOG_CNTL").update(queue_hang_count=100, inst=inst) # 10s, 100ms per unit
self.adev.reg(f"regSDMA{pipe}_UTCL1_CNTL").update(resp_mode=3, redo_delay=9, inst=inst)
# rd=noa, wr=bypass
self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=2, wr_l2_policy=3, **({'llc_noalloc':1} if self.sdma_name == "F32" else {}))
self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0})
self.adev.reg(f"regSDMA{pipe}_UTCL1_PAGE").update(rd_l2_policy=2, wr_l2_policy=3, **({'llc_noalloc':1} if self.sdma_name == "F32" else {}),
inst=inst)
self.adev.reg(f"regSDMA{pipe}_{self.sdma_name}_CNTL").update(halt=0, **{f"{'th1_' if self.sdma_name == 'F32' else ''}reset":0}, inst=inst)
self.adev.reg(f"regSDMA{pipe}_CNTL").update(ctxempty_int_enable=1, trap_enable=1,
**({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}))
**({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}), inst=inst)
if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
self.adev.regDOORBELL0_CTRL_ENTRY_1.write(bif_doorbell1_range_offset_entry=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2,
bif_doorbell1_range_size_entry=4)
for i in range(16): self.adev.reg(f"regDOORBELL0_CTRL_ENTRY_{i+1}").write(**{f"bif_doorbell{i+1}_range_size_entry":4,
f"bif_doorbell{i+1}_range_offset_entry":(am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0 + i * 0xA) * 2})
self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x1, offset=0xe, size=4)
else: self.adev.soc.doorbell_enable(port=2, awid=0xe, awaddr_31_28_value=0x3, offset=am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0*2, size=4)
def fini_hw(self):
reg, inst = ("regSDMA_GFX", 0) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else ("regSDMA0_QUEUE0", 0)
for reg, inst in self.sdma_reginst:
self.adev.reg(f"{reg}_RB_CNTL").update(rb_enable=0, inst=inst)
self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=0, inst=inst)
self.adev.reg(f"{reg}_RB_CNTL").update(rb_enable=0, inst=inst)
self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=0, inst=inst)
if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
time.sleep(0.01)
@@ -446,7 +449,8 @@ class AM_SDMA(AM_IP):
def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, doorbell:int, pipe:int, queue:int) -> int:
# Setup the ring
reg, inst = ("regSDMA_GFX", pipe*4+queue) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
reg, inst = ("regSDMA_GFX", pipe+queue*4) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
self.sdma_reginst.append((reg, inst))
self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x1, inst=inst)
if not self.adev.partial_boot: self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
@@ -460,7 +464,7 @@ class AM_SDMA(AM_IP):
self.adev.reg(f"{reg}_RB_CNTL").write(**({f'{self.sdma_name.lower()}_wptr_poll_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP][:2]!=(4,4) else {}),
rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4, rb_enable=1, rb_priv=1, rb_size=(ring_size//4).bit_length()-1, inst=inst)
self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=1, inst=inst)
return self.adev.reg(f"{reg}_RB_WPTR").read() | (self.adev.reg(f"{reg}_RB_WPTR_HI").read() << 32)
return self.adev.reg(f"{reg}_RB_WPTR").read(inst=inst) | (self.adev.reg(f"{reg}_RB_WPTR_HI").read(inst=inst) << 32)
class AM_PSP(AM_IP):
def init_sw(self):
+1 -1
View File
@@ -45,7 +45,7 @@ class DLL(ctypes.CDLL):
for p in paths:
libpaths = {"posix": ["/usr/lib", "/usr/local/lib"], "nt": os.environ['PATH'].split(os.pathsep),
"darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{p}.framework"],
'linux': ['/lib', f"/lib/{sysconfig.get_config_var('MULTIARCH')}", "/usr/lib/wsl/lib/"]}
'linux': ['/lib', '/lib64', f"/lib/{sysconfig.get_config_var('MULTIARCH')}", "/usr/lib/wsl/lib/"]}
if (pth:=pathlib.Path(p)).is_absolute():
if pth.is_file(): return p
else: continue
+9 -2
View File
@@ -1,4 +1,4 @@
import ctypes, hashlib, tempfile, subprocess, pathlib
import ctypes, hashlib, tempfile, subprocess, pathlib, shutil
from tinygrad.helpers import system
from tinygrad.runtime.autogen import comgr
try:
@@ -12,8 +12,15 @@ from tinygrad.device import Compiler, CompileError
from tinygrad.runtime.support.compiler_cpu import LLVMCompiler
from tinygrad.helpers import OSX, to_char_p_p
def _find_llvm_objdump():
if OSX: return '/opt/homebrew/opt/llvm/bin/llvm-objdump'
# Try ROCm path first, then versioned, then unversioned
for p in ['/opt/rocm/llvm/bin/llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20', 'llvm-objdump']:
if shutil.which(p): return p
raise FileNotFoundError("llvm-objdump not found")
def amdgpu_disassemble(lib:bytes):
asm = system(f"{'/opt/homebrew/opt/llvm/bin/llvm-objdump' if OSX else '/opt/rocm/llvm/bin/llvm-objdump'} -d -", input=lib).splitlines()
asm = system(f"{_find_llvm_objdump()} -d -", input=lib).splitlines()
while asm and ("s_nop 0" in asm[-1] or "s_code_end" in asm[-1]): asm.pop()
print("\n".join(asm))
+5 -6
View File
@@ -2,8 +2,7 @@ import ctypes, platform, sys, subprocess
from tinygrad.device import Compiler
from tinygrad.helpers import OSX, getenv, capstone_flatdump, DEBUG, unwrap
from tinygrad.runtime.support.elf import jit_loader
try: from tinygrad.runtime.autogen import llvm
except (ImportError, FileNotFoundError): llvm = None #type:ignore[assignment]
from tinygrad.runtime.autogen import llvm
class ClangJITCompiler(Compiler):
def __init__(self, cachekey="compile_clang_jit"): super().__init__(cachekey)
@@ -30,7 +29,7 @@ def expect(x, err, ret=None):
class LLVMCompiler(Compiler):
jit = True
target_arch = {'arm64': 'AArch64', 'aarch64': 'AArch64', 'x86_64': 'X86', 'AMD64': 'X86', 'riscv64': 'riscv64'}[platform.machine()]
def __init__(self, processor:str, feats:str):
def __init__(self, processor:str, feats:str, cache_key=None):
for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter']: getattr(llvm, f'LLVMInitialize{self.target_arch}{component}')()
triple = {'AArch64': b'aarch64-none-unknown-elf', 'X86': b'x86_64-none-unknown-elf', 'AMDGPU': b'amdgcn-amd-amdhsa'}[self.target_arch]
@@ -60,7 +59,7 @@ class LLVMCompiler(Compiler):
self.diag_msgs.append(msg)
self.handle_diag = handle_diag
llvm.LLVMContextSetDiagnosticHandler(self.context, handle_diag, None)
super().__init__(f"compile_llvm_{processor}_{feats}{'_jit' if self.jit else ''}{'_opt' if opt else ''}")
super().__init__(cache_key or f"compile_llvm_{processor}_{feats}{'_jit' if self.jit else ''}{'_opt' if opt else ''}")
def __del__(self):
llvm.LLVMDisposePassBuilderOptions(self.pbo)
@@ -84,7 +83,7 @@ class LLVMCompiler(Compiler):
def disassemble(self, lib:bytes): capstone_flatdump(lib)
class CPULLVMCompiler(LLVMCompiler):
def __init__(self):
def __init__(self, cache_key=None):
# +reserve-x18 here does the same thing as -ffixed-x18 in ops_cpu.py, see comments there for why it's needed on arm osx
cpu, feats = ctypes.string_at(llvm.LLVMGetHostCPUName()), (b'+reserve-x18,' if OSX else b'') + ctypes.string_at(llvm.LLVMGetHostCPUFeatures())
super().__init__(cpu.decode(), feats.decode())
super().__init__(cpu.decode(), feats.decode(), cache_key)
+9 -25
View File
@@ -1,10 +1,10 @@
import base64, ctypes, pathlib, tempfile, hashlib, sys
from tinygrad.device import Compiler
from tinygrad.helpers import cpu_objdump, system, data64
from tinygrad.runtime.autogen import mesa
from tinygrad.runtime.autogen import mesa, llvm
from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, expect, cerr
try: from tinygrad.runtime.autogen import llvm
except (ImportError, FileNotFoundError): llvm = None #type:ignore[assignment]
# NB: compilers assume mesa's glsl type cache is managed externally with mesa.glsl_type_singleton_init_or_ref() and mesa.glsl_type_singleton_decref()
def rzalloc(typ, ctx=None, **kwargs):
s = ctypes.cast(mesa.rzalloc_size(ctypes.cast(ctx, ctypes.c_void_p), ctypes.sizeof(typ)), ctypes.POINTER(typ))
@@ -16,20 +16,8 @@ def deserialize(enc_src, opts):
mesa.blob_reader_init(blobreader, src:=base64.b64decode(enc_src), len(src))
return mesa.nir_deserialize(None, ctypes.cast(opts, ctypes.POINTER(mesa.nir_shader_compiler_options)), blobreader)
class NIRCompiler(Compiler):
def __init__(self, cache_key):
mesa.glsl_type_singleton_init_or_ref()
super().__init__(cache_key)
def __del__(self): mesa.glsl_type_singleton_decref()
class LVPCompiler(CPULLVMCompiler, NIRCompiler):
def __init__(self, cache_key="lvp"):
CPULLVMCompiler.__init__(self)
NIRCompiler.__init__(self, f"compile_{cache_key}")
def __del__(self):
NIRCompiler.__del__(self)
CPULLVMCompiler.__del__(self)
class LVPCompiler(CPULLVMCompiler):
def __init__(self, cache_key="lvp"): CPULLVMCompiler.__init__(self, cache_key=f"compile_{cache_key}")
def compile(self, src) -> bytes:
shader, ctx = deserialize(src, mesa.lvp_nir_options), llvm.LLVMGetGlobalContext()
@@ -62,16 +50,14 @@ class LVPCompiler(CPULLVMCompiler, NIRCompiler):
def disassemble(self, lib: bytes): cpu_objdump(lib)
class NAKCompiler(NIRCompiler):
class NAKCompiler(Compiler):
def __init__(self, arch, warps_per_sm, cache_key="nak"):
self.arch, self.warps_per_sm = arch, warps_per_sm
self.cc = mesa.nak_compiler_create(mesa.struct_nv_device_info(sm=int(arch[3:]), max_warps_per_mp=warps_per_sm))
self.nir_options = bytes(mesa.nak_nir_options(self.cc).contents)
super().__init__(f"compile_{cache_key}_{arch}")
def __del__(self):
mesa.nak_compiler_destroy(self.cc)
super().__del__()
def __del__(self): mesa.nak_compiler_destroy(self.cc)
def __reduce__(self): return NAKCompiler, (self.arch, self.warps_per_sm)
@@ -102,7 +88,7 @@ def disas_adreno(lib:bytes, gpu_id=630):
tf.seek(0)
print(tf.read())
class IR3Compiler(NIRCompiler):
class IR3Compiler(Compiler):
def __init__(self, chip_id, cache_key="ir3"):
assert sys.version_info >= (3,14), "IR3 requires python 3.14's bitfield fixes"
self.dev_id = mesa.struct_fd_dev_id(((chip_id >> 24) & 0xFF) * 100 + ((chip_id >> 16) & 0xFF) * 10 + ((chip_id >> 8) & 0xFF), chip_id)
@@ -112,9 +98,7 @@ class IR3Compiler(NIRCompiler):
self.nir_options = bytes(mesa.ir3_get_compiler_options(self.cc).contents)
super().__init__(f"compile_{cache_key}")
def __del__(self):
mesa.ir3_compiler_destroy(self.cc)
super().__del__()
def __del__(self): mesa.ir3_compiler_destroy(self.cc)
def __reduce__(self): return IR3Compiler, (self.dev_id.chip_id,)
+5 -4
View File
@@ -8,6 +8,7 @@ from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEve
from tinygrad.uop.ops import sym_infer, sint, UOp
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.support.memory import BumpAllocator
from tinygrad.dtype import ImageDType
class MMIOInterface:
def __init__(self, addr:int, nbytes:int, fmt='B'): self.mv, self.addr, self.nbytes, self.fmt = to_mv(addr, nbytes).cast(fmt), addr, nbytes, fmt
@@ -354,7 +355,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
cpu_devices: list[HCQCompiled] = []
def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:CompilerSet, runtime, signal_t:Type[SignalType],
comp_queue_t:Callable[[], HWQueue], copy_queue_t:Callable[[], HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
from tinygrad.runtime.graph.hcq import HCQGraph
@@ -455,14 +456,14 @@ class HCQCompiled(Compiled, Generic[SignalType]):
if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
class HCQBuffer:
def __init__(self, va_addr:sint, size:int, texture_info:Any=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
def __init__(self, va_addr:sint, size:int, image:ImageDType|None=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
owner:HCQCompiled|None=None):
self.va_addr, self.size, self.texture_info, self.meta, self._base, self.view = va_addr, size, texture_info, meta, _base, view
self.va_addr, self.size, self.image, self.meta, self._base, self.view = va_addr, size, image, meta, _base, view
self._devs, self.owner = ([owner] if owner is not None else []), owner
self._mappings:dict[HCQCompiled, HCQBuffer] = {} # mapping to the other devices
def offset(self, offset:int=0, size:int|None=None) -> HCQBuffer:
return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, texture_info=self.texture_info, meta=self.meta,
return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, image=self.image, meta=self.meta,
_base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
def cpu_view(self) -> MMIOInterface:
+1 -1
View File
@@ -176,7 +176,7 @@ class MemoryManager:
self.boot_allocator = TLSFAllocator(boot_size, base=0)
self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0, base=self.boot_allocator.size)
self.pa_allocator = TLSFAllocator(vram_size - (64 << 20), base=self.boot_allocator.size + self.ptable_allocator.size)
self.pa_allocator = TLSFAllocator(vram_size - (off_sz:=self.boot_allocator.size + self.ptable_allocator.size) - (64 << 20), base=off_sz)
self.root_page_table = pt_t(self.dev, self.palloc(0x1000, zero=not self.dev.smi_dev, boot=True), lv=first_lv)
def _frag_size(self, va, sz, must_cover=True):
+31 -27
View File
@@ -1,6 +1,6 @@
from typing import cast
import functools, itertools, operator
from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv
from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, graph_rewrite_map, graph_rewrite
from tinygrad.device import Device
@@ -35,45 +35,49 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
if not isinstance(buf.device, tuple): return None
assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}"
n_lbs, shape, numel = len(buf.device), buf.shape, prod(buf.shape)
# ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
# fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
use_ring = (RING >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
if DEBUG >= 2: print(f"{'RING ALLREDUCE' if use_ring else 'NAIVE ALLREDUCE'} {n_lbs}x{numel} | {buf.dtype}")
use_all2all = (ALL2ALL >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
use_ring = not use_all2all and (RING >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {n_lbs}x{numel} | {buf.dtype}")
# contiguous before we copy it
buf = buf.contiguous()
# copy to all devices. if you shrink later, that'll be handled
if not use_ring: return functools.reduce(lambda x,y: x.alu(red.arg, y),
[UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(len(buf.device))])
# naive: copy to all devices. if you shrink later, that'll be handled
if not use_ring and not use_all2all:
return functools.reduce(lambda x,y: x.alu(red.arg, y), [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(n_lbs)])
# new ring reduce
# chunk data into n_lbs pieces
factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
base, left = (numel // factor) // n_lbs, (numel // factor) % n_lbs
chunk_sizes = [(base + 1) * factor] * left + [base * factor] * (n_lbs - left)
chunks = list(itertools.pairwise(itertools.accumulate(chunk_sizes, initial=0)))
chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (n_lbs - left), initial=0)))
# extract chunks and scatter-reduce
# reduce-scatter
reduced_chunks = []
for i,(s,e) in enumerate(chunks):
chunk = buf.reshape((numel,)).shrink(((s,e),))
reduced_chunk = chunk
for step in range(n_lbs-1):
src, dest = (i+step)%n_lbs, (i+step+1)%n_lbs
# copy the chunk from the src device to the dest (operating device), and select the chunk on the dest device
reduced_chunk = reduced_chunk.copy_to_device(buf.device[dest], src if isinstance(reduced_chunk.device, tuple) else None) \
.alu(red.arg, chunk.copy_to_device(buf.device[dest], dest))
reduced_chunks.append(reduced_chunk)
if use_all2all:
chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(n_lbs)]
reduced_chunks.append(functools.reduce(lambda x,y: x.alu(red.arg, y), chunks_on_i))
else:
chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),))
for step in range(n_lbs-1):
src, dest = (i+step)%n_lbs, (i+step+1)%n_lbs
cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None)
reduced = cp.alu(red.arg, chunk.copy_to_device(buf.device[dest], dest))
reduced_chunks.append(reduced)
# allgather
copied_chunks = []
for i,c in enumerate(reduced_chunks):
this_chunk: list[UOp|None] = [None] * len(buf.device)
this_chunk[(i+len(buf.device)-1)%n_lbs] = c
for step in range(n_lbs-1):
dest = (i+step)%n_lbs
this_chunk[dest] = c = c.copy_to_device(buf.device[dest])
copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(cast(list[UOp], this_chunk))))
for i,rc in enumerate(reduced_chunks):
if use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(n_lbs))))
else:
this_chunk: list[UOp|None] = [None] * n_lbs
this_chunk[(i+n_lbs-1)%n_lbs] = rc
for step in range(n_lbs-1):
this_chunk[(i+step)%n_lbs] = rc = rc.copy_to_device(buf.device[(i+step)%n_lbs])
copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(cast(list[UOp], this_chunk))))
# reassemble
pads = [((s,numel-e),) for s,e in chunks]
@@ -217,8 +221,8 @@ multi_pm = PatternMatcher([
(UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
src=(UPat(Ops.MULTI, name="multi"), ), name="root"), passthrough_multi),
# multi supports custom kernels with CUSTOM_KERNEL + AFTER
(UPat(Ops.CUSTOM_KERNEL, src=UPat(Ops.MULTI), name="ck"),
lambda ck: ck.replace(src=tuple(m.src[0] for m in ck.src))),
(UPat(Ops.CUSTOM_KERNEL, src=UPat((Ops.MULTI, Ops.CONTIGUOUS)), name="ck"),
lambda ck: ck.replace(src=tuple(m.src[0] if m.op is Ops.MULTI else m for m in ck.src))),
(UPat(Ops.AFTER, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.CUSTOM_KERNEL)), name="a"),
lambda multi,a: a.replace(src=(multi.src[0],)+a.src[1:]).multi(multi.axis))
])+replace_allreduce

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