Compare commits

..
171 Commits
Author SHA1 Message Date
chenyuandGitHub 6b0a9f5ee6 don't strip sink in to_uops_list [pr] (#14111) 2026-01-12 11:19:03 -05:00
chenyuandGitHub cad7feec02 more onnx ops (#14104)
HannWindow, HammingWindow, BlackmanWindow, Hardmax, LpNormalization
2026-01-12 09:11:13 -05:00
nimlgenandGitHub 635ed2df9d system: use pci.PCI_VENDOR_ID instead of const (#14109) 2026-01-12 15:24:09 +03:00
qazalandGitHub 6c0f0e29ff Revert "viz: loading... (#14107)" (#14108)
This reverts commit 9347757c2d.
2026-01-12 20:45:37 +09:00
nimlgenandGitHub 9347757c2d viz: loading... (#14107) 2026-01-12 13:24:24 +03:00
wozeparrotandGitHub 3a92df66ea feat: bump version to 0.12.0 (#14105) 2026-01-11 21:19:49 -08:00
chenyuandGitHub 7c234a9c7c wgsl cleanup [pr] (#14103)
refactor common pack functions
2026-01-11 21:23:45 -05:00
George HotzandGitHub 91bde927ef assembly/amd: split asm.py into asm.py and disasm.py (#14101)
* split asm.py into asm.py and disasm.py

* split decoder

* move to pcode

* tests
2026-01-12 07:22:02 +09:00
George HotzandGitHub 44135e2e84 assembly/amd: always use v_nop in test for rocprof-trace-decoder (#14100)
* assembly/amd: always use v_nop in test for rocprof-trace-decoder

* test touchups
2026-01-12 05:31:58 +09:00
George HotzandGitHub 8b1b15aec0 assembly/amd: SQTT support (#14099)
* assembly/amd: SQTT support

* simpler

* cmp wave

* instruction compare

* rocprof decode

* simpler

* no llvm

* no strcmp
2026-01-12 05:07:17 +09:00
nimlgenandGitHub 8b5ff403fa am: flag successful finalization (#14097)
* am: flag successful finalization

* import
2026-01-11 16:24:53 +03:00
qazalandGitHub d8aba24967 amd: use kernel descriptor struct in AMDProgram (#14096) 2026-01-11 18:25:16 +09:00
chenyuandGitHub 9973a81356 add channels_last to QLinearGlobalAveragePool (#14094)
and other minor cleanups
2026-01-10 18:38:19 -05:00
chenyuandGitHub c5492f8f75 cstyle cleanup [pr] (#14093) 2026-01-10 09:44:50 -05:00
nimlgenandGitHub d5f954858d viz: show precise timings (#14092) 2026-01-10 16:21:08 +03:00
nimlgenandGitHub 3e2c05ee9f hevc: decoder as iterator (#14091) 2026-01-10 14:57:56 +03:00
chenyuandGitHub 35c9701df0 update outdated tests and comments (#14090) 2026-01-10 01:00:48 -05:00
chenyuandGitHub 92246ea731 update tests, WEBGPU=1 pytest . passes (#14089)
* update tests, `WEBGPU=1 pytest .` passes

* minor update
2026-01-10 00:03:02 -05:00
chenyuandGitHub c34c6d9468 fix wgsl packed_store can drop valid (#14088)
* fix wgsl packed_store can drop valid

* fix
2026-01-09 15:22:06 -05:00
chenyuandGitHub eacccc5ace more disk assign tests (#14087)
covers more edge cases
2026-01-09 14:14:52 -05:00
chenyuandGitHub ed295e74dc don't skip gguf test if ggml is not installed (#14086)
* don't skip gguf test if ggml is not installed

should just let it fail

* fix
2026-01-09 12:05:58 -05:00
chenyuandGitHub cff33c8d78 add some disk assign tests (#14085) 2026-01-09 11:50:59 -05:00
chenyuandGitHub 74fa3c7d09 decomp pow for LVP (#14084)
test failed due to undefined behavior, so use decomp instead
2026-01-09 10:50:28 -05:00
0fbc551622 train bert with fp8 (#13874)
* fp8 train

* clean

* lint

* test fix from #13439

* skip first/last layer

* rm __init__, restore unroll <=32 check

* tests

* clean test, remove unused

* multi-gpu test, clean quantize_to_fp8

* remove bert contiguous

* run script

* test: better check

* run script search

* add seed in bert data shuffle

* move script to mi350x folder

---------

Co-authored-by: chenyu <[email protected]>
2026-01-09 09:21:59 -05:00
nimlgenandGitHub ba209d6305 am: utc_l1_enable on all sdma inst (#14083) 2026-01-09 17:17:05 +03:00
nimlgenandGitHub 6b308b89b7 viz: timeline time (#14080)
* viz: timeline time

* less lines

* cut
2026-01-09 16:43:45 +03:00
nimlgenandGitHub 40f9fa2db4 autogen: new kfd (#14082) 2026-01-09 16:08:17 +03:00
qazalandGitHub 2917ed1616 roc: propagate decoder errors to main thread (#14081)
* roc: propagate decoder errors to main thread

* types

* add cause
2026-01-09 21:10:45 +09:00
qazalandGitHub f3f4d9b387 viz: fix disasm node width (#14079) 2026-01-09 16:37:37 +09:00
c70c112254 fix CUDA=1 disassembly (VIZ=1) by stripping null terminator (#14046)
* fix ptxas disassembly bug

* single '

* move fix to get_bytes

* move rstrip

---------

Co-authored-by: qazal <[email protected]>
2026-01-09 15:19:59 +09:00
qazalandGitHub 13e5d00d0e viz: exclude comma in register highlight (#14078)
* viz: exclude comma in register highlight

* simplify
2026-01-09 15:10:30 +09:00
qazalandGitHub a071adffc0 viz: amdgpu disassembly register highlighting UI (#14059)
* viz: amdgpu disassembly register highlighting

* minor details

* details from IDA

* more details from IDA

* refactor token colors

* move tokenizer to python

* simplify

* minimal tokenizer for registers

* all the operand types
2026-01-09 11:27:09 +09:00
chenyuandGitHub b878f9d5a4 reuse Tensor init with const path [pr] (#14076) 2026-01-08 17:49:37 -05:00
chenyuandGitHub efcb32f6a9 unique const when requires_grad is set to True (#14075)
* unique const when requires_grad is set to True

* fix pyrender
2026-01-08 16:30:45 -05:00
chenyuandGitHub b34c637767 support bfloat16 for CL (#14073) 2026-01-08 14:14:29 -05:00
Garret CastroandGitHub 16b652302e skip bf16 test if not supported by device (#14070) 2026-01-08 13:37:24 -05:00
nimlgenandGitHub 3f61a96d79 am: SetSoftMaxByFreq on gfx10+ (#14068) 2026-01-08 17:00:03 +03:00
e7b5d8a434 assembly/amd: more RDNA4 asm (#14062)
* rdna4 more

* asm

* fixes

* assembly/amd: handwritten wmma failing test

* passes

* wmma default hacks

* space

* 0 skips in rdna3/rdna4 disasm

* more RDNA4 tests

---------

Co-authored-by: qazal <[email protected]>
2026-01-08 05:09:37 -08:00
nimlgenandGitHub e372c841ba hevc: beam in decode (#14067)
* hevc: beam in decode

* fine

* g
2026-01-08 15:47:16 +03:00
nimlgenandGitHub 1732a4ec4b am: rework set_clocks (#14065) 2026-01-08 15:33:32 +03:00
nimlgenandGitHub f3aceaa08b hevc: fast decoder (#14057) 2026-01-08 15:20:37 +03:00
qazalandGitHub 309197bca5 assembly/amd: test_roundtrip for cdna/rdna4 (#14066) 2026-01-08 21:03:13 +09:00
qazalandGitHub 15a056715d fix amd assembly IDE tests on macbook (#14063) 2026-01-08 17:27:52 +09:00
wozeparrotandGitHub 027b935269 tk: fix grouped load store (#14035) 2026-01-07 22:38:02 -08:00
George HotzandGitHub 2db04d0696 assembly/amd: start adding RDNA4 support (#14060)
* assembly/amd: start adding RDNA4 support

* rdna4 asm
2026-01-07 21:19:30 -08:00
George HotzandGitHub cb500466c2 assembly/amd: amd_asm_matmul (#13989)
* amd_asm_matmul

* dsl transform

* asm roundtrip

* fixed

* less

* better

* more

* simpler

* simplify

* lil

* simpler

* compact

* work

* cleanups

* simplify

* simpler

* cleanup

* name the regs

* simp

* big simp

* big simp

* simp

* acc grid

* fast

* stuff

* fast

* simpler

* owrks

* save vgprs

* save vgprs

* Compact

* less VGPRs

* after

* SQTT support

* fastest

* faster

* lil faster

* tile regs

* faster

* readable

* one more

* simpler

* lil simpler

* NO_GLOBAL skips early globals

* stock kernel

* cleanups

* cleanups

* one b reg

* safe reg changes

* acc is compact now

* remove confusing stuff

* sregs

* lds cleanups

* vopd
2026-01-07 20:11:05 -08:00
chenyuandGitHub 3caa1e2c98 fix cast HALF with PYTHON backend (#14058) 2026-01-07 16:52:05 -05:00
chenyuandGitHub 5f1ede7f7e clean up test_dtype (#14055)
use less lambda
2026-01-07 15:45:42 -05:00
nimlgenandGitHub 5bd4593eda hevc: cleaner decoder (#14056)
* hevc: cleaner decoder

* nn
2026-01-07 18:29:30 +03:00
b1tgandGitHub 241f0402b4 add seed in bert data shuffle (#14054) 2026-01-07 10:02:05 -05:00
nimlgenandGitHub 25c82dd242 nv: profile nvdec (#14053) 2026-01-07 15:56:54 +03:00
qazalandGitHub 35900290b2 viz: configure text height for cfg (#14052) 2026-01-07 18:58:56 +09:00
chenyuandGitHub 87f4bc5446 update variable names around jit [pr] (#14049)
lbs, st_vars_dtype_device and rawbuffers no more
2026-01-06 22:32:41 -05:00
chenyuandGitHub 2833c5a54b few more jit tests with multi tensor inputs (#14047) 2026-01-06 22:05:22 -05:00
chenyuandGitHub 72a3f78d19 jit includes tensor inputs in containers (#14043)
* jit includes tensor inputs in containers

* cleanup
2026-01-06 19:42:06 -05:00
chenyuandGitHub c714881832 don't allow jit input to be const (#14045)
* don't allow jit input to be unbuffered like const

* just const to fix multi

* fix rnnt
2026-01-06 18:15:22 -05:00
chenyuandGitHub a8896f28e1 test_unrealized_const_input_frozen (#14044)
unrealized const is not replaced in jit
2026-01-06 14:17:43 -05:00
nimlgenandGitHub 325f4006ff amd: copies w/o sdma (#14036)
* amd: copies w/o sdma

* as_args

* fixes

* f
2026-01-06 21:15:58 +03:00
chenyuandGitHub 7fb18f7e47 raise when jit fxn returns non-Tensor output (#14042) 2026-01-06 12:59:20 -05:00
chenyuandGitHub 4491ec0c9e JitError (#14041)
* JitError

* test_symbolic_jit
2026-01-06 12:19:50 -05:00
chenyuandGitHub 6ddddc68af test jit tolist failure (#14040)
also moved tests to test_jit_footguns
2026-01-06 11:16:57 -05:00
chenyuandGitHub b699b9f763 test case for jit a function with item call (#14039)
* test case for jit a function with item call

output is silently wrong now

* no dtype
2026-01-06 10:40:43 -05:00
nimlgenandGitHub 02084f5376 mockdsp: use dsp allocator (#14037)
* mockdsp: use dsp allocator

* fix

* ?
2026-01-06 16:04:47 +03:00
wozeparrotandGitHub 2b3e01e79c tk: support sliced local -> reg load (#14034) 2026-01-06 05:33:24 -05:00
George HotzandGitHub 45f7fd073d assembly/amd: pcode bug fixes (#14032)
* bring over pcode parser

* fixes

* pdf test

* delay alu
2026-01-06 00:15:48 -08:00
wozeparrotandGitHub 21d0f6bb76 tk: flat global -> local load (#14033) 2026-01-05 23:35:53 -08:00
qazalandGitHub 3170365a5b visualize SQTT with the same cfg infrastructure (#13870)
* start

* rough sketch

* post render dag

* art

* intro g key

* work

* custom color scale

* colors

* more blue

* better

* smaller

* use for loop in test
2026-01-06 14:53:20 +09:00
0120d69caa autogen: avcodec (and simplify workflow) (#14031)
* simplify autogen workflow and add avcodec verification

- Consolidate all regeneration into single steps (delete + import)
- Remove continue-on-error and individual diff checks
- Use git diff at end to catch all differences
- Show artifact URL in failure message
- Add avcodec.py verification

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

* patch avcodec

---------

Co-authored-by: Claude Opus 4.5 <[email protected]>
2026-01-05 23:30:25 -05:00
George HotzandGitHub 20653d2996 assembly/amd: make pdf.py code shine (#14029)
* assembly/amd: make pdf.py code shine

* no merge

* pdf2 is the future

* something

* regen enums

* test

* work

* remove junk

* write

* pcode extraction

* pdf2 passes all tests

* simplify

* simpler pdf

* late filter

* remove hacks

* simplify pdf2.py

* field type

* remove defaults

* don't export srcenum

* simple pdf.py

* simpler

* cleaner

* less hack in PDF
2026-01-05 18:49:40 -08:00
qazalandGitHub ea7b149ca5 viz command line tool (#14030) 2026-01-06 10:19:47 +09:00
sirhcmandGitHub f86c728440 load libclang as 'libclang.so' too (#14028) 2026-01-05 16:56:16 -05:00
chenyuandGitHub eda6a73897 clean up canonicalize_device (#14027)
centralize the type check
2026-01-05 10:29:55 -05:00
chenyuandGitHub ce464b147a clean up comments that mentioned outdated terms (#14026)
no MultiLazyBuffer and no ShapeTracker in comments
2026-01-05 09:42:58 -05:00
chenyuandGitHub 83063cc3e4 onnx TensorScatter (#14024) 2026-01-05 09:05:22 -05:00
chenyuandGitHub 9497ec00f2 fix onnx attention permute (#14025)
* fix onnx attention permute

* skip test_attention_4d_fp16_cpu too
2026-01-05 08:58:50 -05:00
qazalandGitHub 5cff5698f7 viz: g key toggles graph and text view (#14023) 2026-01-05 22:41:45 +09:00
chenyuandGitHub 7a81a3cb98 more passed onnx tests (#14022) 2026-01-05 07:46:27 -05:00
kim yongjinandGitHub 34fe105386 remove unused LazySeq (#14020) 2026-01-05 07:38:33 -05:00
qazalandGitHub 4f2f38bf64 viz: split cfg and table render (#14021) 2026-01-05 20:59:08 +09:00
nimlgenandGitHub 70405b4f3c am_smi: mi350 (#14018) 2026-01-05 13:10:56 +03:00
b2a0b9c551 autogen: dump patch in CI (#14010)
* autogen: don't fast-fail, produce patch artifact on differences

All verification steps now use continue-on-error to run completely.
Each job generates a patch artifact containing all differences found.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <[email protected]>

* add gen from header test

* fix tests

* fail if diff

* add forward decl autogen test

* remove confusing/wrong comments

* macos unittests set LIBCLANG_PATH

---------

Co-authored-by: Claude Sonnet 4.5 <[email protected]>
2026-01-04 22:38:12 -05:00
chenyuandGitHub aae08b20e0 enable passed onnx tests (#14017) 2026-01-04 22:12:50 -05:00
chenyuandGitHub 785d04d127 simpler einsum (#14014) 2026-01-04 20:38:59 -05:00
chenyuandGitHub f6a78a29e0 support einsum trace (#14012)
* support einsum trace

* test_einsum_scalar_cpu
2026-01-04 19:27:27 -05:00
George HotzandGitHub 404eed6172 assembly/amd: improve tests for asm (#14007)
* assembly/amd: improve tests for asm

* upd

* skip

* tests

* re bug

* more passing

* cleanups

* cdna fixups

* improve tests, better CDNA parsing

* fix CI

* no defs

* simpler

* all pass

* from pdf

* regen
2026-01-04 15:14:08 -08:00
wozeparrotandGitHub f550f9204c fa: failing test for bwd jit (#14009)
* tk: failing test for bwd jit

* feat: mark expectedFailure

* clean: spaces
2026-01-04 16:57:43 -05:00
George HotzandGitHub 7abf4591ba use bitsize on dtype (#14011)
* use bitsize on dtype [pr]

* bitsize

* bitsize in js export, but might be wrong

* reverts

* revert that
2026-01-04 12:16:21 -08:00
chenyuandGitHub cfb8bf5814 faster image load (#13977)
sometimes image load does not need to init with NAN
2026-01-04 13:09:59 -05:00
George HotzandGitHub 7ebda28692 assembly/amd: add CDNA support to asm (#13982)
* add CDNA support

* more cdna tests

* something

* fix more stuff

* more work

* simpler

* simplier

* cdna

* disasm

* less skip

* fixes

* simpler
2026-01-04 08:53:56 -08:00
chenyuandGitHub ad041416ca delete unused rewrite rule [pr] (#14006) 2026-01-04 09:48:52 -05:00
nimlgenandGitHub bf356ae996 am: mi300 48bit address space (#14004)
* am: mi300 48bit address space

* fix
2026-01-04 15:19:25 +03:00
nimlgenandGitHub 606786e152 am: do not sleep for each hive node during resets (#14003) 2026-01-04 14:02:11 +03:00
George HotzandGitHub 34ea053b26 assembly/amd: clean up pcode, jit pcode instead of static (#14001)
* assembly/amd: clean up pcode

* regen

* lil

* jit the pcode

* sendmsg

* cleanups

* inst prefetch lol
2026-01-03 23:06:15 -08:00
kamilisjonandGitHub 280790e438 Reuse toposort in recursive_property (#13993) 2026-01-03 22:04:13 -08:00
kamilisjonandGitHub 9a9564118c [pr] Delete reverse_toposort (#13987)
* Delete reverse_toposort

* Update comment and profiler name

* Update profiler name
2026-01-03 22:03:44 -08:00
George HotzandGitHub 8328511808 assembly/amd: make the emu.py code shine (#13996)
* assembly/amd: make the code shine

* lil clean

* reg back in pcode

* cleanups

* gen fma_mix

* no writelane hacks

* fn cleanup

* dead vgpr_write

* readable

* smem

* cleanup bench_emu

* speedups

* simpler and faster

* direct inst._fn

* split fxn

* Revert "simpler and faster"

This reverts commit e85f6594b3.

* move lds to wavestate

* dispatcher

* pc in dispatch

* literal isn't wavestate

* cleanups + program

* one readlane

* exec_vop3sd in exec_vop

* cleaner exec_vopd

* fully merge VOP3P

* no special paths

* no SliceProxy

* low=0

* no bigint

* failing tests

* fma on python 3.13
2026-01-03 20:33:09 -08:00
qazalandGitHub bdb421f13e process_replay: passthrough sink arg for Ops.PROGRAM input (#14000) 2026-01-04 13:09:39 +09:00
GalaxandGitHub 66caa9fe1d fix: library linking for fedora systems (#13999) 2026-01-03 17:40:56 -08:00
chenyuandGitHub 8003db2a28 test case of NOOP store load folding (#13997) 2026-01-03 14:39:26 -05:00
chenyuandGitHub c1b8644a3f test removing expander rules [pr] (#13994) 2026-01-03 12:38:01 -05:00
sirhcmandGitHub 35c2870b1f gate image_conv2d pitch hacks on IMAGE==1 (#13995)
* gate image_conv2d pitch hacks on IMAGE==1

* fix opencl image copies

* cleanup
2026-01-03 12:27:31 -05:00
nimlgenandGitHub a49924a0e9 hcq: _sleep report status (#13992)
* hcq: _sleep report status

* msg

* print all
2026-01-03 14:28:28 +03:00
nimlgenandGitHub 3b354bc11f hcq: better queue managment (#13991) 2026-01-03 13:11:15 +03:00
nimlgenandGitHub efb2ae87c6 hcq sync aql (#13756)
* hcq sync aql

* w
2026-01-03 12:59:24 +03:00
qazalandGitHub bd55507ee4 RDNA3 fp16 assembly gemm 85 TFLOPS (#13990) 2026-01-03 18:34:23 +09:00
wozeparrotandGitHub 6242a9d151 tk: no global copy and clear ranges (#13988) 2026-01-02 23:45:15 -08:00
wozeparrotandGitHub 9f082e8e25 fa: split kv bwd into 2 kernels (#13981) 2026-01-02 18:45:51 -08:00
qazalandGitHub 2cc64d71b0 simplify mi350x gemm / viz asm tests (#13984)
* mi350x gemm cleanup

* asm tests work

* simpler asm tests
2026-01-03 11:11:07 +09:00
chenyuandGitHub 7cbafb2ef1 update hypothesis min version (#13983)
there was a local_constants perf regression that made hypothesis related tests slow
2026-01-02 21:01:57 -05:00
sirhcmandGitHub 9dc524536f IMAGE=1 creates "dynamic" images (#13769)
* remove image from BufferSpec

* cl tiny_gemm (64) works

* mypy

* padding

* openpilot CL

* reshape properly

* remove extra qcom checks

* pad output

* mypy

* update compile test

* move undo

* TestImageCopy valid images

* TestImageRealization valid images

* TestImageDType valid images

* cleanups

* test_renderer_failures

* ruff

* mypy

* simplify ops_qcom

* bump step time

* Revert "bump step time"

This reverts commit 75a037c7d0.

* "dynamic textures" are optional

* a start

* IMAGE=1 works, no FLOAT16

* fast but wrong

* mypy

* some fixes

* better

* works

* refactor

* oops
2026-01-02 16:22:39 -05:00
sirhcmandGitHub 61dc70f1a8 add driving_vision IMAGE=1 benchmark (#13979) 2026-01-02 13:58:27 -05:00
George HotzandGitHub 0e282025ff assembly/amd: split test_emu into hw tests (#13966)
* assmebly/amd: split test_emu into hw tests

* hw tests

* bugfixes

* more tests and fix
2026-01-02 08:04:56 -08:00
chenyuandGitHub 2e2b5fed12 fix misspellings (#13976) 2026-01-02 10:37:38 -05:00
nietrasandGitHub f49e4714af Fix spelling errors in README for AMD assembly (#13975) 2026-01-02 10:15:20 -05:00
b1tgandGitHub a78fcc55a4 amd tc 1616128 (#13439)
* amd tc 1616128

* fix test

* remove hardcoded check in test
2026-01-02 09:01:05 -05:00
chenyuandGitHub fcbb896e05 remove unused to_struct [pr] (#13973) 2026-01-02 08:54:57 -05:00
nimlgenandGitHub ff7853a65a am: fix aid doorbells (#13971) 2026-01-02 15:53:44 +03:00
nimlgenandGitHub 42abb0586c am: fix aid doorbells (#13972) 2026-01-02 15:53:13 +03:00
nimlgenandGitHub ebbaad6bfd am: enable all sdma engines (#13970) 2026-01-02 15:25:15 +03:00
qazalandGitHub 5f52266225 mi350x gemm: use Tensor.custom_kernel in asm test (#13969)
* mi350x gemm: use Tensor.custom_kernel in asm test

* A @ B for baseline
2026-01-02 18:30:50 +09:00
George HotzandGitHub 5a1a561e0f assembly/amd: rdna4 autogen (#13967)
* assembly/amd: add pcode ds ops

* refactors

* fix ds op

* update autogen

* fix flat bug

* more tests

* fix emu test

* that's a hack

* generic

* fix all tests

* two tests

* fix test failure

* better

* remove __all__

* assembly/amd: fix autogen for RDNA4
2026-01-01 23:12:18 -05:00
wozeparrotandGitHub b27527f05a fix: missed inner tracked range (#13964) 2026-01-01 18:09:57 -08:00
wozeparrotandGitHub ecbac8a338 tk: fa cleanups + causal test (#13963) 2026-01-01 18:05:00 -08:00
chenyuandGitHub af0392efea only set DiskDevice.size if it opens successfully (#13962) 2026-01-01 19:33:26 -05:00
chenyuandGitHub e036d6df89 properly fix DiskDevice reuse (#13961) 2026-01-01 18:08:23 -05:00
George HotzandGitHub dfb813b760 assembly/amd: add pcode ds ops (#13939)
* assembly/amd: add pcode ds ops

* refactors

* fix ds op

* update autogen

* fix flat bug

* more tests

* fix emu test

* that's a hack

* generic

* fix all tests

* two tests

* fix test failure

* better

* remove __all__
2026-01-01 16:24:13 -05:00
chenyuandGitHub cb7c76a3bd update test_fuzz_failure to not contruct full UOp (#13960) 2026-01-01 15:09:58 -05:00
chenyuandGitHub 51398edf9c fix indirect import (#13958)
also deleted old external tests
2026-01-01 14:22:45 -05:00
chenyuandGitHub 8e416df438 simpler InvalidType [pr] (#13957)
simpler singleton pattern
2026-01-01 13:55:51 -05:00
nimlgenandGitHub b8ea0d779c am: remove pipe, queue from setup_ring (#13947) 2026-01-01 21:06:41 +03:00
chenyuandGitHub 4d5c4d256d update tqdm for edge case (#13956)
1.00kit/s and not 1000it/s for value 999.5
2026-01-01 11:37:26 -05:00
chenyuandGitHub ed222070f7 update xlog2 fp16 decomp to not use fp32 (#13955) 2026-01-01 11:18:29 -05:00
chenyuandGitHub ce84a23142 remove tee in benchmark (#13954) 2026-01-01 10:55:36 -05:00
24723327ac fix tc_up in search (#13438)
* tensor_core is missing from Scheduler

* test upcast max

---------

Co-authored-by: chenyu <[email protected]>
2026-01-01 10:25:08 -05:00
qazalandGitHub 9726500de8 enable using assembly in Tensor.custom_kernel (#13895) 2026-01-02 00:12:01 +09:00
qazalandGitHub c0f52c9dcb split assembly gemm to per arch directory (#13953) 2026-01-02 00:10:22 +09:00
chenyuandGitHub c69470be52 fix test_symbolic_arange_sym_step (#13952) 2026-01-01 09:41:07 -05:00
chenyuandGitHub b91b46091c delete test_tensor_uop (#13951)
old test for shape tracker. also update tests that refer shapetracker

names
2026-01-01 09:25:05 -05:00
chenyuandGitHub 17ef4af72c new ceildiv that fixed symbolic conv (#13944)
* new ceildiv that fixed symbolic conv

* smaller test case
2026-01-01 09:02:41 -05:00
qazalandGitHub 6a5430ab00 correct args order in mi350x gemm (#13949) 2026-01-01 23:01:46 +09:00
chenyuandGitHub baff10d32c clean up Tensor.svd slices (#13948) 2026-01-01 08:18:45 -05:00
nimlgenandGitHub 1c5ed8e8b5 am: remove doorbells from setup_ring (#13946) 2026-01-01 14:39:21 +03:00
haofeiandGitHub 526fd4ec71 Fix SVD rank‑1 Jacobi rotation when tau == 0 (#13945) 2026-01-01 00:30:18 -05:00
haofeiandGitHub 20777f30b9 Fix QR/SVD NaNs on zero/orthogonal inputs (#13943) 2025-12-31 23:40:09 -05:00
chenyuandGitHub 0ed58c1fcd clean up some functions in helpers [pr] (#13942) 2025-12-31 18:29:16 -05:00
chenyuandGitHub e2987001ee unify pre-commit mypy and ci mypy (#13940) 2025-12-31 17:51:51 -05:00
chenyuandGitHub 8bf7c9c1d2 no-op cleanups for ptx [pr] (#13938) 2025-12-31 17:28:39 -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
196 changed files with 32372 additions and 72399 deletions
+62 -92
View File
@@ -14,10 +14,12 @@ on:
paths:
- 'tinygrad/runtime/autogen/**/*'
- 'tinygrad/runtime/support/autogen.py'
- '.github/workflows/autogen.yml'
workflow_dispatch:
paths:
- 'tinygrad/runtime/autogen/**/*'
- 'tinygrad/runtime/support/autogen.py'
- '.github/workflows/autogen.yml'
jobs:
autogen:
@@ -39,102 +41,45 @@ jobs:
pydeps: 'pyyaml mako'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
- name: Verify OpenCL autogen
- name: Regenerate autogen files
run: |
mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
rm tinygrad/runtime/autogen/opencl.py
python3 -c "from tinygrad.runtime.autogen import opencl"
diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
- name: Verify CUDA autogen
run: |
mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
rm tinygrad/runtime/autogen/{cuda,nvrtc,nvjitlink,nv_570,nv}.py
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
- name: Verify AMD autogen
run: |
mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
- name: Verify Linux autogen
run: |
mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
rm tinygrad/runtime/autogen/{comgr,hsa,hip,amd_gpu,sqtt,rocprof}.py
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof"
rm tinygrad/runtime/autogen/am/{am,pm4_soc15,pm4_nv,sdma_4_0_0,sdma_5_0_0,sdma_6_0_0,smu_v13_0_0,smu_v14_0_2}.py
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
rm tinygrad/runtime/autogen/{libc,kfd,io_uring,ib,pci,vfio}.py
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
- name: Verify LLVM autogen
run: |
mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
rm tinygrad/runtime/autogen/llvm.py
python3 -c "from tinygrad.runtime.autogen import llvm"
diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
- name: Verify WebGPU autogen
run: |
mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
rm tinygrad/runtime/autogen/webgpu.py
python3 -c "from tinygrad.runtime.autogen import webgpu"
diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
- name: Verify Qualcomm autogen
run: |
mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
rm tinygrad/runtime/autogen/{kgsl,qcom_dsp}.py
python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
- name: Verify libusb autogen
run: |
mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
rm tinygrad/runtime/autogen/libusb.py
python3 -c "from tinygrad.runtime.autogen import libusb"
diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
- name: Verify mesa autogen
run: |
mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
rm tinygrad/runtime/autogen/mesa.py
python3 -c "from tinygrad.runtime.autogen import mesa"
diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
- name: Verify libclang autogen
run: |
cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
rm tinygrad/runtime/autogen/avcodec.py
python3 -c "from tinygrad.runtime.autogen import avcodec"
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
- name: Check for differences
run: |
if ! git diff --quiet; then
git diff > autogen-ubuntu.patch
echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
exit 1
fi
- name: Upload patch artifact
if: failure()
uses: actions/upload-artifact@v4
with:
name: autogen-ubuntu-patch
path: autogen-ubuntu.patch
autogen-mac:
name: In-tree Autogen (macos)
runs-on: macos-14
@@ -146,11 +91,24 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
llvm: 'true'
- name: Verify macos autogen
- name: Regenerate autogen files
run: |
mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
rm tinygrad/runtime/autogen/metal.py
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
- name: Check for differences
run: |
if ! git diff --quiet; then
git diff > autogen-macos.patch
echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
exit 1
fi
- name: Upload patch artifact
if: failure()
uses: actions/upload-artifact@v4
with:
name: autogen-macos-patch
path: autogen-macos.patch
autogen-comgr-3:
name: In-tree Autogen (comgr 3)
runs-on: ubuntu-24.04
@@ -169,8 +127,20 @@ jobs:
echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
sudo apt -qq update || true
sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
- name: Verify comgr (3) autogen
- name: Regenerate autogen files
run: |
mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
rm tinygrad/runtime/autogen/comgr_3.py
python3 -c "from tinygrad.runtime.autogen import comgr_3"
diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
- name: Check for differences
run: |
if ! git diff --quiet; then
git diff > autogen-comgr3.patch
echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
exit 1
fi
- name: Upload patch artifact
if: failure()
uses: actions/upload-artifact@v4
with:
name: autogen-comgr3-patch
path: autogen-comgr3.patch
+98 -212
View File
@@ -49,19 +49,19 @@ jobs:
- name: Print macOS version
run: sw_vers
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
- name: Run Stable Diffusion without fp16
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
- name: Run Stable Diffusion v2
# TODO: very slow step time
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
# process replay can't capture this, the graph is too large
- name: Run SDXL
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
- name: Run model inference benchmark
run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
- name: Test speed vs torch
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
- name: Test tensor cores
run: METAL=1 python3.11 test/opt/test_tensor_cores.py
- name: Test AMX tensor cores
@@ -71,84 +71,59 @@ jobs:
DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
- name: Run Tensor Core GEMM (float)
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
- name: Run Tensor Core GEMM (half)
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
- name: Run Tensor Core GEMM (bfloat16)
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
- name: Fuzz Padded Tensor Core GEMM
run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
- name: Run LLaMA
run: |
BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run LLaMA with BEAM
run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run quantized LLaMA
run: |
BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8
BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
- name: Run quantized LLaMA3
run: |
BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8
BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4
#- name: Run LLaMA 7B on 4 (virtual) GPUs
# run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
# run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run GPT2
run: |
BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run GPT2 w HALF
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
- name: Run GPT2 w HALF/BEAM
run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
- name: Run OLMoE
run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
- name: Train MNIST
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py
# NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
# the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
#- name: Run 10 CIFAR training steps
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
# run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
#- name: Run 10 CIFAR training steps w HALF
# run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
# run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py
#- name: Run 10 CIFAR training steps w BF16
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
# run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
- uses: actions/upload-artifact@v4
with:
name: Speed (Mac)
path: |
onnx_inference_speed.csv
torch_speed.txt
llama_unjitted.txt
llama_jitted.txt
llama_beam.txt
llama_int8.txt
llama_nf4.txt
llama3_int8.txt
llama3_nf4.txt
llama_four_gpu.txt
gpt2_unjitted.txt
gpt2_jitted.txt
gpt2_half.txt
gpt2_half_beam.txt
matmul.txt
matmul_half.txt
matmul_bfloat16.txt
sd.txt
sd_no_fp16.txt
sdv2.txt
sdxl.txt
beautiful_mnist.txt
train_cifar.txt
train_cifar_half.txt
train_cifar_bf16.txt
train_cifar_wino.txt
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py
@@ -215,7 +190,7 @@ jobs:
- name: Run model inference benchmark
run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
- name: Test speed vs torch
run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
- name: Test speed vs theoretical
run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
- name: Test benchmark allreduce
@@ -226,79 +201,58 @@ jobs:
NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
- name: Run Tensor Core GEMM (CUDA)
run: |
CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
- name: Run Tensor Core GEMM (PTX)
run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
- name: Run Tensor Core GEMM (NV)
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
- name: Test NV=1
run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
- name: Test CUDA=1
run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
# TODO: too slow
# - name: Run SDXL
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
# run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
- name: Run LLaMA
run: |
BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run LLaMA with BEAM
run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
# - name: Run LLaMA 7B on 4 GPUs
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
# - name: Run LLaMA 7B on 6 GPUs
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run LLaMA-3 8B BEAM
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
- name: Run LLaMA-3 8B on 4 GPUs with BEAM
run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
- name: Run quantized LLaMA3
run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
# - name: Run LLaMA-3 8B on 6 GPUs
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
# - name: Run LLaMA-2 70B
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_2_70B.txt
# run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run Mixtral 8x7B
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
- name: Run GPT2
run: |
BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run GPT2 w HALF
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
- name: Run GPT2 w HALF/BEAM
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
- uses: actions/upload-artifact@v4
with:
name: Speed (NVIDIA)
path: |
onnx_inference_speed.csv
torch_speed.txt
matmul.txt
matmul_bfloat16.txt
matmul_tf32.txt
matmul_ptx.txt
matmul_nv.txt
sd.txt
sdxl.txt
llama_unjitted.txt
llama_jitted.txt
llama_beam.txt
llama3_beam.txt
llama3_four_gpu.txt
llama3_six_gpu.txt
llama3_fp8.txt
llama_2_70B.txt
mixtral.txt
gpt2_unjitted.txt
gpt2_jitted.txt
gpt2_half.txt
gpt2_half_beam.txt
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -337,44 +291,30 @@ jobs:
# - name: Fuzz Padded Tensor Core GEMM (PTX)
# run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
- name: HEVC Decode Benchmark
run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
- name: Train MNIST
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
- name: Run 10 CIFAR training steps w BF16
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run MLPerf resnet eval on training data
run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (6 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (NVIDIA Training)
path: |
beautiful_mnist.txt
train_cifar.txt
train_cifar_half.txt
train_cifar_bf16.txt
train_cifar_wino.txt
train_cifar_one_gpu.txt
train_cifar_six_gpu.txt
train_resnet.txt
train_resnet_one_gpu.txt
train_bert.txt
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -426,7 +366,7 @@ jobs:
#- name: Test speed vs torch
# run: |
# python3 -c "import torch; print(torch.__version__)"
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
- name: Test speed vs theoretical
run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
- name: Test tensor cores AMD_LLVM=0
@@ -437,7 +377,7 @@ jobs:
- name: Run Tensor Core GEMM (AMD)
run: |
AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
- name: Test AMD=1
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
#- name: Test HIP=1
@@ -452,61 +392,39 @@ jobs:
- name: Test AM warm start time
run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
- name: Run Stable Diffusion
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
- name: Run SDXL
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing
- name: Run LLaMA 7B
run: |
BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run LLaMA 7B with BEAM
run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
# - name: Run LLaMA 7B on 4 GPUs
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_four_gpu.txt
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0 --timing
# - name: Run LLaMA 7B on 6 GPUs
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run LLaMA-3 8B BEAM
run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
- name: Run LLaMA-3 8B on 4 GPUs with BEAM
run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
# - name: Run LLaMA-3 8B on 6 GPUs
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
#- name: Restore amdgpu
# run: sudo modprobe amdgpu
# - name: Run LLaMA-2 70B
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_2_70B.txt
# run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run Mixtral 8x7B
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
- name: Run GPT2
run: |
BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
- name: Run GPT2 w HALF
run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
- name: Run GPT2 w HALF/BEAM
run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AMD)
path: |
onnx_inference_speed.csv
torch_speed.txt
llama_unjitted.txt
llama_jitted.txt
llama_beam.txt
llama3_beam.txt
llama3_four_gpu.txt
llama3_six_gpu.txt
llama_2_70B.txt
gpt2_unjitted.txt
gpt2_jitted.txt
gpt2_half.txt
gpt2_half_beam.txt
matmul.txt
matmul_amd.txt
sd.txt
sdxl.txt
mixtral.txt
run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -543,31 +461,20 @@ jobs:
- name: reset process replay
run: test/external/process_replay/reset.py
- name: Train MNIST
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
- name: Run 10 CIFAR training steps
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
- name: Run 10 CIFAR training steps w HALF
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
# - name: Run 10 CIFAR training steps w BF16
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
# run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
# TODO: too slow
# - name: Run 10 CIFAR training steps w winograd
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
# run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AMD Training)
path: |
beautiful_mnist.txt
train_cifar.txt
train_cifar_half.txt
train_cifar_bf16.txt
train_cifar_wino.txt
train_cifar_one_gpu.txt
train_cifar_six_gpu.txt
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -606,19 +513,12 @@ jobs:
- name: Run MLPerf resnet eval
run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (6 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AMD MLPerf)
path: |
train_resnet.txt
train_resnet_one_gpu.txt
train_bert.txt
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -648,6 +548,8 @@ jobs:
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.10.1 driving_vision
run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.10.1 driving_policy
@@ -708,7 +610,7 @@ jobs:
# AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
# AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
- name: Run Tensor Core GEMM (AMD)
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
- name: Test AMD=1
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
- name: Test DISK copy time
@@ -718,20 +620,12 @@ jobs:
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (1 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (AM Driver)
path: |
am_matmul_amd.txt
am_train_cifar_one_gpu.txt
am_train_resnet_one_gpu.txt
am_train_bert_one_gpu.txt
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -778,21 +672,13 @@ jobs:
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
- name: Test LLAMA-3
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (1 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
- uses: actions/upload-artifact@v4
with:
name: Speed (NV Driver)
path: |
nv_llama3_beam.txt
nv_train_cifar_one_gpu.txt
nv_train_resnet_one_gpu.txt
nv_train_bert_one_gpu.txt
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+27 -32
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
@@ -219,7 +218,6 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 10
# TODO: run the pre-commit hook to replace a lot of this
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -233,13 +231,14 @@ jobs:
run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string=' ' --recursive=y .
- name: Lint with ruff
run: |
pip3 install --upgrade --force-reinstall ruff==0.11.0
python3 -m ruff check .
pip3 install --upgrade --force-reinstall ruff==0.14.10
pre-commit run ruff --all-files
python3 -m ruff check examples/mlperf/ --ignore E501
python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
python3 -m ruff check extra/torch_backend/backend.py
- name: Run mypy
run: |
python -m mypy --strict-equality --lineprecision-report .
python -m mypy --lineprecision-report .
cat lineprecision.txt
- name: Run TYPED=1
run: TYPED=1 python -c "import tinygrad"
@@ -258,6 +257,7 @@ jobs:
key: unittest-12
pydeps: "pillow numpy ftfy regex"
deps: testing_unit
llvm: 'true'
- name: Check Device.DEFAULT
run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
- name: Run unit tests
@@ -310,7 +310,7 @@ jobs:
deps: testing_unit
python-version: '3.14'
- name: Test SPEC=2
run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
fuzzing:
name: Fuzzing
@@ -473,6 +473,8 @@ jobs:
name: Test LLM
runs-on: ubuntu-24.04
timeout-minutes: 15
env:
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -603,9 +605,7 @@ jobs:
WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run selected webgpu tests
run: |
WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit \
--ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py \
--ignore=test/test_fuzz_shape_ops.py --durations=20
WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -667,6 +667,11 @@ jobs:
key: rdna3-emu
deps: testing_minimal
amd: 'true'
python-version: '3.13'
- name: Verify AMD autogen is up to date
run: |
python -m extra.assembly.amd.pdf
git diff --exit-code extra/assembly/amd/autogen/
- name: Install LLVM 21
run: |
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
@@ -675,31 +680,19 @@ jobs:
sudo apt-get install llvm-21 llvm-21-tools cloc
- name: RDNA3 Line Count
run: cloc --by-file extra/assembly/amd/*.py
- name: Install rocprof-trace-decoder
run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
- name: Run RDNA3 emulator tests
run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 dtype tests
run: PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
run: PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
testamdautogen:
name: AMD autogen
runs-on: ubuntu-24.04
timeout-minutes: 10
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: rdna3-autogen
pydeps: "pdfplumber"
- name: Verify AMD autogen is up to date
run: |
python -m extra.assembly.amd.generate --arch all
git diff --exit-code extra/assembly/amd/autogen/
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
# TODO: run all once emulator is faster
- name: Run RDNA3 ops tests
run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
testnvidia:
strategy:
@@ -788,6 +781,8 @@ jobs:
ocelot: 'true'
llvm: 'true'
- name: Run unit tests
env:
LIBCLANG_PATH: '/opt/homebrew/opt/llvm@20/lib/libclang.dylib'
run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Run ONNX
run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+1 -1
View File
@@ -16,7 +16,7 @@ repos:
pass_filenames: false
- id: mypy
name: mypy
entry: python3 -m mypy tinygrad/ --strict-equality
entry: python3 -m mypy
language: system
always_run: true
pass_filenames: false
+13 -2
View File
@@ -192,9 +192,12 @@ When optimizing tinygrad internals:
9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
## Pattern Matching Profiling
## Pattern Matching Analysis
Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
**Use the right tool:**
- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs
```bash
TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
@@ -209,6 +212,14 @@ Key patterns to watch (from ResNet50 benchmark):
Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
```bash
# Save the trace
VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
# Explore it
./extra/viz/cli.py --help
```
## AMD Performance Counter Profiling
Set VIZ to `-2` to save performance counters traces for the AMD backend.
+1 -1
View File
@@ -70,7 +70,7 @@ AMD backend supports several interfaces for communicating with devices:
* `KFD`: uses the amdgpu driver
* `PCI`: uses the [AM driver](developer/am.md)
* `USB`: USB3 interafce for asm24xx chips.
* `USB`: USB3 interface for asm24xx chips.
You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
+3 -2
View File
@@ -213,12 +213,13 @@ class InterleavedDataset:
self.queues[queue_index].queue.extend(load_file(file))
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
def batch_load_train_bert(BS:int):
def batch_load_train_bert(BS:int, seed:int|None=None):
from extra.datasets.wikipedia import get_wiki_train_files
rng = random.Random(seed)
fs = sorted(get_wiki_train_files())
train_files = []
while fs: # TF shuffle
random.shuffle(fs)
rng.shuffle(fs)
train_files.append(fs.pop(0))
cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
+12 -1
View File
@@ -219,7 +219,18 @@ def get_mlperf_bert_model():
config = get_mlperf_bert_config()
if getenv("DISABLE_DROPOUT", 0):
config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
return BertForPretraining(**config)
model = BertForPretraining(**config)
if getenv("FP8_TRAIN"):
from extra.fp8.fp8_linear import convert_to_float8_training
def module_filter_fn(mod, fqn):
if isinstance(mod, LinearBert):
skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
return True
return False
convert_to_float8_training(model, module_filter_fn)
return model
def get_fake_data_bert(BS:int):
return {
+4 -3
View File
@@ -1008,6 +1008,7 @@ def train_bert():
config["DISABLE_DROPOUT"] = getenv("DISABLE_DROPOUT", 0)
config["TRAIN_BEAM"] = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
config["EVAL_BEAM"] = EVAL_BEAM = getenv("EVAL_BEAM", BEAM.value)
config["FP8_TRAIN"] = getenv("FP8_TRAIN", 0)
Tensor.manual_seed(seed) # seed for weight initialization
@@ -1085,7 +1086,7 @@ def train_bert():
if RUNMLPERF:
# only load real data with RUNMLPERF
eval_it = iter(batch_load_val_bert(EVAL_BS))
train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
train_it = iter(tqdm(batch_load_train_bert(BS, seed=seed), total=train_steps, disable=BENCHMARK))
for _ in range(start_step): next(train_it) # Fast forward
else:
# repeat fake data
@@ -1147,7 +1148,7 @@ def train_bert():
device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
loss = loss.item()
assert not math.isnan(loss)
if not getenv("FP8_TRAIN"): assert not math.isnan(loss)
lr = lr.item()
cl = time.perf_counter()
@@ -1160,7 +1161,7 @@ def train_bert():
if WANDB:
wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
"train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
"train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
"train/mem":GlobalCounters.mem_used / 1e9, "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
train_data, next_data = next_data, None
i += 1
@@ -0,0 +1,24 @@
#!/bin/bash
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
export BEAM_TIMEOUT_SEC=15
export FP8_TRAIN=1
# search
IGNORE_BEAM_CACHE=1 BENCHMARK=10 BERT_LAYERS=2 RUNMLPERF=0 python3 examples/mlperf/model_train.py
export WANDB=1 PARALLEL=0
RUNMLPERF=1 python3 examples/mlperf/model_train.py
+1 -1
View File
@@ -7,7 +7,7 @@ if __name__ == "__main__":
with open(fetch(sys.argv[1]), "rb") as f:
run_onnx_jit = pickle.load(f)
input_name = run_onnx_jit.captured.expected_names[0]
device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
device = run_onnx_jit.captured.expected_input_info[0][-1]
print(f"input goes into {input_name=} on {device=}")
hit = 0
for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
+10 -11
View File
@@ -153,8 +153,7 @@ class SMICtx:
tables = {}
for dev in self.devs:
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableX_t
case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
return tables
@@ -165,17 +164,17 @@ class SMICtx:
def get_gfx_activity(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
case _: return metrics.SmuMetrics.AverageGfxActivity
def get_mem_activity(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
case _: return metrics.SmuMetrics.AverageUclkActivity
def get_temps(self, dev, metrics, compact=False):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6):
case (13,0,6)|(13,0,12):
temps = {
"Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
"HBM": self._smuq10_round(metrics.MaxHbmTemperature),
@@ -191,7 +190,7 @@ class SMICtx:
def get_voltage(self, dev, metrics, compact=False):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return {}
case (13,0,6)|(13,0,12): return {}
case _:
voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
@@ -205,33 +204,33 @@ class SMICtx:
def get_gfx_freq(self, dev, metrics):
if metrics is None: return 0
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return self._smuq10_round(metrics.GfxclkFrequency[0])
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.GfxclkFrequency[0])
case _:
return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
def get_mem_freq(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return self._smuq10_round(metrics.UclkFrequency)
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.UclkFrequency)
case _:
return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
metrics.SmuMetrics.AverageMemclkFrequencyPreDs
def get_fckl_freq(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return self._smuq10_round(metrics.FclkFrequency)
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.FclkFrequency)
case _:
return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
metrics.SmuMetrics.AverageFclkFrequencyPreDs
def get_fan_rpm_pwm(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return None, None
case (13,0,6)|(13,0,12): return None, None
case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
def get_power(self, dev, metrics):
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
def get_mem_usage(self, dev):
+39
View File
@@ -0,0 +1,39 @@
An integrated environment for AMD GPU assembly and emulation
Test with `PYTHONPATH="." pytest -n12 extra/assembly/amd/`
`AMD_LLVM=1 PYTHONPATH="." pytest -n12 extra/assembly/amd/`
* pdf.py -- extract assembly format + instruction pseudocode from AMD PDF
* dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
* pcode.py -- pseudocode execution environment. pseudocode should be transformed as little as possible.
* asm.py -- an asm/disasm function to transform to and from AMD assembly syntax
* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
The autogen folder is autogenerated from the AMD PDFs with `python3 -m extra.assembly.amd.pdf --arch all`
test_emu.py has a good set of instruction tests for the emulation, with USE_HW=1 it will compare to real hardware.
Whenever an instruction is fixed, regression tests should be added here and confirmed with real hardware.
test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same as LLVM.
tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
`PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_dtype_alu.py test/test_dtype.py`
The ops tests also pass, but they are very slow, so you should run them one at a time.
`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=12 test/test_ops.py`
`SKIP_SLOW_TEST=1 PYTHONPATH="." AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=12 test/test_ops.py`
When something is caught by main tinygrad tests, a local regression test should be added to `extra/assembly/amd/test`.
While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
If a test is failing with `AMD=1 PYTHON_REMU=1 MOCKGPU=1` it's because an instruction is emulated incorrectly.
You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~2000 lines.
Get line count with `cloc --by-file extra/assembly/amd/*.py`
+254 -477
View File
@@ -1,479 +1,45 @@
# RDNA3 assembler and disassembler
# RDNA3/RDNA4/CDNA assembler
from __future__ import annotations
import re
from extra.assembly.amd.dsl import Inst, RawImm, Reg, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
from extra.assembly.amd.dsl import RawImm, SrcMod, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory
from extra.assembly.amd.dsl import VCC_LO, VCC_HI, VCC, EXEC_LO, EXEC_HI, EXEC, SCC, M0, NULL, OFF
from extra.assembly.amd.autogen.rdna3 import VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP
from extra.assembly.amd.autogen.rdna3 import VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, VINTERPOp
from extra.assembly.amd.autogen.rdna3 import SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp
from extra.assembly.amd.dsl import FLOAT_ENC
from extra.assembly.amd.autogen.rdna3 import ins
from extra.assembly.amd.autogen.rdna3.ins import VOP2Op, VOPDOp, SOPKOp
from extra.assembly.amd.autogen.rdna3.enum import BufFmt
from extra.assembly.amd.autogen.rdna4 import ins as rdna4_ins
# VOP3SD opcodes that share VOP3 encoding
VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
def detect_format(data: bytes) -> type[Inst]:
"""Detect instruction format from machine code bytes."""
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
word = int.from_bytes(data[:4], 'little')
hi2 = (word >> 30) & 0x3
if hi2 == 0b11:
enc = (word >> 26) & 0xf
if enc == 0b0010: return VOPD
if enc == 0b0011: return VOP3P
if enc == 0b0100: return VINTERP
if enc == 0b0101: return VOP3SD if ((word >> 16) & 0x3ff) in VOP3SD_OPS else VOP3
if enc == 0b0110: return DS
if enc == 0b0111: return FLAT
if enc == 0b1000: return MUBUF
if enc == 0b1010: return MTBUF
if enc == 0b1100 or enc == 0b1111: return MIMG
if enc == 0b1101: return SMEM
if enc == 0b1110: return EXP
raise ValueError(f"unknown 64-bit format enc={enc:#06b} word={word:#010x}")
if hi2 == 0b10:
enc = (word >> 23) & 0x7f
if enc == 0b1111101: return SOP1
if enc == 0b1111110: return SOPC
if enc == 0b1111111: return SOPP
return SOPK if ((word >> 28) & 0xf) == 0b1011 else SOP2
# hi2 == 0b00 or 0b01: VOP1/VOP2/VOPC (bit 31 = 0)
assert (word >> 31) == 0, f"expected bit 31 = 0 for VOP, got word={word:#010x}"
enc = (word >> 25) & 0x7f
if enc == 0b0111110: return VOPC
if enc == 0b0111111: return VOP1
if enc <= 0b0111101: return VOP2
raise ValueError(f"unknown VOP format enc={enc:#09b} word={word:#010x}")
# Re-export disasm for backwards compatibility
from extra.assembly.amd.disasm import disasm, HWREG, HWREG_RDNA4
# ═══════════════════════════════════════════════════════════════════════════════
# CONSTANTS
# ═══════════════════════════════════════════════════════════════════════════════
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
HWREG_IDS = {v.lower(): k for k, v in HWREG.items()}
MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
# ═══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def decode_src(val: int) -> str:
if val <= 105: return f"s{val}"
if val in SPECIAL_DEC: return SPECIAL_DEC[val]
if 108 <= val <= 123: return f"ttmp{val - 108}"
if 128 <= val <= 192: return str(val - 128)
if 193 <= val <= 208: return str(-(val - 192))
if 256 <= val <= 511: return f"v{val - 256}"
return "lit" if val == 255 else f"?{val}"
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{b}" if n == 1 else f"{p}[{b}:{b+n-1}]"
def _sreg(b: int, n: int = 1) -> str: return _reg("s", b, n)
def _vreg(b: int, n: int = 1) -> str: return _reg("v", b, n)
def _hl(v: int, hi_thresh: int = 128) -> str: return 'h' if v >= hi_thresh else 'l'
def _fmt_sdst(v: int, n: int = 1) -> str:
if v == 124: return "null"
if 108 <= v <= 123: return _reg("ttmp", v - 108, n)
if n > 1: return SPECIAL_PAIRS.get(v) or _sreg(v, n)
return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
def _fmt_src(v: int, n: int = 1) -> str:
if n == 1: return decode_src(v)
if v >= 256: return _vreg(v - 256, n)
if v <= 105: return _sreg(v, n)
if n == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if 108 <= v <= 123: return _reg("ttmp", v - 108, n)
return decode_src(v)
def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
return f"v{(v - base) & 0x7f}.{_hl(v, hi_thresh)}"
# RDNA unified buffer format
BUF_FMT = {e.name: e.value for e in BufFmt}
_BUF_FMT_EXT = {'BUF_FMT_32_32_32_32_SINT': 62, 'BUF_FMT_32_32_32_32_FLOAT': 63, 'BUF_FMT_8_FLOAT': 108}
BUF_FMT.update(_BUF_FMT_EXT)
def _parse_buf_fmt_combo(s: str) -> int:
parts = [p.strip().replace('BUF_DATA_FORMAT_', '').replace('BUF_NUM_FORMAT_', '') for p in s.split(',')]
return BUF_FMT.get(f'BUF_FMT_{parts[0]}_{parts[1]}') if len(parts) == 2 else None
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
def _is16(op: str) -> bool: return _has(op, 'f16', 'i16', 'u16', 'b16') and not _has(op, '_f32', '_i32')
def _is64(op: str) -> bool: return _has(op, 'f64', 'i64', 'u64', 'b64')
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool, any_hi: bool) -> str:
"""Format VOP3 source operand with modifiers."""
if n > 1: s = _fmt_src(v, n)
elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else (f"v{v - 256}.l" if any_hi else inst.lit(v))
else: s = inst.lit(v)
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
"""Format op_sel modifier string."""
if not need: return ""
if is16_d and (opsel & 8): return f" op_sel:[1,1,1{',1' if n == 3 else ''}]"
if n == 3: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
# ═══════════════════════════════════════════════════════════════════════════════
# DISASSEMBLER
# ═══════════════════════════════════════════════════════════════════════════════
def _disasm_vop1(inst: VOP1) -> str:
op = VOP1Op(inst.op)
if op in (VOP1Op.V_NOP, VOP1Op.V_PIPEFLUSH): return op.name.lower()
F64_OPS = {VOP1Op.V_CEIL_F64, VOP1Op.V_FLOOR_F64, VOP1Op.V_FRACT_F64, VOP1Op.V_FREXP_MANT_F64, VOP1Op.V_RCP_F64, VOP1Op.V_RNDNE_F64, VOP1Op.V_RSQ_F64, VOP1Op.V_SQRT_F64, VOP1Op.V_TRUNC_F64}
is_f64_d = op in F64_OPS or op in (VOP1Op.V_CVT_F64_F32, VOP1Op.V_CVT_F64_I32, VOP1Op.V_CVT_F64_U32)
is_f64_s = op in F64_OPS or op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64, VOP1Op.V_FREXP_EXP_I32_F64)
name = op.name.lower()
parts = name.split('_')
is_16d = any(p in ('f16','i16','u16','b16') for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in ('f16','i16','u16','b16') and 'cvt' not in name)
is_16s = parts[-1] in ('f16','i16','u16','b16') and 'sat_pk' not in name
if op == VOP1Op.V_READFIRSTLANE_B32: return f"v_readfirstlane_b32 {decode_src(inst.vdst)}, v{inst.src0 - 256 if inst.src0 >= 256 else inst.src0}"
dst = _vreg(inst.vdst, 2) if is_f64_d else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
src = _fmt_src(inst.src0, 2) if is_f64_s else _fmt_v16(inst.src0) if is_16s and inst.src0 >= 256 else inst.lit(inst.src0)
return f"{name}_e32 {dst}, {src}"
def _disasm_vop2(inst: VOP2) -> str:
op = VOP2Op(inst.op)
name = op.name.lower()
suf = "" if op == VOP2Op.V_DOT2ACC_F32_F16 else "_e32"
is16 = _is16(name) and 'pk_' not in name
# fmaak: dst = src0 * vsrc1 + K, fmamk: dst = src0 * K + vsrc1
if op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{inst._literal:x}"
if op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16): return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{inst._literal:x}, v{inst.vsrc1}"
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_fmt_v16(inst.src0) if inst.src0 >= 256 else inst.lit(inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (", vcc_lo" if op == VOP2Op.V_CNDMASK_B32 else "")
VOPC_CLASS = {VOPCOp.V_CMP_CLASS_F16, VOPCOp.V_CMP_CLASS_F32, VOPCOp.V_CMP_CLASS_F64,
VOPCOp.V_CMPX_CLASS_F16, VOPCOp.V_CMPX_CLASS_F32, VOPCOp.V_CMPX_CLASS_F64}
def _disasm_vopc(inst: VOPC) -> str:
op = VOPCOp(inst.op)
name = op.name.lower()
is64, is16 = _is64(name), _is16(name)
s0 = _fmt_src(inst.src0, 2) if is64 else _fmt_v16(inst.src0) if is16 and inst.src0 >= 256 else inst.lit(inst.src0)
s1 = _vreg(inst.vsrc1, 2) if is64 and op not in VOPC_CLASS else _fmt_v16(inst.vsrc1, 0, 128) if is16 else f"v{inst.vsrc1}"
return f"{name}_e32 {s0}, {s1}" if op.value >= 128 else f"{name}_e32 vcc_lo, {s0}, {s1}"
NO_ARG_SOPP = {SOPPOp.S_ENDPGM, SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE}
def _disasm_sopp(inst: SOPP) -> str:
op, name = SOPPOp(inst.op), SOPPOp(inst.op).name.lower()
if op in NO_ARG_SOPP: return name
if op == SOPPOp.S_WAITCNT:
vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if op == SOPPOp.S_DELAY_ALU:
deps, skips = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3'], ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
return f"{name} {inst.simm16}" if name.startswith(('s_cbranch', 's_branch')) else f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
op = SMEMOp(inst.op)
name = op.name.lower()
if op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
off_s = f"{decode_src(inst.soffset)} offset:0x{inst.offset:x}" if inst.offset and inst.soffset != 124 else f"0x{inst.offset:x}" if inst.offset else decode_src(inst.soffset)
sbase_idx, sbase_count = inst.sbase * 2, 4 if (8 <= inst.op <= 12 or name == 's_atc_probe_buffer') else 2
sbase_str = _fmt_src(sbase_idx, sbase_count) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(inst.op, 1)
return f"{name} {_fmt_sdst(inst.sdata, width)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (inst.dlc, " dlc"))
def _disasm_flat(inst: FLAT) -> str:
name = FLATOp(inst.op).name.lower()
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
suffix = name.split('_')[-1]
w = {'b32':1,'b64':2,'b96':3,'b128':4,'u8':1,'i8':1,'u16':1,'i16':1,'u32':1,'i32':1,'u64':2,'i64':2,'f32':1,'f64':2}.get(suffix, 1)
if 'cmpswap' in name: w *= 2
if name.endswith('_x2') or 'x2' in suffix: w = max(w, 2)
mods = f"{f' offset:{off_val}' if off_val else ''}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
# saddr
if seg == 'flat' or inst.saddr == 0x7F: saddr_s = ""
elif inst.saddr == 124: saddr_s = ", off"
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr)}"
elif inst.saddr in SPECIAL_PAIRS: saddr_s = f", {SPECIAL_PAIRS[inst.saddr]}"
elif 108 <= inst.saddr <= 123: saddr_s = f", {_reg('ttmp', inst.saddr - 108, 2)}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr)}"
# addtid: no addr
if 'addtid' in name: return f"{instr} v{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
# addr width
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, 1 if seg == 'scratch' or (inst.saddr not in (0x7F, 124)) else 2)
data_s, vdst_s = _vreg(inst.data, w), _vreg(inst.vdst, w // 2 if 'cmpswap' in name else w)
if 'atomic' in name:
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if inst.glc else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
return f"{instr} {_vreg(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
def _disasm_ds(inst: DS) -> str:
op, name = DSOp(inst.op), DSOp(inst.op).name.lower()
gds = " gds" if inst.gds else ""
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
off2 = f" offset0:{inst.offset0} offset1:{inst.offset1}" if inst.offset0 or inst.offset1 else ""
w = 4 if '128' in name else 3 if '96' in name else 2 if (name.endswith('64') or 'gs_reg' in name) else 1
d0, d1, dst, addr = _vreg(inst.data0, w), _vreg(inst.data1, w), _vreg(inst.vdst, w), f"v{inst.addr}"
if op == DSOp.DS_NOP: return name
if op == DSOp.DS_BVH_STACK_RTN_B32: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}, {_vreg(inst.data1, 4)}{off}{gds}"
if 'gws_sema' in name and op != DSOp.DS_GWS_SEMA_BR: return f"{name}{off}{gds}"
if 'gws_' in name: return f"{name} {addr}{off}{gds}"
if op in (DSOp.DS_CONSUME, DSOp.DS_APPEND): return f"{name} v{inst.vdst}{off}{gds}"
if 'gs_reg' in name: return f"{name} {_vreg(inst.vdst, 2)}, v{inst.data0}{off}{gds}"
if '2addr' in name:
if 'load' in name: return f"{name} {_vreg(inst.vdst, w*2)}, {addr}{off2}{gds}"
if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
return f"{name} {_vreg(inst.vdst, w*2)}, {addr}, {d0}, {d1}{off2}{gds}"
if 'load' in name: return f"{name} v{inst.vdst}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
if 'store' in name and not _has(name, 'cmp', 'xchg'):
return f"{name} v{inst.data0}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
if 'swizzle' in name or op == DSOp.DS_ORDERED_COUNT: return f"{name} v{inst.vdst}, {addr}{off}{gds}"
if 'permute' in name: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}{off}{gds}"
if 'condxchg' in name: return f"{name} {_vreg(inst.vdst, 2)}, {addr}, {_vreg(inst.data0, 2)}{off}{gds}"
if _has(name, 'cmpstore', 'mskor', 'wrap'):
return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
def _disasm_vop3(inst: VOP3) -> str:
op = VOP3SDOp(inst.op) if inst.op in VOP3SD_OPS else VOP3Op(inst.op)
name = op.name.lower()
# VOP3SD (shared encoding)
if inst.op in VOP3SD_OPS:
sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
is64, mad64 = 'f64' in name, _has(name, 'mad_i64_i32', 'mad_u64_u32')
def src(v, neg, ext=False): s = _fmt_src(v, 2) if ext or is64 else inst.lit(v); return f"-{s}" if neg else s
s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4, mad64)
dst = _vreg(inst.vdst, 2) if is64 or mad64 else f"v{inst.vdst}"
if op in (VOP3SDOp.V_ADD_CO_U32, VOP3SDOp.V_SUB_CO_U32, VOP3SDOp.V_SUBREV_CO_U32): return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {s0}, {s1}"
if op in (VOP3SDOp.V_ADD_CO_CI_U32, VOP3SDOp.V_SUB_CO_CI_U32, VOP3SDOp.V_SUBREV_CO_CI_U32): return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {s0}, {s1}, {s2}"
return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {s0}, {s1}, {s2}" + _omod(inst.omod)
# Detect operand sizes
is64 = _is64(name)
is64_src, is64_dst = False, False
is16_d = is16_s = is16_s2 = False
if 'cvt_pk' in name: is16_s = name.endswith('16')
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', name):
is16_d, is16_s = _has(m.group(1), 'f16','i16','u16','b16'), _has(m.group(2), 'f16','i16','u16','b16')
is64_src, is64_dst = '64' in m.group(2), '64' in m.group(1)
is16_s2, is64 = is16_s, False
elif re.match(r'v_mad_[iu]32_[iu]16', name): is16_s = True
elif 'pack_b32' in name: is16_s = is16_s2 = True
else: is16_d = is16_s = is16_s2 = _is16(name) and not _has(name, 'dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad')
# Source counts
shift64 = 'rev' in name and '64' in name and name.startswith('v_')
ldexp64 = op == VOP3Op.V_LDEXP_F64
trig = op == VOP3Op.V_TRIG_PREOP_F64
sad64, mqsad = _has(name, 'qsad_pk', 'mqsad_pk'), 'mqsad_u32' in name
s0n = 2 if ((is64 and not shift64) or sad64 or mqsad or is64_src) else 1
s1n = 2 if (is64 and not _has(name, 'class') and not ldexp64 and not trig) else 1
s2n = 4 if mqsad else 2 if (is64 or sad64) else 1
any_hi = inst.opsel != 0
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, s0n, is16_s, any_hi)
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, s1n, is16_s, any_hi)
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, s2n, is16_s2, any_hi)
# Destination
dn = 4 if mqsad else 2 if (is64 or sad64 or is64_dst) else 1
if op == VOP3Op.V_READLANE_B32: dst = _fmt_sdst(inst.vdst, 1)
elif dn > 1: dst = _vreg(inst.vdst, dn)
elif is16_d: dst = f"v{inst.vdst}.h" if (inst.opsel & 8) else f"v{inst.vdst}.l" if any_hi else f"v{inst.vdst}"
else: dst = f"v{inst.vdst}"
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
nonvgpr_opsel = (inst.src0 < 256 and (inst.opsel & 1)) or (inst.src1 < 256 and (inst.opsel & 2)) or (inst.src2 < 256 and (inst.opsel & 4))
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
if inst.op < 256: # VOPC
return f"{name}_e64 {s0}, {s1}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}"
if inst.op < 384: # VOP2
os = _opsel_str(inst.opsel, 3, need_opsel, is16_d) if 'cndmask' in name else _opsel_str(inst.opsel, 2, need_opsel, is16_d)
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if 'cndmask' in name else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
if inst.op < 512: # VOP1
return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{_opsel_str(inst.opsel, 1, need_opsel, is16_d)}{cl}{om}"
# Native VOP3
is3 = _has(name, 'fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube', 'bfe', 'bfi',
'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add', 'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit')
os = _opsel_str(inst.opsel, 3 if is3 else 2, need_opsel, is16_d)
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if is3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
def _disasm_vop3sd(inst: VOP3SD) -> str:
op, name = VOP3SDOp(inst.op), VOP3SDOp(inst.op).name.lower()
is64, mad64 = 'f64' in name, _has(name, 'mad_i64_i32', 'mad_u64_u32')
def src(v, neg, ext=False): s = _fmt_src(v, 2) if ext or is64 else inst.lit(v); return f"-{s}" if neg else s
s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4, mad64)
dst, is2src = _vreg(inst.vdst, 2) if is64 or mad64 else f"v{inst.vdst}", op in (VOP3SDOp.V_ADD_CO_U32, VOP3SDOp.V_SUB_CO_U32, VOP3SDOp.V_SUBREV_CO_U32)
suffix = "_e64" if name.startswith('v_') and 'co_' in name else ""
return f"{name}{suffix} {dst}, {_fmt_sdst(inst.sdst, 1)}, {s0}, {s1}{'' if is2src else f', {s2}'}{' clamp' if inst.clmp else ''}{_omod(inst.omod)}"
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal or inst.literal
vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), VOPDOp(inst.opx).name.lower(), VOPDOp(inst.opy).name.lower()
def half(n, vd, s0, vs1): return f"{n} v{vd}, {inst.lit(s0)}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}" if 'mov' in n else f"{n} v{vd}, {inst.lit(s0)}, v{vs1}{f', 0x{lit:x}' if lit and _has(n, 'fmaak', 'fmamk') else ''}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
def _disasm_vop3p(inst: VOP3P) -> str:
name = VOP3POp(inst.op).name.lower()
is_wmma, is_3src, is_fma_mix = 'wmma' in name, _has(name, 'fma', 'mad', 'dot', 'wmma'), 'fma_mix' in name
if is_wmma:
sc = 2 if 'iu4' in name else 4 if 'iu8' in name else 8
src0, src1, src2, dst = _fmt_src(inst.src0, sc), _fmt_src(inst.src1, sc), _fmt_src(inst.src2, 8), _vreg(inst.vdst, 8)
else: src0, src1, src2, dst = _fmt_src(inst.src0, 1), _fmt_src(inst.src1, 1), _fmt_src(inst.src2, 1), f"v{inst.vdst}"
n, opsel_hi = 3 if is_3src else 2, inst.opsel_hi | (inst.opsel_hi2 << 2)
if is_fma_mix:
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if inst.clmp else [])
else:
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != (7 if is_3src else 3) else []) + \
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if is_3src else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
def _disasm_buf(inst: MUBUF | MTBUF) -> str:
op = MTBUFOp(inst.op) if isinstance(inst, MTBUF) else MUBUFOp(inst.op)
name = op.name.lower()
if op in (MUBUFOp.BUFFER_GL0_INV, MUBUFOp.BUFFER_GL1_INV): return name
w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
{'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], 1)
if inst.tfe: w += 1
vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else f"v{inst.vaddr}" if inst.offen or inst.idxen else "off"
srsrc = _reg("ttmp", inst.srsrc*4 - 108, 4) if 108 <= inst.srsrc*4 <= 123 else _sreg(inst.srsrc*4, 4)
mods = ([f"format:{inst.format}"] if isinstance(inst, MTBUF) else []) + [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.glc,"glc"),(inst.dlc,"dlc"),(inst.slc,"slc"),(inst.tfe,"tfe")] if c]
return f"{name} {_vreg(inst.vdata, w)}, {vaddr}, {srsrc}, {decode_src(inst.soffset)}{' ' + ' '.join(mods) if mods else ''}"
def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
"""Calculate vaddr register count for MIMG sample/gather operations."""
# 1d,2d,3d,cube,1d_arr,2d_arr,2d_msaa,2d_msaa_arr
base = [1, 2, 3, 3, 2, 3, 3, 4][dim] # address coords
grad = [1, 2, 3, 2, 1, 2, 2, 2][dim] # gradient coords (for derivatives)
if 'get_resinfo' in name: return 1 # only mip level
packed, unpacked = 0, 0
if '_mip' in name: packed += 1
elif 'sample' in name or 'gather' in name:
if '_o' in name: unpacked += 1 # offset
if re.search(r'_c(_|$)', name): unpacked += 1 # compare (not _cl)
if '_d' in name: unpacked += (grad + 1) & ~1 if '_g16' in name else grad*2 # derivatives
if '_b' in name: unpacked += 1 # bias
if '_l' in name and '_cl' not in name and '_lz' not in name: packed += 1 # LOD
if '_cl' in name: packed += 1 # clamp
return (base + packed + 1) // 2 + unpacked if a16 else base + packed + unpacked
def _disasm_mimg(inst: MIMG) -> str:
name = MIMGOp(inst.op).name.lower()
srsrc_base = inst.srsrc * 4
srsrc_str = _reg("ttmp", srsrc_base - 108, 8) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 8)
# BVH intersect ray: special case with 4 SGPR srsrc
if 'bvh' in name:
vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
srsrc = _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {srsrc}{' a16' if inst.a16 else ''}"
# vdata width from dmask (gather4/msaa_load always 4), d16 packs, tfe adds 1
vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
if inst.d16: vdata = (vdata + 1) // 2
if inst.tfe: vdata += 1
# vaddr width
dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"dim_{inst.dim}"
vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
vaddr_str = f"v{inst.vaddr}" if vaddr == 1 else _vreg(inst.vaddr, vaddr)
# modifiers
mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask and (inst.dmask != 15 or 'atomic' in name) else []
mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
for flag, mod in [(inst.unrm,"unorm"),(inst.glc,"glc"),(inst.slc,"slc"),(inst.dlc,"dlc"),(inst.r128,"r128"),
(inst.a16,"a16"),(inst.tfe,"tfe"),(inst.lwe,"lwe"),(inst.d16,"d16")]:
if flag: mods.append(mod)
# ssamp for sample/gather/get_lod
ssamp_str = ""
if 'sample' in name or 'gather' in name or 'get_lod' in name:
ssamp_base = inst.ssamp * 4
ssamp_str = ", " + (_reg("ttmp", ssamp_base - 108, 4) if 108 <= ssamp_base <= 123 else _sreg(ssamp_base, 4))
return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
def _sop_widths(name: str) -> tuple[int, int, int]:
"""Return (dst_width, src0_width, src1_width) in register count for SOP instructions."""
if name in ('s_bitset0_b64', 's_bitset1_b64', 's_bfm_b64'): return 2, 1, 1
if name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return 2, 2, 1
if name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return 1, 2, 1
if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', name): return 2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1, 1
if m := re.search(r'_(b|i|u)(32|64)$', name): sz = 2 if m.group(2) == '64' else 1; return sz, sz, sz
return 1, 1, 1
def _disasm_sop1(inst: SOP1) -> str:
op, name = SOP1Op(inst.op), SOP1Op(inst.op).name.lower()
if op == SOP1Op.S_GETPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
if op in (SOP1Op.S_SETPC_B64, SOP1Op.S_RFE_B64): return f"{name} {_fmt_src(inst.ssrc0, 2)}"
if op == SOP1Op.S_SWAPPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {_fmt_src(inst.ssrc0, 2)}"
if op in (SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64): return f"{name} {_fmt_sdst(inst.sdst, 2 if 'b64' in name else 1)}, sendmsg({MSG.get(inst.ssrc0, str(inst.ssrc0))})"
dn, s0n, _ = _sop_widths(name)
return f"{name} {_fmt_sdst(inst.sdst, dn)}, {inst.lit(inst.ssrc0) if s0n == 1 else _fmt_src(inst.ssrc0, s0n)}"
def _disasm_sop2(inst: SOP2) -> str:
name = SOP2Op(inst.op).name.lower()
dn, s0n, s1n = _sop_widths(name)
return f"{name} {_fmt_sdst(inst.sdst, dn)}, {inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, s0n)}, {inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, s1n)}"
def _disasm_sopc(inst: SOPC) -> str:
name = SOPCOp(inst.op).name.lower()
_, s0n, s1n = _sop_widths(name)
return f"{name} {_fmt_src(inst.ssrc0, s0n)}, {_fmt_src(inst.ssrc1, s1n)}"
def _disasm_sopk(inst: SOPK) -> str:
op, name = SOPKOp(inst.op), SOPKOp(inst.op).name.lower()
if op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
if op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32):
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
hs = f"0x{inst.simm16:x}" if hid in (16, 17) else f"hwreg({HWREG.get(hid, str(hid))}, {hoff}, {hsz})"
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1)}" if op == SOPKOp.S_SETREG_B32 else f"{name} {_fmt_sdst(inst.sdst, 1)}, {hs}"
dn, _, _ = _sop_widths(name)
return f"{name} {_fmt_sdst(inst.sdst, dn)}, 0x{inst.simm16:x}"
def _disasm_vinterp(inst: VINTERP) -> str:
name = VINTERPOp(inst.op).name.lower()
src0 = f"-{inst.lit(inst.src0)}" if inst.neg & 1 else inst.lit(inst.src0)
src1 = f"-{inst.lit(inst.src1)}" if inst.neg & 2 else inst.lit(inst.src1)
src2 = f"-{inst.lit(inst.src2)}" if inst.neg & 4 else inst.lit(inst.src2)
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
return f"{name} v{inst.vdst}, {src0}, {src1}, {src2}" + (" " + mods if mods else "")
def _disasm_generic(inst: Inst) -> str:
name = f"op_{inst.op}"
def format_field(field_name, val):
val = unwrap(val)
if field_name in SRC_FIELDS: return inst.lit(val) if val != 255 else "0xff"
return f"{'s' if field_name == 'sdst' else 'v'}{val}" if field_name in ('sdst', 'vdst') else f"v{val}" if field_name == 'vsrc1' else f"0x{val:x}" if field_name == 'simm16' else str(val)
operands = [format_field(field_name, inst._values.get(field_name, 0)) for field_name in inst._fields if field_name not in ('encoding', 'op')]
return f"{name} {', '.join(operands)}" if operands else name
DISASM_HANDLERS = {VOP1: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, MUBUF: _disasm_buf, MTBUF: _disasm_buf,
MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk}
def disasm(inst: Inst) -> str: return DISASM_HANDLERS.get(type(inst), _disasm_generic)(inst)
# ═══════════════════════════════════════════════════════════════════════════════
# ASSEMBLER
# ═══════════════════════════════════════════════════════════════════════════════
SPEC_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'vcc': RawImm(106), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125),
'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'exec': RawImm(126), 'scc': RawImm(253), 'src_scc': RawImm(253)}
FLOATS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
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'}
SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b96', 's_load_b128', 's_load_b256', 's_load_b512',
's_load_i8', 's_load_u8', 's_load_i16', 's_load_u16',
's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b96', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512',
's_buffer_load_i8', 's_buffer_load_u8', 's_buffer_load_i16', 's_buffer_load_u16',
's_atc_probe', 's_atc_probe_buffer'}
SPEC_DSL = {'vcc_lo': 'VCC_LO', 'vcc_hi': 'VCC_HI', 'vcc': 'VCC_LO', 'null': 'NULL', 'off': 'OFF', 'm0': 'M0',
'exec_lo': 'EXEC_LO', 'exec_hi': 'EXEC_HI', 'exec': 'EXEC_LO', 'scc': 'SCC', 'src_scc': 'SCC'}
@@ -481,6 +47,7 @@ def _op2dsl(op: str) -> str:
op = op.strip()
neg = op.startswith('-') and not (op[1:2].isdigit() or (len(op) > 2 and op[1] == '0' and op[2] in 'xX'))
if neg: op = op[1:]
if op.startswith('neg(') and op.endswith(')'): neg = True; op = op[4:-1]
abs_ = (op.startswith('|') and op.endswith('|')) or (op.startswith('abs(') and op.endswith(')'))
if abs_: op = op[1:-1] if op.startswith('|') else op[4:-1]
hi = ".h" if op.endswith('.h') else ".l" if op.endswith('.l') else ""
@@ -507,29 +74,99 @@ def _parse_ops(s: str) -> list[str]:
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():]
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(), []
# Instruction aliases: LLVM uses different names for some instructions
_ALIASES = {
'v_cmp_tru_f16': 'v_cmp_t_f16', 'v_cmp_tru_f32': 'v_cmp_t_f32', 'v_cmp_tru_f64': 'v_cmp_t_f64',
'v_cmpx_tru_f16': 'v_cmpx_t_f16', 'v_cmpx_tru_f32': 'v_cmpx_t_f32', 'v_cmpx_tru_f64': 'v_cmpx_t_f64',
'v_cvt_flr_i32_f32': 'v_cvt_floor_i32_f32', 'v_cvt_rpi_i32_f32': 'v_cvt_nearest_i32_f32',
'v_ffbh_i32': 'v_cls_i32', 'v_ffbh_u32': 'v_clz_i32_u32', 'v_ffbl_b32': 'v_ctz_i32_b32',
'v_cvt_pkrtz_f16_f32': 'v_cvt_pk_rtz_f16_f32', 'v_fmac_legacy_f32': 'v_fmac_dx9_zero_f32', 'v_mul_legacy_f32': 'v_mul_dx9_zero_f32',
's_load_dword': 's_load_b32', 's_load_dwordx2': 's_load_b64', 's_load_dwordx4': 's_load_b128',
's_load_dwordx8': 's_load_b256', 's_load_dwordx16': 's_load_b512',
's_buffer_load_dword': 's_buffer_load_b32', 's_buffer_load_dwordx2': 's_buffer_load_b64',
's_buffer_load_dwordx4': 's_buffer_load_b128', 's_buffer_load_dwordx8': 's_buffer_load_b256',
's_buffer_load_dwordx16': 's_buffer_load_b512',
'v_cvt_pknorm_i16_f16': 'v_cvt_pk_norm_i16_f16', 'v_cvt_pknorm_u16_f16': 'v_cvt_pk_norm_u16_f16',
'v_add3_nc_u32': 'v_add3_u32', 'v_xor_add_u32': 'v_xad_u32',
'v_interp_p2_new_f32': 'v_interp_p2_f32',
's_ff1_i32_b32': 's_ctz_i32_b32', 's_ff1_i32_b64': 's_ctz_i32_b64',
's_flbit_i32_b32': 's_clz_i32_u32', 's_flbit_i32_b64': 's_clz_i32_u64', 's_flbit_i32': 's_cls_i32', 's_flbit_i32_i64': 's_cls_i32_i64',
's_andn1_saveexec_b32': 's_and_not0_saveexec_b32', 's_andn1_saveexec_b64': 's_and_not0_saveexec_b64',
's_andn1_wrexec_b32': 's_and_not0_wrexec_b32', 's_andn1_wrexec_b64': 's_and_not0_wrexec_b64',
's_andn2_saveexec_b32': 's_and_not1_saveexec_b32', 's_andn2_saveexec_b64': 's_and_not1_saveexec_b64',
's_andn2_wrexec_b32': 's_and_not1_wrexec_b32', 's_andn2_wrexec_b64': 's_and_not1_wrexec_b64',
's_orn1_saveexec_b32': 's_or_not0_saveexec_b32', 's_orn1_saveexec_b64': 's_or_not0_saveexec_b64',
's_orn2_saveexec_b32': 's_or_not1_saveexec_b32', 's_orn2_saveexec_b64': 's_or_not1_saveexec_b64',
's_andn2_b32': 's_and_not1_b32', 's_andn2_b64': 's_and_not1_b64',
's_orn2_b32': 's_or_not1_b32', 's_orn2_b64': 's_or_not1_b64',
'v_dot2c_f32_f16': 'v_dot2acc_f32_f16',
'v_fma_legacy_f32': 'v_fma_dx9_zero_f32',
'ds_read_b32': 'ds_load_b32', 'ds_read_b64': 'ds_load_b64', 'ds_read_b96': 'ds_load_b96', 'ds_read_b128': 'ds_load_b128',
'ds_read_i8': 'ds_load_i8', 'ds_read_u8': 'ds_load_u8', 'ds_read_i16': 'ds_load_i16', 'ds_read_u16': 'ds_load_u16',
'ds_read_i8_d16': 'ds_load_i8_d16', 'ds_read_u8_d16': 'ds_load_u8_d16', 'ds_read_i8_d16_hi': 'ds_load_i8_d16_hi', 'ds_read_u8_d16_hi': 'ds_load_u8_d16_hi',
'ds_read_u16_d16': 'ds_load_u16_d16', 'ds_read_u16_d16_hi': 'ds_load_u16_d16_hi',
'ds_read2_b32': 'ds_load_2addr_b32', 'ds_read2_b64': 'ds_load_2addr_b64',
'ds_read2st64_b32': 'ds_load_2addr_stride64_b32', 'ds_read2st64_b64': 'ds_load_2addr_stride64_b64',
'ds_read_addtid_b32': 'ds_load_addtid_b32', 'ds_write_addtid_b32': 'ds_store_addtid_b32',
'ds_write_b32': 'ds_store_b32', 'ds_write_b64': 'ds_store_b64', 'ds_write_b96': 'ds_store_b96', 'ds_write_b128': 'ds_store_b128',
'ds_write_b8': 'ds_store_b8', 'ds_write_b16': 'ds_store_b16',
'ds_write_b8_d16_hi': 'ds_store_b8_d16_hi', 'ds_write_b16_d16_hi': 'ds_store_b16_d16_hi',
'ds_write2_b32': 'ds_store_2addr_b32', 'ds_write2_b64': 'ds_store_2addr_b64',
'ds_write2st64_b32': 'ds_store_2addr_stride64_b32', 'ds_write2st64_b64': 'ds_store_2addr_stride64_b64',
'ds_wrxchg_rtn_b32': 'ds_storexchg_rtn_b32', 'ds_wrxchg_rtn_b64': 'ds_storexchg_rtn_b64',
'ds_wrxchg2_rtn_b32': 'ds_storexchg_2addr_rtn_b32', 'ds_wrxchg2_rtn_b64': 'ds_storexchg_2addr_rtn_b64',
'ds_wrxchg2st64_rtn_b32': 'ds_storexchg_2addr_stride64_rtn_b32', 'ds_wrxchg2st64_rtn_b64': 'ds_storexchg_2addr_stride64_rtn_b64',
}
def _apply_alias(text: str) -> str:
mn = text.split()[0].lower() if ' ' in text else text.lower().rstrip('_')
for m in (mn, mn.removesuffix('_e32'), mn.removesuffix('_e64')):
if m in _ALIASES: return _ALIASES[m] + text[len(m):]
return text
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
def get_dsl(text: str, arch: str = "rdna3") -> str:
text, kw = _apply_alias(text.strip()), []
# Extract modifiers
for pat, val in [(r'\s+mul:2(?:\s|$)', 1), (r'\s+mul:4(?:\s|$)', 2), (r'\s+div:2(?:\s|$)', 3)]:
if (m := _extract(text, pat))[0]: kw.append(f'omod={val}'); text = m[1]; break
if (m := _extract(text, r'\s+clamp(?:\s|$)'))[0]: kw.append('clmp=1'); text = m[1]
clamp_found = False
if (m := _extract(text, r'\s+clamp(?:\s|$)'))[0]: clamp_found = True; text = m[1]
opsel, m, text = None, *_extract(text, r'\s+op_sel:\[([^\]]+)\]')
if m:
bits, mn = [int(x.strip()) for x in m.group(1).split(',')], text.split()[0].lower()
is3p = mn.startswith(('v_pk_', 'v_wmma_', 'v_dot'))
is3p = mn.startswith(('v_pk_', 'v_wmma_', 'v_dot', 'v_fma_mix'))
opsel = (bits[0] | (bits[1] << 1) | (bits[2] << 2)) if len(bits) == 3 and is3p else \
(bits[0] | (bits[1] << 1) | (bits[2] << 3)) if len(bits) == 3 else sum(b << i for i, b in enumerate(bits))
opsel_hi_val, m, text = None, *_extract(text, r'\s+op_sel_hi:\[([^\]]+)\]')
if m: opsel_hi_val = [int(x.strip()) for x in m.group(1).split(',')]
m, text = _extract(text, r'\s+wait_exp:(\d+)'); waitexp = m.group(1) if m else None
m, text = _extract(text, r'\s+offset:(0x[0-9a-fA-F]+|-?\d+)'); off_val = m.group(1) if m else None
m, text = _extract(text, r'\s+dlc(?:\s|$)'); dlc = 1 if m else None
m, text = _extract(text, r'\s+glc(?:\s|$)'); glc = 1 if m else None
m, text = _extract(text, r'\s+slc(?:\s|$)'); slc = 1 if m else None
m, text = _extract(text, r'\s+tfe(?:\s|$)'); tfe = 1 if m else None
m, text = _extract(text, r'\s+offen(?:\s|$)'); offen = 1 if m else None
m, text = _extract(text, r'\s+idxen(?:\s|$)'); idxen = 1 if m else None
m, text = _extract(text, r'\s+format:\[([^\]]+)\]'); fmt_val = m.group(1) if m else None
m, text = _extract(text, r'\s+format:(\d+)'); fmt_val = m.group(1) if m and not fmt_val else fmt_val
m, text = _extract(text, r'\s+neg_lo:\[([^\]]+)\]'); neg_lo = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
m, text = _extract(text, r'\s+neg_hi:\[([^\]]+)\]'); neg_hi = sum(int(x.strip()) << i for i, x in enumerate(m.group(1).split(','))) if m else None
m, text = _extract(text, r'\s+byte_sel:(\d+)'); byte_sel = int(m.group(1)) if m else None
m, text = _extract(text, r'\s+offset0:(\d+)'); ds_off0 = int(m.group(1)) if m else None
m, text = _extract(text, r'\s+offset1:(\d+)'); ds_off1 = int(m.group(1)) if m else None
m, text = _extract(text, r'\s+index_key:(\d+)'); index_key = int(m.group(1)) if m else None
if waitexp: kw.append(f'waitexp={waitexp}')
if byte_sel is not None:
if opsel is None: opsel = 0
opsel |= (byte_sel << 2)
if ds_off0 is not None: kw.append(f'offset0={ds_off0}')
if ds_off1 is not None: kw.append(f'offset1={ds_off1}')
if index_key is not None: kw.append(f'opsel={index_key}')
parts = text.replace(',', ' ').split()
if not parts: raise ValueError("empty instruction")
@@ -560,24 +197,103 @@ def get_dsl(text: str) -> str:
# Special instructions
if mn == 's_setreg_imm32_b32': raise ValueError(f"unsupported: {mn}")
sop1_no_dest = ('s_alloc_vgpr', 's_barrier_init', 's_barrier_join', 's_barrier_signal', 's_barrier_signal_isfirst', 's_sleep_var')
if mn in sop1_no_dest:
return f"{mn}(sdst=RawImm(128), ssrc0={args[0]})"
if mn in ('s_setpc_b64', 's_rfe_b64'): return f"{mn}(ssrc0={args[0]})"
if mn in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'): return f"{mn}(sdst={args[0]}, ssrc0=RawImm({args[1].strip()}))"
if mn == 's_version': return f"{mn}(simm16={args[0]})"
if mn == 's_setreg_b32': return f"{mn}(simm16={args[0]}, sdst={args[1]})"
# Export instructions (RDNA4 VEXPORT)
if mn == 'export':
target_map = {**{f'mrt{i}': i for i in range(8)}, 'mrtz': 8, **{f'pos{i}': 12+i for i in range(4)}}
m, exp_str = _extract(op_str, r'\s+done(?:\s|$)')
done_val = 1 if m else 0
exp_parts = exp_str.replace(',', ' ').split()
target_name = exp_parts[0].lower().strip()
target = target_map.get(target_name, 0)
vsrcs, en = [], 0
for i, o in enumerate(exp_parts[1:5]):
o = o.strip().lower()
if o == 'off': vsrcs.append('v[0]')
else: vsrcs.append(_op2dsl(o)); en |= (1 << i)
return f"VEXPORT(target={target}, en={en}, vsrc0={vsrcs[0]}, vsrc1={vsrcs[1]}, vsrc2={vsrcs[2]}, vsrc3={vsrcs[3]}, done={done_val})"
# SMEM
if mn in SMEM_OPS:
gs, ds = ", glc=1" if glc else "", ", dlc=1" if dlc else ""
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})"
off_field = "ioffset" if arch == "rdna4" else "offset"
th_s, scope_s, smem_str = "", "", op_str
if arch == "rdna4":
m, smem_str = _extract(op_str, r'\s+th:TH_(\w+)')
th_val = {'LOAD_RT': 0, 'LOAD_NT': 1, 'LOAD_HT': 2, 'LOAD_LU': 3, 'STORE_RT': 0, 'STORE_NT': 1, 'STORE_HT': 2, 'STORE_LU': 3}.get(m.group(1), 0) if m else None
m, smem_str = _extract(smem_str, r'\s+scope:SCOPE_(\w+)')
scope_val = {'CU': 0, 'SE': 1, 'DEV': 2, 'SYS': 3}.get(m.group(1), 0) if m else None
if scope_val is None:
m, smem_str = _extract(smem_str, r'\s+scope:(0?x?[0-9a-fA-F]+)')
scope_val = int(m.group(1), 0) if m else None
th_s = f", th={th_val}" if th_val else ""
scope_s = f", scope={scope_val}" if scope_val else ""
smem_ops = _parse_ops(smem_str)
smem_args = [_op2dsl(o) for o in smem_ops]
if len(smem_ops) >= 3 and re.match(r'^-?[0-9]|^-?0x', smem_ops[2].strip().lower()):
return f"{mn}(sdata={smem_args[0]}, sbase={smem_args[1]}, {off_field}={smem_ops[2].strip()}, soffset=RawImm(124){gs}{ds}{th_s}{scope_s})"
if off_val and len(smem_ops) >= 3: return f"{mn}(sdata={smem_args[0]}, sbase={smem_args[1]}, {off_field}={off_val}, soffset={smem_args[2]}{gs}{ds}{th_s}{scope_s})"
if len(smem_ops) >= 3: return f"{mn}(sdata={smem_args[0]}, sbase={smem_args[1]}, soffset={smem_args[2]}{gs}{ds}{th_s}{scope_s})"
# Buffer
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)'})"
# Buffer (MUBUF/MTBUF/VBUFFER) instructions
if mn.startswith(('buffer_', 'tbuffer_')):
is_tbuf = mn.startswith('tbuffer_')
fmt_num = None
if fmt_val is not None:
if fmt_val.isdigit(): fmt_num = int(fmt_val)
else: fmt_num = BUF_FMT.get(fmt_val.replace(' ', '')) or _parse_buf_fmt_combo(fmt_val)
if mn in ('buffer_gl0_inv', 'buffer_gl1_inv', 'buffer_wbl2', 'buffer_inv'): return f"{mn}()"
if arch == "rdna4":
m, buf_text = _extract(op_str, r'\s+th:TH_(\w+)')
th_val = {'LOAD_RT': 0, 'LOAD_NT': 1, 'LOAD_HT': 2, 'LOAD_BYPASS': 3, 'LOAD_LU': 4, 'LOAD_RT_NT': 5, 'LOAD_NT_HT': 6, 'LOAD_RT_WB': 7,
'STORE_RT': 0, 'STORE_NT': 1, 'STORE_HT': 2, 'STORE_BYPASS': 3, 'STORE_LU': 4, 'STORE_RT_NT': 5, 'STORE_NT_HT': 6,
'ATOMIC_RT': 0, 'ATOMIC_NT': 1, 'ATOMIC_RETURN': 1, 'ATOMIC_RT_RETURN': 1, 'ATOMIC_NT_RETURN': 3, 'ATOMIC_CASCADE_RT': 6, 'ATOMIC_CASCADE_NT': 6}.get(m.group(1), 0) if m else 0
m, buf_text = _extract(buf_text, r'\s+scope:SCOPE_(\w+)')
scope_val = {'CU': 0, 'SE': 1, 'DEV': 2, 'SYS': 3}.get(m.group(1), 0) if m else 0
buf_ops = _parse_ops(buf_text)
buf_args = [_op2dsl(o) for o in buf_ops]
vbuf_mods = "".join([f", ioffset={off_val}" if off_val else "", ", offen=1" if offen else "", ", idxen=1" if idxen else "",
f", th={th_val}" if th_val else "", f", scope={scope_val}" if scope_val else "",
", tfe=1" if tfe else ""])
if is_tbuf and fmt_num is not None: vbuf_mods = f", format={fmt_num}" + vbuf_mods
elif is_tbuf: vbuf_mods = ", format=1" + vbuf_mods
else: vbuf_mods = ", format=1" + vbuf_mods
vaddr_idx = 1
if len(buf_ops) > vaddr_idx and buf_ops[vaddr_idx].strip().lower() == 'off': vaddr_val = "v[0]"
else: vaddr_val = buf_args[vaddr_idx] if len(buf_args) > vaddr_idx else "v[0]"
rsrc_idx, soff_idx = (2, 3) if len(buf_ops) > 1 else (1, 2)
rsrc_raw = buf_ops[rsrc_idx].strip() if len(buf_ops) > rsrc_idx else "s[0:3]"
if m := re.match(r's\[(\d+):\d+\]', rsrc_raw.lower()): rsrc_val = m.group(1)
elif m := re.match(r's(\d+)', rsrc_raw.lower()): rsrc_val = m.group(1)
elif m := re.match(r'ttmp\[(\d+):\d+\]', rsrc_raw.lower()): rsrc_val = str(108 + int(m.group(1)))
elif m := re.match(r'ttmp(\d+)', rsrc_raw.lower()): rsrc_val = str(108 + int(m.group(1)))
else: rsrc_val = "0"
soff_raw = buf_ops[soff_idx].strip() if len(buf_ops) > soff_idx else "0"
soff_lower = soff_raw.lower()
if soff_lower == 'm0': soff_val = "RawImm(125)"
elif soff_lower in ('null', 'off'): soff_val = "RawImm(124)"
elif m := re.match(r's(\d+)', soff_lower): soff_val = f"RawImm({m.group(1)})"
else: soff_val = f"RawImm({soff_raw})"
return f"{mn}(vdata={buf_args[0]}, vaddr={vaddr_val}, rsrc={rsrc_val}, soffset={soff_val}{vbuf_mods})"
buf_mods = "".join([f", offset={off_val}" if off_val else "", ", glc=1" if glc else "", ", dlc=1" if dlc else "",
", slc=1" if slc else "", ", tfe=1" if tfe else "", ", offen=1" if offen else "", ", idxen=1" if idxen else ""])
if is_tbuf and fmt_num is not None: buf_mods = f", format={fmt_num}" + buf_mods
vaddr_idx = 1
if len(ops) > vaddr_idx and ops[vaddr_idx].strip().lower() == 'off': vaddr_val = "v[0]"
else: vaddr_val = args[vaddr_idx] if len(args) > vaddr_idx else "v[0]"
srsrc_idx, soff_idx = (2, 3) if len(ops) > 1 else (1, 2)
srsrc_val = args[srsrc_idx] if len(args) > srsrc_idx else "s[0:3]"
soff_val = args[soff_idx] if len(args) > soff_idx else "0"
return f"{mn}(vdata={args[0]}, vaddr={vaddr_val}, srsrc={srsrc_val}, soffset={soff_val}{buf_mods})"
# FLAT/GLOBAL/SCRATCH load/store/atomic - saddr needs RawImm(124) for off/null
# FLAT/GLOBAL/SCRATCH load/store/atomic
def _saddr(a): return 'RawImm(124)' if a in ('OFF', 'NULL') else a
flat_mods = f"{f', offset={off_val}' if off_val else ''}{', glc=1' if glc else ''}{', slc=1' if slc else ''}{', dlc=1' if dlc else ''}"
for pre, flds in [('flat_load','vdst,addr,saddr'), ('global_load','vdst,addr,saddr'), ('scratch_load','vdst,addr,saddr'),
@@ -592,7 +308,12 @@ def get_dsl(text: str) -> str:
# 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")
if ds_off0 is not None or ds_off1 is not None:
off0, off1 = str(ds_off0 or 0), str(ds_off1 or 0)
elif off_val:
off0, off1 = str(int(off_val, 0) & 0xff), str((int(off_val, 0) >> 8) & 0xff)
else:
off0, off1 = "0", "0"
gds_s = ", gds=1" if 'gds' in text.lower().split()[-1:] else ""
off_kw = f", offset0={off0}, offset1={off1}{gds_s}"
if mn == 'ds_nop' or mn in ('ds_gws_sema_v', 'ds_gws_sema_p', 'ds_gws_sema_release_all'): return f"{mn}({off_kw.lstrip(', ')})"
@@ -618,13 +339,33 @@ def get_dsl(text: str) -> str:
lit_s = ""
if mn in ('v_fmaak_f32', 'v_fmaak_f16') and len(args) == 4: lit_s, args = f", literal={args[3].strip()}", args[:3]
elif mn in ('v_fmamk_f32', 'v_fmamk_f16') and len(args) == 4: lit_s, args = f", literal={args[2].strip()}", [args[0], args[1], args[3]]
elif mn in ('s_fmaak_f32',) and len(args) == 4: lit_s, args = f", literal={args[3].strip()}", args[:3]
elif mn in ('s_fmamk_f32',) and len(args) == 4: lit_s, args = f", literal={args[2].strip()}", [args[0], args[1], args[3]]
elif mn in ('v_cndmask_b32', 'v_cndmask_b32_e32') and len(args) == 4 and ops[3].strip().lower() in ('vcc_lo', 'vcc'):
mn, args = 'v_cndmask_b32_e32', args[:3]
# VCC ops cleanup
_SGPR_NAMES = {'vcc_lo': 106, 'vcc_hi': 107, 'vcc': 106, 'null': 124, 'm0': 125, 'exec_lo': 126, 'exec_hi': 127}
vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'}
if mn.replace('_e32', '') in vcc_ops and len(args) >= 5: mn, args = mn.replace('_e32', '') + '_e32', [args[0], args[2], args[3]]
if mn.replace('_e32', '') in vcc_ops and len(args) >= 5:
carry_in = ops[4].strip().lower() if len(ops) > 4 else 'vcc_lo'
carry_out = ops[1].strip().lower() if len(ops) > 1 else 'vcc_lo'
if carry_in in ('vcc_lo', 'vcc') and carry_out in ('vcc_lo', 'vcc'):
mn, args = mn.replace('_e32', '') + '_e32', [args[0], args[2], args[3]]
else:
mn_base = mn.replace('_e32', '').replace('_e64', '')
sdst = _SGPR_NAMES.get(carry_out, 124) if carry_out in _SGPR_NAMES else (int(carry_out[1:]) if carry_out.startswith('s') and carry_out[1:].isdigit() else 124)
src2 = _SGPR_NAMES.get(carry_in, 0) if carry_in in _SGPR_NAMES else (int(carry_in[1:]) if carry_in.startswith('s') and carry_in[1:].isdigit() else 0)
return f"{mn_base}(vdst={args[0]}, sdst=RawImm({sdst}), src0={args[2]}, src1={args[3]}, src2=RawImm({src2}))"
if mn.replace('_e64', '') in vcc_ops and mn.endswith('_e64'): mn = mn.replace('_e64', '')
if mn.startswith('v_cmp') and not mn.endswith('_e64') and len(args) >= 3 and ops[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'): args = args[1:]
if 'cmpx' in mn and mn.endswith('_e64') and len(args) == 2: args = ['RawImm(126)'] + args
if ((mn.startswith('v_cmp') and 'cmpx' not in mn and mn.endswith('_e64')) or mn.startswith('v_s_') or mn in ('v_readlane_b32', 'v_readfirstlane_b32')) and len(args) >= 1:
dst = ops[0].strip().lower()
if dst.startswith('s') and dst[1:].isdigit(): args[0] = f'RawImm({int(dst[1:])})'
elif dst.startswith('s[') and ':' in dst: args[0] = f'RawImm({int(dst[2:].split(":")[0])})'
elif dst.startswith('ttmp') and dst[4:].isdigit(): args[0] = f'RawImm({108 + int(dst[4:])})'
elif dst.startswith('ttmp[') and ':' in dst: args[0] = f'RawImm({108 + int(dst[5:].split(":")[0])})'
elif dst in _SGPR_NAMES: args[0] = f'RawImm({_SGPR_NAMES[dst]})'
fn = mn.replace('.', '_')
if opsel is not None: args = [re.sub(r'\.[hl]$', '', a) for a in args]
@@ -651,16 +392,52 @@ def get_dsl(text: str) -> str:
if neg_lo is not None: all_kw.append(f'neg={neg_lo}')
if neg_hi is not None: all_kw.append(f'neg_hi={neg_hi}')
if 'bvh' in mn and 'intersect_ray' in mn: all_kw.extend(['dmask=15', 'unrm=1', 'r128=1'])
vop3p_ops = {'v_pk_', 'v_dot2', 'v_dot4', 'v_dot8', 'v_wmma', 'v_swmmac'}
is_vop3p = any(mn.startswith(p) for p in vop3p_ops)
is_fma_mix = 'fma_mix' in mn
if opsel_hi_val is not None:
opsel_hi_enc = opsel_hi_val[0] | (opsel_hi_val[1] << 1) if len(opsel_hi_val) >= 2 else opsel_hi_val[0]
opsel_hi2_enc = opsel_hi_val[2] if len(opsel_hi_val) >= 3 else (0 if is_fma_mix else 1)
all_kw.extend([f'opsel_hi={opsel_hi_enc}', f'opsel_hi2={opsel_hi2_enc}'])
elif is_vop3p and not is_fma_mix:
all_kw.extend(['opsel_hi=3', 'opsel_hi2=1'])
if clamp_found:
if arch == 'rdna4': all_kw.append('cm=1')
else: all_kw.append('clmp=1')
a_str, kw_str = ', '.join(args), ', '.join(all_kw)
return f"{fn}({a_str}, {kw_str})" if kw_str and a_str else f"{fn}({kw_str})" if kw_str else f"{fn}({a_str})"
def asm(text: str) -> Inst:
from extra.assembly.amd.autogen import rdna3 as ag
dsl = get_dsl(text)
ns = {n: getattr(ag, n) for n in dir(ag) if not n.startswith('_')}
def _hwreg(id_, offset=0, size=32): return id_ | (offset << 6) | ((size - 1) << 11)
def _sendmsg(id_, op=0, stream=0): return id_ | (op << 4) | (stream << 8)
_HWREG_NAMES = {'HW_REG_MODE': 1, 'HW_REG_STATUS': 2, 'HW_REG_TRAPSTS': 3, 'HW_REG_HW_ID': 4, 'HW_REG_GPR_ALLOC': 5,
'HW_REG_LDS_ALLOC': 6, 'HW_REG_IB_STS': 7, 'HW_REG_PC_LO': 8, 'HW_REG_PC_HI': 9, 'HW_REG_INST_DW0': 10, 'HW_REG_INST_DW1': 11,
'HW_REG_IB_DBG0': 12, 'HW_REG_IB_DBG1': 13, 'HW_REG_FLUSH_IB': 14, 'HW_REG_SH_MEM_BASES': 15, 'HW_REG_SQ_SHADER_TBA_LO': 16,
'HW_REG_SQ_SHADER_TBA_HI': 17, 'HW_REG_SQ_SHADER_TMA_LO': 18, 'HW_REG_SQ_SHADER_TMA_HI': 19, 'HW_REG_FLAT_SCR_LO': 20,
'HW_REG_FLAT_SCR_HI': 21, 'HW_REG_XNACK_MASK': 22, 'HW_REG_HW_ID1': 23, 'HW_REG_HW_ID2': 24, 'HW_REG_POPS_PACKER': 25,
'HW_REG_PERF_SNAPSHOT_DATA': 26, 'HW_REG_PERF_SNAPSHOT_PC_LO': 27, 'HW_REG_PERF_SNAPSHOT_PC_HI': 28, 'HW_REG_SHADER_CYCLES': 29,
'HW_REG_SHADER_CYCLES_HI': 30, 'HW_REG_WAVE_MODE': 31, 'HW_REG_WAVE_SCRATCH_BASE': 32}
_HWREG_NAMES_RDNA4 = {v: k for k, v in HWREG_RDNA4.items()}
_SENDMSG_NAMES = {'MSG_INTERRUPT': 1, 'MSG_GS': 2, 'MSG_GS_DONE': 3, 'MSG_SAVEWAVE': 4, 'MSG_STALL_WAVE_GEN': 5,
'MSG_HALT_WAVES': 6, 'MSG_ORDERED_PS_DONE': 7, 'MSG_EARLY_PRIM_DEALLOC': 8, 'MSG_GS_ALLOC_REQ': 9, 'MSG_GET_DOORBELL': 10,
'MSG_GET_DDID': 11, 'MSG_HS_TESSFACTOR': 2, 'MSG_DEALLOC_VGPRS': 10, 'MSG_RTN_GET_DOORBELL': 128, 'MSG_RTN_GET_DDID': 129,
'MSG_RTN_GET_TMA': 130, 'MSG_RTN_GET_REALTIME': 131, 'MSG_RTN_SAVE_WAVE': 132, 'MSG_RTN_GET_TBA': 133,
'MSG_RTN_GET_TBA_TO_PC': 134, 'MSG_RTN_GET_SE_AID_ID': 135}
def asm(text: str, arch: str = "rdna3") -> Inst:
dsl = get_dsl(text, arch)
if arch == "rdna4":
ns = {n: getattr(rdna4_ins, n) for n in dir(rdna4_ins) if not n.startswith('_')}
hwreg_names = _HWREG_NAMES_RDNA4
else:
ns = {n: getattr(ins, n) for n in dir(ins) if not n.startswith('_')}
hwreg_names = _HWREG_NAMES
def hwreg(id_, offset=0, size=32): return _hwreg(hwreg_names.get(id_, id_) if isinstance(id_, str) else id_, offset, size)
def sendmsg(id_, op=0, stream=0): return _sendmsg(_SENDMSG_NAMES.get(id_, id_) if isinstance(id_, str) else id_, op, stream)
ns.update({'s': s, 'v': v, 'ttmp': ttmp, 'abs': abs, 'RawImm': RawImm, 'SrcMod': SrcMod, 'VGPR': VGPR, 'SGPR': SGPR, 'TTMP': TTMP,
'VCC_LO': VCC_LO, 'VCC_HI': VCC_HI, 'VCC': VCC, 'EXEC_LO': EXEC_LO, 'EXEC_HI': EXEC_HI, 'EXEC': EXEC, 'SCC': SCC, 'M0': M0, 'NULL': NULL, 'OFF': OFF})
'VCC_LO': VCC_LO, 'VCC_HI': VCC_HI, 'VCC': VCC, 'EXEC_LO': EXEC_LO, 'EXEC_HI': EXEC_HI, 'EXEC': EXEC, 'SCC': SCC, 'M0': M0, 'NULL': NULL, 'OFF': OFF,
'hwreg': hwreg, 'sendmsg': sendmsg, **{k: k for k in hwreg_names}, **{k: k for k in _SENDMSG_NAMES}})
try: return eval(dsl, ns)
except NameError:
if m := re.match(r'^(v_\w+)(\(.*\))$', dsl): return eval(f"{m.group(1)}_e32{m.group(2)}", ns)
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
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
+62
View File
@@ -0,0 +1,62 @@
# Instruction format detection and decoding
from __future__ import annotations
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.autogen.rdna3.ins import VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP2 as R4_VOP2, VOP3 as R4_VOP3, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
SMEM as R4_SMEM, DS as R4_DS, VBUFFER as R4_VBUFFER, VEXPORT as R4_VEXPORT)
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP2 as C_VOP2, VOPC as C_VOPC, VOP3A, VOP3B, VOP3P as C_VOP3P,
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
FLAT as C_FLAT, MUBUF as C_MUBUF, MTBUF as C_MTBUF, SDWA, DPP)
def _matches_encoding(word: int, cls: type[Inst]) -> bool:
"""Check if word matches the encoding pattern of an instruction class."""
if cls._encoding is None: return False
bf, val = cls._encoding
return ((word >> bf.lo) & bf.mask()) == val
# Order matters: more specific encodings first, VOP2 last (it's a catch-all for bit31=0)
_RDNA_FORMATS_64 = [VOPD, VOP3P, VINTERP, VOP3, DS, FLAT, MUBUF, MTBUF, MIMG, SMEM, EXP]
_RDNA_FORMATS_32 = [SOP1, SOPC, SOPP, SOPK, VOPC, VOP1, SOP2, VOP2] # SOP2/VOP2 are catch-alls
_CDNA_FORMATS_64 = [C_VOP3P, VOP3A, C_DS, C_FLAT, C_MUBUF, C_MTBUF, C_SMEM]
_CDNA_FORMATS_32 = [SDWA, DPP, C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_VOPC, C_VOP1, C_SOP2, C_VOP2]
_CDNA_VOP3B_OPS = {281, 282, 283, 284, 285, 286, 480, 481, 488, 489} # VOP3B opcodes
_RDNA4_FORMATS_64 = [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3, R4_DS, R4_VBUFFER, R4_SMEM, R4_VEXPORT]
_RDNA4_FORMATS_32 = [R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_VOPC, R4_VOP1, R4_SOP2, R4_VOP2]
_RDNA4_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
"""Detect instruction format from machine code bytes."""
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
word = int.from_bytes(data[:4], 'little')
if arch == "cdna":
if (word >> 30) == 0b11:
for cls in _CDNA_FORMATS_64:
if _matches_encoding(word, cls):
return VOP3B if cls is VOP3A and ((word >> 16) & 0x3ff) in _CDNA_VOP3B_OPS else cls
raise ValueError(f"unknown CDNA 64-bit format word={word:#010x}")
for cls in _CDNA_FORMATS_32:
if _matches_encoding(word, cls): return cls
raise ValueError(f"unknown CDNA 32-bit format word={word:#010x}")
if arch == "rdna4":
if (word >> 30) == 0b11:
for cls in _RDNA4_FORMATS_64:
if _matches_encoding(word, cls):
return R4_VOP3SD if cls is R4_VOP3 and ((word >> 16) & 0x3ff) in _RDNA4_VOP3SD_OPS else cls
raise ValueError(f"unknown RDNA4 64-bit format word={word:#010x}")
for cls in _RDNA4_FORMATS_32:
if _matches_encoding(word, cls): return cls
raise ValueError(f"unknown RDNA4 32-bit format word={word:#010x}")
# RDNA3 (default)
if (word >> 30) == 0b11:
for cls in _RDNA_FORMATS_64:
if _matches_encoding(word, cls):
return VOP3SD if cls is VOP3 and ((word >> 16) & 0x3ff) in Inst._VOP3SD_OPS else cls
raise ValueError(f"unknown 64-bit format word={word:#010x}")
for cls in _RDNA_FORMATS_32:
if _matches_encoding(word, cls): return cls
raise ValueError(f"unknown 32-bit format word={word:#010x}")
def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
"""Decode machine code bytes into an instruction."""
return detect_format(data, arch).from_bytes(data)
+745
View File
@@ -0,0 +1,745 @@
# RDNA3/RDNA4/CDNA disassembler
from __future__ import annotations
import re
from extra.assembly.amd.dsl import Inst, decode_src, SPECIAL_GPRS, SPECIAL_GPRS_CDNA, SPECIAL_PAIRS, SPECIAL_PAIRS_CDNA
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, VOPD, VINTERP, SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, DS, FLAT, MUBUF, MTBUF, MIMG, EXP,
VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOPDOp, SOP1Op, SOPKOp, SOPPOp, SMEMOp, DSOp, MUBUFOp)
from extra.assembly.amd.autogen.rdna3.enum import BufFmt
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP2 as R4_VOP2, VOP3 as R4_VOP3, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOPC as R4_SOPC, SOPK as R4_SOPK, SOPP as R4_SOPP,
SMEM as R4_SMEM, DS as R4_DS, VBUFFER as R4_VBUFFER, VEXPORT as R4_VEXPORT, VOPDOp as R4_VOPDOp)
from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, MUBUF as C_MUBUF, MTBUF as C_MTBUF
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
# ═══════════════════════════════════════════════════════════════════════════════
# CONSTANTS
# ═══════════════════════════════════════════════════════════════════════════════
HWREG = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI', 22: 'HW_REG_XNACK_MASK',
23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
HWREG_RDNA4 = {1: 'HW_REG_WAVE_MODE', 2: 'HW_REG_WAVE_STATUS', 4: 'HW_REG_WAVE_STATE_PRIV', 5: 'HW_REG_WAVE_GPR_ALLOC',
6: 'HW_REG_WAVE_LDS_ALLOC', 7: 'HW_REG_IB_STS', 10: 'HW_REG_PERF_SNAPSHOT_DATA', 11: 'HW_REG_PERF_SNAPSHOT_PC_LO',
12: 'HW_REG_PERF_SNAPSHOT_PC_HI', 15: 'HW_REG_PERF_SNAPSHOT_DATA1', 16: 'HW_REG_PERF_SNAPSHOT_DATA2',
17: 'HW_REG_WAVE_EXCP_FLAG_PRIV', 18: 'HW_REG_WAVE_EXCP_FLAG_USER', 19: 'HW_REG_WAVE_TRAP_CTRL',
20: 'HW_REG_WAVE_SCRATCH_BASE_LO', 21: 'HW_REG_WAVE_SCRATCH_BASE_HI', 23: 'HW_REG_WAVE_HW_ID1',
24: 'HW_REG_WAVE_HW_ID2', 26: 'HW_REG_WAVE_SCHED_MODE', 29: 'HW_REG_SHADER_CYCLES_LO',
30: 'HW_REG_SHADER_CYCLES_HI', 31: 'HW_REG_WAVE_DVGPR_ALLOC_LO', 32: 'HW_REG_WAVE_DVGPR_ALLOC_HI'}
MSG = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA',
134: 'MSG_RTN_GET_TBA_TO_PC', 135: 'MSG_RTN_GET_SE_AID_ID'}
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
# ═══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{b}" if n == 1 else f"{p}[{b}:{b+n-1}]"
def _sreg(b: int, n: int = 1) -> str: return _reg("s", b, n)
def _vreg(b: int, n: int = 1) -> str: return _reg("v", b, n)
def _areg(b: int, n: int = 1) -> str: return _reg("a", b, n) # accumulator registers for GFX90a
def _ttmp(b: int, n: int = 1) -> str: return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _sreg_or_ttmp(b: int, n: int = 1) -> str: return _ttmp(b, n) or _sreg(b, n)
def _fmt_sdst(v: int, n: int = 1, cdna: bool = False) -> str:
if t := _ttmp(v, n): return t
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
if n > 1: return pairs.get(v) or gprs.get(v) or _sreg(v, n) # also check gprs for null/m0
return gprs.get(v, f"s{v}")
def _fmt_src(v: int, n: int = 1, cdna: bool = False) -> str:
if n == 1: return decode_src(v, cdna)
if v >= 256: return _vreg(v - 256, n)
if v <= 101: return _sreg(v, n) # s0-s101 can be pairs, but 102+ are special on CDNA
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
if n == 2 and v in pairs: return pairs[v]
if v <= 105: return _sreg(v, n) # s102-s105 regular pairs for RDNA
if t := _ttmp(v, n): return t
return decode_src(v, cdna)
def _fmt_v16(v: int, base: int = 256, hi_thresh: int = 384) -> str:
return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
def _src16(inst, v: int) -> str: return _fmt_v16(v) if v >= 256 else inst.lit(v) # format 16-bit src: vgpr.h/l or literal
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool) -> str:
"""Format VOP3 source operand with modifiers."""
if v == 255: s = inst.lit(v) # literal constant takes priority
elif n > 1: s = _fmt_src(v, n)
elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else f"v{v - 256}.l"
elif v == 253: s = "src_scc" # VOP3 sources use src_scc not scc
else: s = inst.lit(v)
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
"""Format op_sel modifier string."""
if not need: return ""
dst_hi = (opsel >> 3) & 1
if n == 1: return f" op_sel:[{opsel & 1},{dst_hi}]"
if n == 2: return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{dst_hi}]"
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{dst_hi}]"
# ═══════════════════════════════════════════════════════════════════════════════
# DISASSEMBLER
# ═══════════════════════════════════════════════════════════════════════════════
def _disasm_vop1(inst: VOP1) -> str:
name, cdna = inst.op_name.lower() or f'vop1_op_{inst.op}', _is_cdna(inst)
suf = "" if cdna else "_e32"
if name in ('v_nop', 'v_pipeflush', 'v_clrexcp'): return name # no operands
if 'readfirstlane' in name:
src = f"v{inst.src0 - 256}" if inst.src0 >= 256 else decode_src(inst.src0, cdna)
return f"{name} {_fmt_sdst(inst.vdst, 1, cdna)}, {src}"
# 16-bit dst: uses .h/.l suffix for RDNA (CDNA uses plain vN)
parts = name.split('_')
is_16d = not cdna and (any(p in ('f16','i16','u16','b16') for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in ('f16','i16','u16','b16') and 'cvt' not in name))
# v_cvt_pk_f32_fp8 and v_cvt_pk_f32_bf8 output to 2 VGPRs, and take 16-bit src
is_pk_fp8 = 'cvt_pk_f32_fp8' in name or 'cvt_pk_f32_bf8' in name
dregs = 2 if is_pk_fp8 else inst.dst_regs()
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _fmt_v16(inst.vdst, 0, 128) if is_16d else f"v{inst.vdst}"
src = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, inst.src_regs(0), cdna) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if not cdna and (inst.is_src_16(0) or is_pk_fp8) and 'sat_pk' not in name else inst.lit(inst.src0)
return f"{name}{suf} {dst}, {src}"
_VOP2_CARRY_OUT = {'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'} # carry out only
_VOP2_CARRY_INOUT = {'v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'} # carry in and out (CDNA)
_VOP2_CARRY_INOUT_RDNA = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'} # carry in and out (RDNA)
def _disasm_vop2(inst: VOP2) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if cdna: name = _CDNA_DISASM_ALIASES.get(name, name) # apply CDNA aliases
suf = "" if cdna or (not cdna and inst.op == VOP2Op.V_DOT2ACC_F32_F16) else "_e32"
lit = getattr(inst, '_literal', None)
is16 = not cdna and inst.is_16bit()
# fmaak/madak: dst = src0 * vsrc1 + K, fmamk/madmk: dst = src0 * K + vsrc1
if 'fmaak' in name or 'madak' in name or (not cdna and inst.op in (VOP2Op.V_FMAAK_F32, VOP2Op.V_FMAAK_F16)):
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}, 0x{lit:x}"
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}, 0x{lit:x}"
if 'fmamk' in name or 'madmk' in name or (not cdna and inst.op in (VOP2Op.V_FMAMK_F32, VOP2Op.V_FMAMK_F16)):
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, 0x{lit:x}, {_fmt_v16(inst.vsrc1, 0, 128)}"
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, 0x{lit:x}, v{inst.vsrc1}"
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst, 0, 128)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1, 0, 128)}"
vcc = "vcc" if cdna else "vcc_lo"
# CDNA carry ops output vcc after vdst
if cdna and name in _VOP2_CARRY_OUT: return f"{name}{suf} v{inst.vdst}, {vcc}, {inst.lit(inst.src0)}, v{inst.vsrc1}"
if cdna and name in _VOP2_CARRY_INOUT: return f"{name}{suf} v{inst.vdst}, {vcc}, {inst.lit(inst.src0)}, v{inst.vsrc1}, {vcc}"
# RDNA carry-in/out ops: v_add_co_ci_u32, etc. - format: vdst, vcc_lo, src0, vsrc1, vcc_lo
if not cdna and name in _VOP2_CARRY_INOUT_RDNA: return f"{name}{suf} v{inst.vdst}, {vcc}, {inst.lit(inst.src0)}, v{inst.vsrc1}, {vcc}"
# Handle 64-bit register operands (v_add_f64, v_mul_f64, etc.)
dn, sn0, sn1 = inst.dst_regs(), inst.src_regs(0), inst.src_regs(1)
if dn > 1 or sn0 > 1 or sn1 > 1:
dst = _vreg(inst.vdst, dn) if dn > 1 else f"v{inst.vdst}"
src0 = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, sn0, cdna)
src1 = _vreg(inst.vsrc1, sn1) if sn1 > 1 else f"v{inst.vsrc1}"
return f"{name} {dst}, {src0}, {src1}"
return f"{name}{suf} v{inst.vdst}, {inst.lit(inst.src0)}, v{inst.vsrc1}" + (f", {vcc}" if name == 'v_cndmask_b32' else "")
def _disasm_vopc(inst: VOPC) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if cdna:
s0 = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, inst.src_regs(0), cdna)
s1 = _vreg(inst.vsrc1, inst.src_regs(1)) if inst.src_regs(1) > 1 else f"v{inst.vsrc1}"
return f"{name} vcc, {s0}, {s1}" # CDNA VOPC always outputs vcc
# RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
has_vcc = 'cmpx' not in name
s0 = inst.lit(inst.src0) if inst.src0 == 255 else _fmt_src(inst.src0, inst.src_regs(0)) if inst.src_regs(0) > 1 else _src16(inst, inst.src0) if inst.is_16bit() else inst.lit(inst.src0)
s1 = _vreg(inst.vsrc1, inst.src_regs(1)) if inst.src_regs(1) > 1 else _fmt_v16(inst.vsrc1, 0, 128) if inst.is_16bit() else f"v{inst.vsrc1}"
return f"{name}_e32 vcc_lo, {s0}, {s1}" if has_vcc else f"{name}_e32 {s0}, {s1}"
NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE, SOPPOp.S_TTRACEDATA}
_CDNA_NO_ARG_SOPP = {'s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_nop', 's_sethalt', 's_sleep',
's_setprio', 's_trap', 's_incperflevel', 's_decperflevel', 's_sendmsg', 's_sendmsghalt'}
def _disasm_sopp(inst: SOPP) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if cdna:
if name == 's_endpgm': return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
if name in ('s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata'): return name
if name == 's_waitcnt':
vm, lgkm, exp = inst.simm16 & 0xf, (inst.simm16 >> 8) & 0x3f, (inst.simm16 >> 4) & 0x7
p = [f"vmcnt({vm})" if vm != 0xf else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
return f"{name} 0x{inst.simm16:x}" if inst.simm16 else name
# RDNA
if inst.op in NO_ARG_SOPP: return name
if inst.op == SOPPOp.S_ENDPGM: return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
if inst.op == SOPPOp.S_WAITCNT:
vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if inst.op == SOPPOp.S_DELAY_ALU:
deps, skips = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3'], ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
return f"{name} {inst.simm16}" if name.startswith(('s_cbranch', 's_branch')) else f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if inst.op in (SMEMOp.S_GL1_INV, SMEMOp.S_DCACHE_INV): return name
soe, imm = getattr(inst, 'soe', 0), getattr(inst, 'imm', 1)
is_rdna4 = 'rdna4' in inst.__class__.__module__
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
if cdna:
if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
elif imm: off_s = f"0x{offset:x}"
elif offset < 256: off_s = decode_src(offset, cdna)
else: off_s = decode_src(inst.soffset, cdna)
elif offset and inst.soffset != 124: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
elif offset: off_s = f"0x{offset:x}"
else: off_s = decode_src(inst.soffset, cdna)
is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
sbase_idx, sbase_count = inst.sbase * 2, 4 if is_buffer else 2
sbase_str = _fmt_src(sbase_idx, sbase_count, cdna) if sbase_count == 2 else _sreg(sbase_idx, sbase_count) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_count)
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {inst.sdata}, {sbase_str}, {off_s}"
if 'prefetch' in name:
off = getattr(inst, 'ioffset', inst.offset)
if off >= 0x800000: off = off - 0x1000000
off_s = f"0x{off:x}" if off > 255 else str(off)
soff_s = decode_src(inst.soffset, cdna) if inst.soffset != 124 else "null"
if 'pc_rel' in name: return f"{name} {off_s}, {soff_s}, {inst.sdata}"
return f"{name} {sbase_str}, {off_s}, {soff_s}, {inst.sdata}"
th, scope = getattr(inst, 'th', 0), getattr(inst, 'scope', 0)
if th or scope:
th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
return f"{name} {_fmt_sdst(inst.sdata, inst.dst_regs(), cdna)}, {sbase_str}, {off_s}{mods}"
return f"{name} {_fmt_sdst(inst.sdata, inst.dst_regs(), cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
def _disasm_flat(inst: FLAT) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
off_val = inst.offset if seg == 'flat' else (inst.offset if inst.offset < 4096 else inst.offset - 8192)
w = inst.dst_regs() * (2 if '_x2' in name else 1) * (2 if 'cmpswap' in name else 1)
off_s = f" offset:{off_val}" if off_val else ""
if cdna: mods = f"{off_s}{' glc' if inst.sc0 else ''}{' slc' if inst.nt else ''}"
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
if seg == 'flat' or inst.saddr == 0x7F: saddr_s = ""
elif inst.saddr == 124: saddr_s = ", off"
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
elif inst.saddr in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS): saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[inst.saddr]}"
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if inst.saddr < 106 else decode_src(inst.saddr, cdna)}"
if 'addtid' in name: return f"{instr} {'a' if acc else 'v'}{inst.data if 'store' in name else inst.vdst}{saddr_s}{mods}"
if cdna: addr_w = 1 if seg == 'scratch' else 2
else: addr_w = 1 if seg == 'scratch' or (inst.saddr not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
glc_or_sc0 = inst.sc0 if cdna else inst.glc
if 'atomic' in name:
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
def _disasm_ds(inst: DS) -> str:
op, name = inst.op, inst.op_name.lower()
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
rp = 'a' if acc else 'v'
gds = " gds" if inst.gds else ""
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
off2 = (" offset0:" + str(inst.offset0) if inst.offset0 else "") + (" offset1:" + str(inst.offset1) if inst.offset1 else "")
w = inst.dst_regs()
d0, d1, dst, addr = reg_fn(inst.data0, w), reg_fn(inst.data1, w), reg_fn(inst.vdst, w), f"v{inst.addr}"
if op == DSOp.DS_NOP: return name
if op == DSOp.DS_BVH_STACK_RTN_B32: return f"{name} v{inst.vdst}, {addr}, v{inst.data0}, {_vreg(inst.data1, 4)}{off}{gds}"
if 'bvh_stack_push' in name:
d1_regs = 8 if 'push8' in name else 4
vdst_regs = 2 if 'pop2' in name else 1
vdst_s = _vreg(inst.vdst, vdst_regs) if vdst_regs > 1 else f"v{inst.vdst}"
return f"{name} {vdst_s}, {addr}, v{inst.data0}, {_vreg(inst.data1, d1_regs)}{off}{gds}"
if 'gws_sema' in name and op != DSOp.DS_GWS_SEMA_BR: return f"{name}{off}{gds}"
if 'gws_' in name: return f"{name} {addr}{off}{gds}"
if op in (DSOp.DS_CONSUME, DSOp.DS_APPEND): return f"{name} {rp}{inst.vdst}{off}{gds}"
if 'gs_reg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {rp}{inst.data0}{off}{gds}"
if '2addr' in name:
if 'load' in name: return f"{name} {reg_fn(inst.vdst, w*2)}, {addr}{off2}{gds}"
if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
return f"{name} {reg_fn(inst.vdst, w*2)}, {addr}, {d0}, {d1}{off2}{gds}"
if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
if 'read2' in name: return f"{name} {reg_fn(inst.vdst, w*2)}, {addr}{off2}{gds}"
if 'load' in name: return f"{name} {rp}{inst.vdst}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
if 'store' in name and not _has(name, 'cmp', 'xchg'):
return f"{name} {rp}{inst.data0}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
if 'swizzle' in name or op == DSOp.DS_ORDERED_COUNT: return f"{name} {rp}{inst.vdst}, {addr}{off}{gds}"
if 'permute' in name: return f"{name} {rp}{inst.vdst}, {addr}, {rp}{inst.data0}{off}{gds}"
if 'condxchg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {addr}, {reg_fn(inst.data0, 2)}{off}{gds}"
if _has(name, 'cmpstore', 'mskor', 'wrap'):
return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
def _disasm_vop3(inst: VOP3) -> str:
op, name = inst.op, inst.op_name.lower()
# RDNA4 v_s_* scalar VOP3 instructions
if name.startswith('v_s_'):
src = inst.lit(inst.src0) if inst.src0 == 255 else ("src_scc" if inst.src0 == 253 else _fmt_src(inst.src0, inst.src_regs(0)))
if inst.neg & 1: src = f"-{src}"
if inst.abs & 1: src = f"|{src}|"
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
return f"{name} s{inst.vdst}, {src}" + (" clamp" if clamp else "") + _omod(inst.omod)
# VOP3SD (shared encoding)
if isinstance(op, VOP3SDOp):
sdst = (inst.clmp << 7) | (inst.opsel << 3) | inst.abs
def src(v, neg, n):
s = inst.lit(v) if v == 255 else ("src_scc" if v == 253 else (_fmt_src(v, n) if n > 1 else inst.lit(v)))
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, inst.src_regs(2))
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
return f"{name} {dst}, {_fmt_sdst(sdst, 1)}, {srcs}" + _omod(inst.omod)
# Detect 16-bit operand sizes
is16_d = is16_s = is16_s2 = False
if 'cvt_pk' in name: is16_s = name.endswith('16')
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', name):
is16_d, is16_s = _has(m.group(1), 'f16','i16','u16','b16'), _has(m.group(2), 'f16','i16','u16','b16')
is16_s2 = is16_s
elif re.match(r'v_mad_[iu]32_[iu]16', name): is16_s = True
elif 'pack_b32' in name: is16_s = is16_s2 = True
elif 'sat_pk' in name: is16_d = True
else: is16_d = is16_s = is16_s2 = inst.is_16bit()
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, inst.src_regs(0), is16_s)
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, inst.src_regs(1), is16_s)
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, inst.src_regs(2), is16_s2)
# Destination
dn = inst.dst_regs()
if op == VOP3Op.V_READLANE_B32: dst = _fmt_sdst(inst.vdst, 1)
elif dn > 1: dst = _vreg(inst.vdst, dn)
elif is16_d: dst = f"v{inst.vdst}.h" if (inst.opsel & 8) else f"v{inst.vdst}.l"
else: dst = f"v{inst.vdst}"
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
cl, om = " clamp" if clamp else "", _omod(inst.omod)
nonvgpr_opsel = (inst.src0 < 256 and (inst.opsel & 1)) or (inst.src1 < 256 and (inst.opsel & 2)) or (inst.src2 < 256 and (inst.opsel & 4))
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
if inst.op < 256: # VOPC
return f"{name}_e64 {s0}, {s1}{cl}" if name.startswith('v_cmpx') else f"{name}_e64 {_fmt_sdst(inst.vdst, 1)}, {s0}, {s1}{cl}"
if inst.op < 384: # VOP2
n = inst.num_srcs()
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}_e64 {dst}, {s0}, {s1}{os}{cl}{om}"
if inst.op < 512: # VOP1
if re.match(r'v_cvt_f32_(bf|fp)8', name) and inst.opsel:
os = f" byte_sel:{((inst.opsel & 1) << 1) | ((inst.opsel >> 1) & 1)}"
else:
os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)
return f"{name}_e64" if op in (VOP3Op.V_NOP, VOP3Op.V_PIPEFLUSH) else f"{name}_e64 {dst}, {s0}{os}{cl}{om}"
# Native VOP3
n = inst.num_srcs()
if 'cvt_sr' in name and inst.opsel:
os = f" byte_sel:{inst.opsel >> 2}"
else:
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
def _disasm_vop3sd(inst: VOP3SD) -> str:
name = inst.op_name.lower()
src2_n = 2 if '_co_' in name and '64' in name else inst.src_regs(2)
def src(v, neg, n):
s = inst.lit(v) if v == 255 else ("src_scc" if v == 253 else (_fmt_src(v, n) if n > 1 else inst.lit(v)))
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
s0, s1, s2 = src(inst.src0, inst.neg & 1, inst.src_regs(0)), src(inst.src1, inst.neg & 2, inst.src_regs(1)), src(inst.src2, inst.neg & 4, src2_n)
dst = _vreg(inst.vdst, inst.dst_regs()) if inst.dst_regs() > 1 else f"v{inst.vdst}"
srcs = f"{s0}, {s1}, {s2}" if inst.num_srcs() == 3 else f"{s0}, {s1}"
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
return f"{name} {dst}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if clamp else ''}{_omod(inst.omod)}"
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal or inst.literal
is_rdna4 = 'rdna4' in inst.__class__.__module__
op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
vdst_y, nx, ny = (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
def half(n, vd, s0, vs1):
if 'mov' in n: return f"{n} v{vd}, {inst.lit(s0)}"
if 'fmamk' in n and lit: return f"{n} v{vd}, {inst.lit(s0)}, 0x{lit:x}, v{vs1}"
if 'fmaak' in n and lit: return f"{n} v{vd}, {inst.lit(s0)}, v{vs1}, 0x{lit:x}"
return f"{n} v{vd}, {inst.lit(s0)}, v{vs1}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
def _swmmac_regs(name: str) -> tuple[int, int, int, int]:
"""Return (dst, src0, src1, src2) register counts for SWMMAC instructions."""
if 'f16_16x16x32' in name or 'bf16_16x16x32' in name: return (4, 4, 8, 1)
if 'f32_16x16x32_f16' in name or 'f32_16x16x32_bf16' in name: return (8, 4, 8, 1)
if 'i32_16x16x32_iu4' in name: return (8, 1, 2, 1)
if 'i32_16x16x64_iu4' in name: return (8, 2, 4, 1)
if 'i32_16x16x32_iu8' in name or 'f32_16x16x32_fp8' in name or 'f32_16x16x32_bf8' in name: return (8, 2, 4, 1)
return (8, 8, 8, 8)
def _disasm_vop3p(inst: VOP3P) -> str:
name = inst.op_name.lower()
is_wmma, is_swmmac, n, is_fma_mix = 'wmma' in name, 'swmmac' in name, inst.num_srcs(), 'fma_mix' in name
def get_src(v, sc): return inst.lit(v) if v == 255 else _fmt_src(v, sc)
if is_swmmac:
dn, s0n, s1n, s2n = _swmmac_regs(name)
src0, src1, src2, dst = get_src(inst.src0, s0n), get_src(inst.src1, s1n), get_src(inst.src2, s2n), _vreg(inst.vdst, dn)
elif is_wmma:
is_rdna4_wmma = 'rdna4' in inst.__class__.__module__
sc = 1 if '16x16x16_iu4' in name else 2 if ('iu4' in name or 'iu8' in name or 'fp8' in name or 'bf8' in name) else 4
if not is_rdna4_wmma: sc *= 2
dc = 8 if not is_rdna4_wmma else (4 if ('f16_16x16' in name or 'bf16_16x16' in name) and 'f32' not in name else 8)
src0, src1, src2, dst = get_src(inst.src0, sc), get_src(inst.src1, sc), get_src(inst.src2, dc), _vreg(inst.vdst, dc)
else: src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), f"v{inst.vdst}"
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
clamp = inst.cm if 'cm' in inst._fields else getattr(inst, 'clmp', 0)
if is_fma_mix:
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else [])
elif is_swmmac:
mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
else:
opsel_hi_default = 7 if n == 3 else 3
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else []) + \
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
def _disasm_buf(inst: MUBUF | MTBUF) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
if cdna and name in ('buffer_wbl2', 'buffer_inv'): return name
if not cdna and inst.op in (MUBUFOp.BUFFER_GL0_INV, MUBUFOp.BUFFER_GL1_INV): return name
w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
{'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], 1)
if hasattr(inst, 'tfe') and inst.tfe: w += 1
vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else f"v{inst.vaddr}" if inst.offen or inst.idxen else "off"
srsrc = _sreg_or_ttmp(inst.srsrc*4, 4)
is_mtbuf = isinstance(inst, MTBUF) or isinstance(inst, C_MTBUF)
if is_mtbuf:
dfmt, nfmt = inst.format & 0xf, (inst.format >> 4) & 0x7
if acc: fmt_s = f" dfmt:{dfmt}, nfmt:{nfmt},"
elif not cdna: fmt_s = f" format:{inst.format}" if inst.format else ""
else:
dfmt_names = ['INVALID', '8', '16', '8_8', '32', '16_16', '10_11_11', '11_11_10', '10_10_10_2', '2_10_10_10', '8_8_8_8', '32_32', '16_16_16_16', '32_32_32', '32_32_32_32', 'RESERVED_15']
nfmt_names = ['UNORM', 'SNORM', 'USCALED', 'SSCALED', 'UINT', 'SINT', 'RESERVED_6', 'FLOAT']
if dfmt == 1 and nfmt == 0: fmt_s = ""
elif nfmt == 0: fmt_s = f" format:[BUF_DATA_FORMAT_{dfmt_names[dfmt]}]"
elif dfmt == 1: fmt_s = f" format:[BUF_NUM_FORMAT_{nfmt_names[nfmt]}]"
else: fmt_s = f" format:[BUF_DATA_FORMAT_{dfmt_names[dfmt]},BUF_NUM_FORMAT_{nfmt_names[nfmt]}]"
else: fmt_s = ""
if cdna: mods = [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.sc0,"glc"),(inst.nt,"slc"),(inst.sc1,"sc1")] if c]
else: mods = [m for c, m in [(inst.idxen,"idxen"),(inst.offen,"offen"),(inst.offset,f"offset:{inst.offset}"),(inst.glc,"glc"),(inst.dlc,"dlc"),(inst.slc,"slc"),(inst.tfe,"tfe")] if c]
soffset_s = decode_src(inst.soffset, cdna)
if cdna and not acc and is_mtbuf: return f"{name} {reg_fn(inst.vdata, w)}, {vaddr}, {srsrc}, {soffset_s}{fmt_s}{' ' + ' '.join(mods) if mods else ''}"
return f"{name} {reg_fn(inst.vdata, w)}, {vaddr}, {srsrc},{fmt_s} {soffset_s}{' ' + ' '.join(mods) if mods else ''}"
def _mimg_vaddr_width(name: str, dim: int, a16: bool) -> int:
base = [1, 2, 3, 3, 2, 3, 3, 4][dim]
grad = [1, 2, 3, 2, 1, 2, 2, 2][dim]
if 'get_resinfo' in name: return 1
packed, unpacked = 0, 0
if '_mip' in name: packed += 1
elif 'sample' in name or 'gather' in name:
if '_o' in name: unpacked += 1
if re.search(r'_c(_|$)', name): unpacked += 1
if '_d' in name: unpacked += (grad + 1) & ~1 if '_g16' in name else grad*2
if '_b' in name: unpacked += 1
if '_l' in name and '_cl' not in name and '_lz' not in name: packed += 1
if '_cl' in name: packed += 1
return (base + packed + 1) // 2 + unpacked if a16 else base + packed + unpacked
def _disasm_mimg(inst: MIMG) -> str:
name = inst.op_name.lower()
srsrc_base = inst.srsrc * 4
srsrc_str = _sreg_or_ttmp(srsrc_base, 8)
if 'bvh' in name:
vaddr = (9 if '64' in name else 8) if inst.a16 else (12 if '64' in name else 11)
return f"{name} {_vreg(inst.vdata, 4)}, {_vreg(inst.vaddr, vaddr)}, {_sreg_or_ttmp(srsrc_base, 4)}{' a16' if inst.a16 else ''}"
vdata = 4 if 'gather4' in name or 'msaa_load' in name else (bin(inst.dmask).count('1') or 1)
if inst.d16: vdata = (vdata + 1) // 2
if inst.tfe: vdata += 1
dim_names = ['1d', '2d', '3d', 'cube', '1d_array', '2d_array', '2d_msaa', '2d_msaa_array']
dim = dim_names[inst.dim] if inst.dim < len(dim_names) else f"dim_{inst.dim}"
vaddr = _mimg_vaddr_width(name, inst.dim, inst.a16)
vaddr_str = f"v{inst.vaddr}" if vaddr == 1 else _vreg(inst.vaddr, vaddr)
mods = [f"dmask:0x{inst.dmask:x}"] if inst.dmask and (inst.dmask != 15 or 'atomic' in name) else []
mods.append(f"dim:SQ_RSRC_IMG_{dim.upper()}")
for flag, mod in [(inst.unrm,"unorm"),(inst.glc,"glc"),(inst.slc,"slc"),(inst.dlc,"dlc"),(inst.r128,"r128"),
(inst.a16,"a16"),(inst.tfe,"tfe"),(inst.lwe,"lwe"),(inst.d16,"d16")]:
if flag: mods.append(mod)
ssamp_str = ""
if 'sample' in name or 'gather' in name or 'get_lod' in name:
ssamp_str = ", " + _sreg_or_ttmp(inst.ssamp * 4, 4)
return f"{name} {_vreg(inst.vdata, vdata)}, {vaddr_str}, {srsrc_str}{ssamp_str} {' '.join(mods)}"
def _disasm_sop1(inst: SOP1) -> str:
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
src = inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0), cdna)
if not cdna:
if op == SOP1Op.S_GETPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
if op in (SOP1Op.S_SETPC_B64, SOP1Op.S_RFE_B64): return f"{name} {src}"
if op == SOP1Op.S_SWAPPC_B64: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {src}"
if op in (SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64): return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs())}, sendmsg({MSG.get(inst.ssrc0, str(inst.ssrc0))})"
sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST', 'S_BARRIER_INIT', 'S_BARRIER_JOIN')
if inst.op_name in sop1_src_only: return f"{name} {src}"
return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs(), cdna)}, {src}"
def _disasm_sop2(inst: SOP2) -> str:
cdna, name = _is_cdna(inst), inst.op_name.lower()
lit = getattr(inst, '_literal', None)
s0 = inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0), cdna)
s1 = inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, inst.src_regs(1), cdna)
dst = _fmt_sdst(inst.sdst, inst.dst_regs(), cdna)
if 'fmamk' in name and lit is not None: return f"{name} {dst}, {s0}, 0x{lit:x}, {s1}"
if 'fmaak' in name and lit is not None: return f"{name} {dst}, {s0}, {s1}, 0x{lit:x}"
return f"{name} {dst}, {s0}, {s1}"
def _disasm_sopc(inst: SOPC) -> str:
cdna = _is_cdna(inst)
s0 = inst.lit(inst.ssrc0) if inst.ssrc0 == 255 else _fmt_src(inst.ssrc0, inst.src_regs(0), cdna)
s1 = inst.lit(inst.ssrc1) if inst.ssrc1 == 255 else _fmt_src(inst.ssrc1, inst.src_regs(1), cdna)
return f"{inst.op_name.lower()} {s0}, {s1}"
def _disasm_sopk(inst: SOPK) -> str:
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
hw = HWREG
def fmt_hwreg(hid, hoff, hsz):
if hid not in hw: return f"0x{inst.simm16:x}"
hr_name = str(hid) if is_rdna4 else hw[hid]
return f"hwreg({hr_name})" if hoff == 0 and hsz == 32 else f"hwreg({hr_name}, {hoff}, {hsz})"
if name == 's_setreg_imm32_b32' or (not cdna and op == SOPKOp.S_SETREG_IMM32_B32):
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
return f"{name} {fmt_hwreg(hid, hoff, hsz)}, 0x{inst._literal:x}"
if not cdna and op == SOPKOp.S_VERSION: return f"{name} 0x{inst.simm16:x}"
if (not cdna and op in (SOPKOp.S_SETREG_B32, SOPKOp.S_GETREG_B32)) or (cdna and name in ('s_setreg_b32', 's_getreg_b32')):
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
hs = fmt_hwreg(hid, hoff, hsz)
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1, cdna)}" if 'setreg' in name else f"{name} {_fmt_sdst(inst.sdst, 1, cdna)}, {hs}"
if not cdna and op in (SOPKOp.S_SUBVECTOR_LOOP_BEGIN, SOPKOp.S_SUBVECTOR_LOOP_END):
return f"{name} {_fmt_sdst(inst.sdst, 1)}, 0x{inst.simm16:x}"
return f"{name} {_fmt_sdst(inst.sdst, inst.dst_regs(), cdna)}, 0x{inst.simm16:x}"
def _disasm_vinterp(inst: VINTERP) -> str:
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
return f"{inst.op_name.lower()} v{inst.vdst}, {inst.lit(inst.src0, inst.neg & 1)}, {inst.lit(inst.src1, inst.neg & 2)}, {inst.lit(inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
EXP_TARGETS = {0: 'mrt0', 1: 'mrt1', 2: 'mrt2', 3: 'mrt3', 4: 'mrt4', 5: 'mrt5', 6: 'mrt6', 7: 'mrt7',
8: 'mrtz', 9: 'null', 12: 'pos0', 13: 'pos1', 14: 'pos2', 15: 'pos3', 16: 'pos4',
32: 'param0', 33: 'param1', 34: 'param2', 35: 'param3', 36: 'param4', 37: 'param5'}
def _disasm_vexport(inst) -> str:
tgt = EXP_TARGETS.get(inst.target, f'{inst.target}')
srcs = [f'v{getattr(inst, f"vsrc{i}")}' if inst.en & (1 << i) else 'off' for i in range(4)]
mods = _mods((inst.done, "done"), (inst.row, "row_en"))
return f"export {tgt} {', '.join(srcs)}" + (" " + mods if mods else "")
def _disasm_vbuffer(inst) -> str:
name = inst.op_name.lower().replace('buffer_', 'buffer_').replace('tbuffer_', 'tbuffer_')
w = (2 if _has(name, 'xyz', 'xyzw') else 1) if 'd16' in name else \
((2 if _has(name, 'b64', 'u64', 'i64') else 1) * (2 if 'cmpswap' in name else 1)) if 'atomic' in name else \
{'b32':1,'b64':2,'b96':3,'b128':4,'b16':1,'x':1,'xy':2,'xyz':3,'xyzw':4}.get(name.split('_')[-1], inst.dst_regs())
if getattr(inst, 'tfe', 0): w += 1
vdata = _vreg(inst.vdata, w) if w else f'v{inst.vdata}'
vaddr = _vreg(inst.vaddr, 2) if inst.offen and inst.idxen else (f'v{inst.vaddr}' if inst.offen or inst.idxen else 'off')
srsrc = f'ttmp[{inst.rsrc - 108}:{inst.rsrc - 108 + 3}]' if inst.rsrc >= 108 else f's[{inst.rsrc}:{inst.rsrc + 3}]'
soff = decode_src(inst.soffset) if inst.soffset >= 106 else f's{inst.soffset}'
fmt = getattr(inst, 'format', 0)
fmt_names = {e.value: e.name for e in BufFmt}
fmt_s = f" format:[{fmt_names[fmt]}]" if fmt > 1 and fmt in fmt_names else (f" format:{fmt}" if fmt > 1 else "")
if 'atomic' in name: th_names = {1: 'TH_ATOMIC_RETURN', 6: 'TH_ATOMIC_CASCADE_NT'}
elif 'store' in name: th_names = {3: 'TH_STORE_BYPASS', 6: 'TH_STORE_NT_HT'}
else: th_names = {3: 'TH_LOAD_BYPASS', 6: 'TH_LOAD_NT_HT'}
scope_names = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
mods = _mods((inst.idxen, "idxen"), (inst.offen, "offen"), (inst.ioffset, f"offset:{inst.ioffset}"),
(inst.th in th_names, f"th:{th_names.get(inst.th, '')}"), (inst.scope in scope_names, f"scope:{scope_names.get(inst.scope, '')}"))
return f"{name} {vdata}, {vaddr}, {srsrc}, {soff}{fmt_s}" + (" " + mods if mods else "")
DISASM_HANDLERS: dict[type, callable] = {
VOP1: _disasm_vop1, VOP2: _disasm_vop2, VOPC: _disasm_vopc, VOP3: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOPD: _disasm_vopd, VOP3P: _disasm_vop3p,
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, MUBUF: _disasm_buf, MTBUF: _disasm_buf,
MIMG: _disasm_mimg, SOP1: _disasm_sop1, SOP2: _disasm_sop2, SOPC: _disasm_sopc, SOPK: _disasm_sopk,
# RDNA4
R4_VOP1: _disasm_vop1, R4_VOP2: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOP3: _disasm_vop3, R4_VOP3SD: _disasm_vop3sd,
R4_VOPD: _disasm_vopd, R4_VOP3P: _disasm_vop3p, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem,
R4_DS: _disasm_ds, R4_SOP1: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOPC: _disasm_sopc, R4_SOPK: _disasm_sopk,
R4_VEXPORT: _disasm_vexport, R4_VBUFFER: _disasm_vbuffer}
def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
# ═══════════════════════════════════════════════════════════════════════════════
# CDNA DISASSEMBLER SUPPORT
# ═══════════════════════════════════════════════════════════════════════════════
try:
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP2 as CDNA_VOP2, VOPC as CDNA_VOPC, VOP3A, VOP3B, VOP3P as CDNA_VOP3P,
SOP1 as CDNA_SOP1, SOP2 as CDNA_SOP2, SOPC as CDNA_SOPC, SOPK as CDNA_SOPK, SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
FLAT as CDNA_FLAT, MUBUF as CDNA_MUBUF, MTBUF as CDNA_MTBUF, SDWA, DPP, VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op, VOPCOp as CDNA_VOPCOp)
def _cdna_src(inst, v, neg, abs_=0, n=1):
s = inst.lit(v) if v == 255 else _fmt_src(v, n, cdna=True)
if abs_: s = f"|{s}|"
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
_CDNA_VOP3_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32'}
def _disasm_vop3a(inst) -> str:
op_val = inst._values.get('op', 0)
if hasattr(op_val, 'value'): op_val = op_val.value
name = inst.op_name.lower() or f'vop3a_op_{op_val}'
from extra.assembly.amd.dsl import spec_num_srcs, spec_regs
n = spec_num_srcs(name) if name else inst.num_srcs()
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
orig_name = name
name = _CDNA_VOP3_ALIASES.get(name, name)
if name != orig_name:
s0, s1 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, 1), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, 1)
s2 = ""
dst = f"v{inst.vdst}"
else:
dregs, r0, r1, r2 = spec_regs(name) if name else (inst.dst_regs(), inst.src_regs(0), inst.src_regs(1), inst.src_regs(2))
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1), _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
dst = _vreg(inst.vdst, dregs) if dregs > 1 else f"v{inst.vdst}"
if op_val >= 512:
return f"{name} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{cl}{om}"
if op_val < 256:
sdst = _fmt_sdst(inst.vdst, 2, cdna=True)
return f"{name}_e64 {sdst}, {s0}, {s1}{cl}"
if 320 <= op_val < 512:
if name in ('v_nop', 'v_clrexcp'): return f"{name}_e64"
return f"{name}_e64 {dst}, {s0}{cl}{om}"
if name == 'v_cndmask_b32':
s2 = _fmt_src(inst.src2, 2, cdna=True)
return f"{name}_e64 {dst}, {s0}, {s1}, {s2}{cl}{om}"
if name in ('v_mul_legacy_f32', 'v_mac_f32'):
return f"{name}_e64 {dst}, {s0}, {s1}{cl}{om}"
suf = "_e64" if op_val < 512 else ""
return f"{name}{suf} {dst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {s0}, {s1}{cl}{om}"
def _disasm_vop3b(inst) -> str:
op_val = inst._values.get('op', 0)
if hasattr(op_val, 'value'): op_val = op_val.value
name = inst.op_name.lower() or f'vop3b_op_{op_val}'
from extra.assembly.amd.dsl import spec_num_srcs, spec_regs
n = spec_num_srcs(name) if name else inst.num_srcs()
dregs, r0, r1, r2 = spec_regs(name) if name else (inst.dst_regs(), inst.src_regs(0), inst.src_regs(1), inst.src_regs(2))
s0, s1, s2 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0), _cdna_src(inst, inst.src1, inst.neg&2, n=r1), _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
dst = _vreg(inst.vdst, dregs) if dregs > 1 else f"v{inst.vdst}"
sdst = _fmt_sdst(inst.sdst, 2, cdna=True)
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
s2 = _fmt_src(inst.src2, 2, cdna=True)
return f"{name}_e64 {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}"
suf = "_e64" if 'co_' in name else ""
return f"{name}{suf} {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name}{suf} {dst}, {sdst}, {s0}, {s1}{cl}{om}"
def _disasm_cdna_vop3p(inst) -> str:
name, n, is_mfma = inst.op_name.lower(), inst.num_srcs(), 'mfma' in inst.op_name.lower() or 'smfmac' in inst.op_name.lower()
get_src = lambda v, sc: inst.lit(v) if v == 255 else _fmt_src(v, sc, cdna=True)
if is_mfma: sc = 2 if 'iu4' in name else 4 if 'iu8' in name or 'i4' in name else 8 if 'f16' in name or 'bf16' in name else 4; src0, src1, src2, dst = get_src(inst.src0, sc), get_src(inst.src1, sc), get_src(inst.src2, 16), _vreg(inst.vdst, 16)
else: src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), f"v{inst.vdst}"
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
mods = ([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else []) + ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != (7 if n == 3 else 3) else []) + \
([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else [])
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}" if n == 3 else f"{name} {dst}, {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
_SEL = {0: 'BYTE_0', 1: 'BYTE_1', 2: 'BYTE_2', 3: 'BYTE_3', 4: 'WORD_0', 5: 'WORD_1', 6: 'DWORD'}
_UNUSED = {0: 'UNUSED_PAD', 1: 'UNUSED_SEXT', 2: 'UNUSED_PRESERVE'}
_DPP = {0x130: "wave_shl:1", 0x134: "wave_rol:1", 0x138: "wave_shr:1", 0x13c: "wave_ror:1", 0x140: "row_mirror", 0x141: "row_half_mirror", 0x142: "row_bcast:15", 0x143: "row_bcast:31"}
def _sdwa_src0(v, is_sgpr, sext=0, neg=0, abs_=0):
s = decode_src(v, cdna=True) if is_sgpr else f"v{v}"
if sext: s = f"sext({s})"
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
def _sdwa_vsrc1(v, sext=0, neg=0, abs_=0):
s = f"v{v}"
if sext: s = f"sext({s})"
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
_OMOD_SDWA = {0: "", 1: " mul:2", 2: " mul:4", 3: " div:2"}
def _disasm_sdwa(inst) -> str:
vop2_op = inst.vop2_op
src0 = _sdwa_src0(inst.src0, inst.s0, inst.src0_sext, inst.src0_neg, inst.src0_abs)
clamp = " clamp" if inst.clmp else ""
omod = _OMOD_SDWA.get(inst.omod, "")
if vop2_op == 63:
try: name = CDNA_VOP1Op(inst.vop_op).name.lower()
except ValueError: name = f"vop1_op_{inst.vop_op}"
dst = f"v{inst.vdst}"
mods = [f"dst_sel:{_SEL[inst.dst_sel]}", f"dst_unused:{_UNUSED[inst.dst_u]}", f"src0_sel:{_SEL[inst.src0_sel]}"]
return f"{name}_sdwa {dst}, {src0}{clamp}{omod} " + " ".join(mods)
elif vop2_op == 62:
try: name = CDNA_VOPCOp(inst.vdst).name.lower()
except ValueError: name = f"vopc_op_{inst.vdst}"
src1 = _sdwa_vsrc1(inst.vop_op, inst.src1_sext, inst.src1_neg, inst.src1_abs)
sdst_enc = inst.dst_sel | (inst.dst_u << 3) | (inst.clmp << 5) | (inst.omod << 6)
if sdst_enc == 0: sdst = "vcc"
else:
sdst_val = sdst_enc - 128 if sdst_enc >= 128 else sdst_enc
sdst = _fmt_sdst(sdst_val, 2, cdna=True)
mods = [f"src0_sel:{_SEL[inst.src0_sel]}", f"src1_sel:{_SEL[inst.src1_sel]}"]
return f"{name}_sdwa {sdst}, {src0}, {src1} " + " ".join(mods)
else:
try: name = CDNA_VOP2Op(vop2_op).name.lower()
except ValueError: name = f"vop2_op_{vop2_op}"
name = _CDNA_DISASM_ALIASES.get(name, name)
dst = f"v{inst.vdst}"
src1 = _sdwa_vsrc1(inst.vop_op, inst.src1_sext, inst.src1_neg, inst.src1_abs)
mods = [f"dst_sel:{_SEL[inst.dst_sel]}", f"dst_unused:{_UNUSED[inst.dst_u]}", f"src0_sel:{_SEL[inst.src0_sel]}", f"src1_sel:{_SEL[inst.src1_sel]}"]
if name == 'v_cndmask_b32':
return f"{name}_sdwa {dst}, {src0}, {src1}, vcc{clamp}{omod} " + " ".join(mods)
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
return f"{name}_sdwa {dst}, vcc, {src0}, {src1}, vcc{clamp}{omod} " + " ".join(mods)
if '_co_' in name:
return f"{name}_sdwa {dst}, vcc, {src0}, {src1}{clamp}{omod} " + " ".join(mods)
return f"{name}_sdwa {dst}, {src0}, {src1}{clamp}{omod} " + " ".join(mods)
def _dpp_src(v, neg=0, abs_=0):
s = f"v{v}" if v < 256 else f"v{v - 256}"
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
def _disasm_dpp(inst) -> str:
vop2_op = inst.vop2_op
ctrl = inst.dpp_ctrl
dpp = f"quad_perm:[{ctrl&3},{(ctrl>>2)&3},{(ctrl>>4)&3},{(ctrl>>6)&3}]" if ctrl < 0x100 else f"row_shl:{ctrl&0xf}" if ctrl < 0x110 else f"row_shr:{ctrl&0xf}" if ctrl < 0x120 else f"row_ror:{ctrl&0xf}" if ctrl < 0x130 else _DPP.get(ctrl, f"dpp_ctrl:0x{ctrl:x}")
src0 = _dpp_src(inst.src0, inst.src0_neg, inst.src0_abs)
mods = [dpp, f"row_mask:0x{inst.row_mask:x}", f"bank_mask:0x{inst.bank_mask:x}"] + (["bound_ctrl:0"] if inst.bound_ctrl else [])
if vop2_op == 63:
try: name = CDNA_VOP1Op(inst.vop_op).name.lower()
except ValueError: name = f"vop1_op_{inst.vop_op}"
return f"{name}_dpp v{inst.vdst}, {src0} " + " ".join(mods)
else:
try: name = CDNA_VOP2Op(vop2_op).name.lower()
except ValueError: name = f"vop2_op_{vop2_op}"
name = _CDNA_DISASM_ALIASES.get(name, name)
src1 = _dpp_src(inst.vop_op, inst.src1_neg, inst.src1_abs)
if name == 'v_cndmask_b32':
return f"{name}_dpp v{inst.vdst}, {src0}, {src1}, vcc " + " ".join(mods)
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
return f"{name}_dpp v{inst.vdst}, vcc, {src0}, {src1}, vcc " + " ".join(mods)
if '_co_' in name:
return f"{name}_dpp v{inst.vdst}, vcc, {src0}, {src1} " + " ".join(mods)
return f"{name}_dpp v{inst.vdst}, {src0}, {src1} " + " ".join(mods)
DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP2: _disasm_vop2, CDNA_VOPC: _disasm_vopc,
CDNA_SOP1: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOPC: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_MUBUF: _disasm_buf, CDNA_MTBUF: _disasm_buf,
VOP3A: _disasm_vop3a, VOP3B: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p, SDWA: _disasm_sdwa, DPP: _disasm_dpp})
except ImportError:
pass
+373 -136
View File
@@ -1,8 +1,122 @@
# library for RDNA3 assembly DSL
# mypy: ignore-errors
from __future__ import annotations
import re, struct
from enum import IntEnum
from functools import cache
from typing import overload, Annotated, TypeVar, Generic
from extra.assembly.amd.autogen.rdna3.enum import (VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, VOPDOp, SOP1Op, SOP2Op,
SOPCOp, SOPKOp, SOPPOp, SMEMOp, DSOp, FLATOp, MUBUFOp, MTBUFOp, MIMGOp, VINTERPOp)
from extra.assembly.amd.autogen.cdna.enum import VOP1Op as CDNA_VOP1Op, VOP2Op as CDNA_VOP2Op
from extra.assembly.amd.autogen.rdna4.enum import (VOP1Op as RDNA4_VOP1Op, VOP2Op as RDNA4_VOP2Op, VOP3Op as RDNA4_VOP3Op,
VOP3SDOp as RDNA4_VOP3SDOp, VOP3POp as RDNA4_VOP3POp, VOPCOp as RDNA4_VOPCOp, VOPDOp as RDNA4_VOPDOp,
SOP1Op as RDNA4_SOP1Op, SOP2Op as RDNA4_SOP2Op, SOPCOp as RDNA4_SOPCOp, SOPKOp as RDNA4_SOPKOp, SOPPOp as RDNA4_SOPPOp,
SMEMOp as RDNA4_SMEMOp, DSOp as RDNA4_DSOp, VBUFFEROp as RDNA4_VBUFFEROp, VINTERPOp as RDNA4_VINTERPOp)
# Source operand encoding - constant across all AMD ISAs
class SrcEnum(IntEnum):
VCC_LO=106; VCC_HI=107; NULL=124; M0=125; EXEC_LO=126; EXEC_HI=127; ZERO=128
DPP8=233; DPP8FI=234; SHARED_BASE=235; SHARED_LIMIT=236; PRIVATE_BASE=237; PRIVATE_LIMIT=238
POS_HALF=240; NEG_HALF=241; POS_ONE=242; NEG_ONE=243; POS_TWO=244; NEG_TWO=245
POS_FOUR=246; NEG_FOUR=247; INV_2PI=248; DPP16=250; VCCZ=251; EXECZ=252; SCC=253; LDS_DIRECT=254
VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, ZERO = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.ZERO
DPP8FI, SHARED_BASE, SHARED_LIMIT, PRIVATE_BASE, PRIVATE_LIMIT = SrcEnum.DPP8FI, SrcEnum.SHARED_BASE, SrcEnum.SHARED_LIMIT, SrcEnum.PRIVATE_BASE, SrcEnum.PRIVATE_LIMIT
POS_HALF, NEG_HALF, POS_ONE, NEG_ONE, POS_TWO, NEG_TWO = SrcEnum.POS_HALF, SrcEnum.NEG_HALF, SrcEnum.POS_ONE, SrcEnum.NEG_ONE, SrcEnum.POS_TWO, SrcEnum.NEG_TWO
POS_FOUR, NEG_FOUR, INV_2PI, VCCZ, EXECZ, SCC, LDS_DIRECT = SrcEnum.POS_FOUR, SrcEnum.NEG_FOUR, SrcEnum.INV_2PI, SrcEnum.VCCZ, SrcEnum.EXECZ, SrcEnum.SCC, SrcEnum.LDS_DIRECT
OFF = NULL
# Common masks
MASK32, MASK64, MASK128 = 0xffffffff, 0xffffffffffffffff, (1 << 128) - 1
# Float/int bit conversion (simple versions for literal encoding)
def _i32(f: float) -> int: return struct.unpack("<I", struct.pack("<f", f))[0]
def _i64(f: float) -> int: return struct.unpack("<Q", struct.pack("<d", f))[0]
# Instruction spec - register counts and dtypes derived from instruction names
_REGS = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16,
'F32': 1, 'I32': 1, 'U32': 1, 'F64': 2, 'I64': 2, 'U64': 2,
'F16': 1, 'I16': 1, 'U16': 1, 'B16': 1, 'I8': 1, 'U8': 1, 'B8': 1,
'DWORD': 1, 'DWORDX2': 2, 'DWORDX3': 3, 'DWORDX4': 4, 'DWORDX8': 8, 'DWORDX16': 16,
'BYTE': 1, 'SHORT': 1, 'UBYTE': 1, 'SBYTE': 1, 'USHORT': 1, 'SSHORT': 1}
_CVT_RE = re.compile(r'CVT_([FIUB]\d+)_([FIUB]\d+)$')
_MAD_MUL_RE = re.compile(r'(?:MAD|MUL)_([IU]\d+)_([IU]\d+)$')
_PACK_RE = re.compile(r'PACK_([FIUB]\d+)_([FIUB]\d+)$')
_DST_SRC_RE = re.compile(r'_([FIUB]\d+)_([FIUB]\d+)$')
_SINGLE_RE = re.compile(r'_([FIUB](?:32|64|16|8|96|128|256|512)|DWORD(?:X(?:2|3|4|8|16))?|[US]?BYTE|[US]?SHORT)$')
@cache
def _suffix(name: str) -> tuple[str | None, str | None]:
name = name.upper()
if m := _CVT_RE.search(name): return m.group(1), m.group(2)
if m := _MAD_MUL_RE.search(name): return m.group(1), m.group(2)
if m := _PACK_RE.search(name): return m.group(1), m.group(2)
if m := _DST_SRC_RE.search(name): return m.group(1), m.group(2)
if m := _SINGLE_RE.search(name): return m.group(1), m.group(1)
return None, None
_SPECIAL_REGS = {
'V_LSHLREV_B64': (2, 1, 2, 1), 'V_LSHRREV_B64': (2, 1, 2, 1), 'V_ASHRREV_I64': (2, 1, 2, 1),
'S_LSHL_B64': (2, 2, 1, 1), 'S_LSHR_B64': (2, 2, 1, 1), 'S_ASHR_I64': (2, 2, 1, 1),
'S_BFE_U64': (2, 2, 1, 1), 'S_BFE_I64': (2, 2, 1, 1), 'S_BFM_B64': (2, 1, 1, 1),
'S_BITSET0_B64': (2, 1, 1, 1), 'S_BITSET1_B64': (2, 1, 1, 1),
'S_BITCMP0_B64': (1, 2, 1, 1), 'S_BITCMP1_B64': (1, 2, 1, 1),
'V_LDEXP_F64': (2, 2, 1, 1), 'V_TRIG_PREOP_F64': (2, 2, 1, 1),
'V_CMP_CLASS_F64': (1, 2, 1, 1), 'V_CMPX_CLASS_F64': (1, 2, 1, 1),
'V_CMP_CLASS_F32': (1, 1, 1, 1), 'V_CMPX_CLASS_F32': (1, 1, 1, 1),
'V_CMP_CLASS_F16': (1, 1, 1, 1), 'V_CMPX_CLASS_F16': (1, 1, 1, 1),
'V_MAD_U64_U32': (2, 1, 1, 2), 'V_MAD_I64_I32': (2, 1, 1, 2),
'V_QSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_PK_U16_U8': (2, 2, 1, 2), 'V_MQSAD_U32_U8': (4, 2, 1, 4),
'V_CVT_PK_F32_BF8': (2, 1, 1, 1), 'V_CVT_PK_F32_FP8': (2, 1, 1, 1),
}
_SPECIAL_DTYPE = {
'V_LSHLREV_B64': ('B64', 'U32', 'B64', None), 'V_LSHRREV_B64': ('B64', 'U32', 'B64', None), 'V_ASHRREV_I64': ('I64', 'U32', 'I64', None),
'S_LSHL_B64': ('B64', 'B64', 'U32', None), 'S_LSHR_B64': ('B64', 'B64', 'U32', None), 'S_ASHR_I64': ('I64', 'I64', 'U32', None),
'S_BFE_U64': ('U64', 'U64', 'U32', None), 'S_BFE_I64': ('I64', 'I64', 'U32', None),
'S_BFM_B64': ('B64', 'U32', 'U32', None), 'S_BITSET0_B64': ('B64', 'U32', None, None), 'S_BITSET1_B64': ('B64', 'U32', None, None),
'S_BITCMP0_B64': ('SCC', 'B64', 'U32', None), 'S_BITCMP1_B64': ('SCC', 'B64', 'U32', None),
'V_LDEXP_F64': ('F64', 'F64', 'I32', None), 'V_TRIG_PREOP_F64': ('F64', 'F64', 'U32', None),
'V_CMP_CLASS_F64': ('VCC', 'F64', 'U32', None), 'V_CMPX_CLASS_F64': ('EXEC', 'F64', 'U32', None),
'V_CMP_CLASS_F32': ('VCC', 'F32', 'U32', None), 'V_CMPX_CLASS_F32': ('EXEC', 'F32', 'U32', None),
'V_CMP_CLASS_F16': ('VCC', 'F16', 'U32', None), 'V_CMPX_CLASS_F16': ('EXEC', 'F16', 'U32', None),
'V_MAD_U64_U32': ('U64', 'U32', 'U32', 'U64'), 'V_MAD_I64_I32': ('I64', 'I32', 'I32', 'I64'),
'V_QSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'), 'V_MQSAD_PK_U16_U8': ('B64', 'B64', 'B64', 'B64'),
'V_MQSAD_U32_U8': ('B128', 'B64', 'B64', 'B128'),
}
@cache
def spec_regs(name: str) -> tuple[int, int, int, int]:
uname = name.upper()
if uname in _SPECIAL_REGS: return _SPECIAL_REGS[uname]
if 'SAD' in uname and 'U8' in uname and 'QSAD' not in uname and 'MQSAD' not in uname: return 1, 1, 1, 1
dst_suf, src_suf = _suffix(name)
return _REGS.get(dst_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1), _REGS.get(src_suf, 1)
@cache
def spec_dtype(name: str) -> tuple[str | None, str | None, str | None, str | None]:
uname = name.upper()
if uname in _SPECIAL_DTYPE: return _SPECIAL_DTYPE[uname]
if 'SAD' in uname and ('U8' in uname or 'U16' in uname) and 'QSAD' not in uname and 'MQSAD' not in uname: return 'U32', 'U32', 'U32', 'U32'
if '_CMP_' in uname or '_CMPX_' in uname:
dst_suf, src_suf = _suffix(name)
return 'EXEC' if '_CMPX_' in uname else 'VCC', src_suf, src_suf, None
dst_suf, src_suf = _suffix(name)
return dst_suf, src_suf, src_suf, src_suf
_F16_RE = re.compile(r'_[FIUB]16(?:_|$)')
_F64_RE = re.compile(r'_[FIUB]64(?:_|$)')
@cache
def spec_is_16bit(name: str) -> bool:
uname = name.upper()
if 'SAD' in uname or 'PACK' in uname or '_PK_' in uname or 'SAT_PK' in uname or 'DOT2' in uname: return False
if '_F32' in uname or '_I32' in uname or '_U32' in uname or '_B32' in uname: return False
return bool(_F16_RE.search(uname))
@cache
def spec_is_64bit(name: str) -> bool: return bool(_F64_RE.search(name.upper()))
_3SRC = {'FMA', 'MAD', 'MIN3', 'MAX3', 'MED3', 'DIV_FIX', 'DIV_FMAS', 'DIV_SCALE', 'SAD', 'LERP', 'ALIGN', 'CUBE', 'BFE', 'BFI',
'PERM_B32', 'PERMLANE', 'CNDMASK', 'XOR3', 'OR3', 'ADD3', 'LSHL_OR', 'AND_OR', 'LSHL_ADD', 'ADD_LSHL', 'XAD', 'MAXMIN',
'MINMAX', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM', 'MINIMUM3', 'MAXIMUM3', 'DOT2', 'DOT4', 'DOT8', 'WMMA', 'CVT_PK_U8', 'MULLIT', 'CO_CI'}
_2SRC = {'FMAC', 'PERMLANE16_VAR', 'PERMLANEX16_VAR'} # FMAC uses dst as implicit accumulator, _VAR permlane only 2 sources
def spec_num_srcs(name: str) -> int:
name = name.upper()
if any(k in name for k in _2SRC): return 2
return 3 if any(k in name for k in _3SRC) else 2
def is_dtype_16(dt: str | None) -> bool: return dt is not None and '16' in dt
def is_dtype_64(dt: str | None) -> bool: return dt is not None and '64' in dt
# Bit field DSL
class BitField:
@@ -30,6 +144,10 @@ class BitField:
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
@@ -112,31 +230,46 @@ def unwrap(val) -> int:
if hasattr(val, 'idx'): return val.idx # Reg
return val
# Encoding helpers
# Encoding/decoding constants
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
FLOAT_DEC = {v: str(k) for k, v in FLOAT_ENC.items()}
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
SPECIAL_GPRS_CDNA = {102: "flat_scratch_lo", 103: "flat_scratch_hi", 104: "xnack_mask_lo", 105: "xnack_mask_hi",
106: "vcc_lo", 107: "vcc_hi", 124: "m0", 126: "exec_lo", 127: "exec_hi",
251: "src_vccz", 252: "src_execz", 253: "src_scc", 254: "src_lds_direct"}
SPECIAL_PAIRS = {106: "vcc", 126: "exec"}
SPECIAL_PAIRS_CDNA = {102: "flat_scratch", 104: "xnack_mask", 106: "vcc", 126: "exec"}
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata', 'vsrc1'}
def _encode_reg(val: Reg) -> int:
if isinstance(val, TTMP): return 108 + val.idx
return val.idx # hi bit is handled via opsel, not in register encoding
def _encode_reg(val: Reg) -> int: return (108 if isinstance(val, TTMP) else 0) + val.idx
def _is_encoded_src(v: int) -> bool: return 106 <= v <= 127 or 128 <= v <= 208 or 240 <= v <= 255 # Special regs (106-127) or inline const
def encode_src(val) -> int:
if isinstance(val, VGPR): return 256 + _encode_reg(val)
if isinstance(val, Reg): return _encode_reg(val)
if isinstance(val, SrcMod) and not isinstance(val, Reg):
# SrcMod wraps either special registers (VCC_LO=106, EXEC_LO=126, etc.) or literals
# Special register values are in valid encoding ranges - return as-is
# Literals (large integers) need 255 marker
v = val.val
# Valid source encoding ranges: 0-127 (SGPRs/special), 128-192 (inline const), 193-208 (neg inline), 240-247 (float), 251-253 (special)
if 0 <= v <= 127 or 240 <= v <= 255: return v # SGPRs, special regs, float constants
if 128 <= v <= 192: return v # Inline positive constants (0-64)
if 193 <= v <= 208: return v # Inline negative constants (-1 to -16)
return 255 # Literal marker - value stored separately
if _is_encoded_src(v): return v # Already encoded (special reg 106-127 or inline const 128-208 or float 240-255)
if isinstance(v, int) and 0 <= v <= 64: return 128 + v # Encode as inline constant
if isinstance(v, int) and -16 <= v <= -1: return 192 - v
return 255 # Literal
if hasattr(val, 'value'): return val.value # IntEnum
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
if isinstance(val, int): return 128 + val if 0 <= val <= 64 else 192 - val if -16 <= val <= -1 else 255
return 255
def decode_src(val: int, cdna: bool = False) -> str:
special = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
if val in special: return special[val]
if val <= 105: return f"s{val}"
if val in FLOAT_DEC: return FLOAT_DEC[val]
if 108 <= val <= 123: return f"ttmp{val - 108}"
if 128 <= val <= 192: return str(val - 128)
if 193 <= val <= 208: return str(-(val - 192))
if 256 <= val <= 511: return f"v{val - 256}"
return "lit" if val == 255 else f"?{val}"
# Instruction base class
class Inst:
@@ -149,127 +282,152 @@ class Inst:
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))}
# Merge fields from parent classes
cls._fields = {}
for base in reversed(cls.__mro__):
if base is Inst or not hasattr(base, '_fields'): continue
cls._fields.update(base._fields)
# Add this class's own fields (overrides parents)
cls._fields.update({n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))})
# Compute size from max bit (exclude optional MIMG NSA fields: addr1/addr2 at bits 64+)
optional_nsa = {'addr1', 'addr2'}
max_bit = max((bf.hi for n, bf in cls._fields.items() if n not in optional_nsa), default=0) if cls._fields else 0
cls._sz = 12 if max_bit > 63 else 8 if max_bit > 31 else 4
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
def _or_field(self, name: str, bit: int):
cur = self._values.get(name, 0)
self._values[name] = (cur.val if isinstance(cur, RawImm) else cur) | bit
def _encode_src(self, name: str, val):
"""Encode a source field, handling modifiers and literals."""
encoded = encode_src(val)
has_opsel = 'opsel' in self._fields
if isinstance(val, Reg) and val.hi and not has_opsel: encoded |= 0x80 # hi bit in src for VOP1/2/C
self._values[name] = RawImm(encoded)
# Handle neg/abs/opsel modifiers
if isinstance(val, SrcMod):
mod_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
if val.neg and 'neg' in self._fields: self._or_field('neg', mod_bit)
if val.abs_ and 'abs' in self._fields: self._or_field('abs', mod_bit)
if isinstance(val, Reg) and val.hi and has_opsel:
self._or_field('opsel', {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0))
# Track literal value if needed
if encoded == 255 and self._literal is None:
import struct
# Check if THIS source uses 64-bit encoding (not just src0)
src_idx = {'src0': 0, 'src1': 1, 'src2': 2, 'ssrc0': 0, 'ssrc1': 1}.get(name, 0)
src_regs = self.src_regs(src_idx)
is_64 = src_regs == 2
if isinstance(val, SrcMod) and not isinstance(val, Reg): lit32 = val.val & MASK32
elif isinstance(val, int) and not isinstance(val, IntEnum): lit32 = val & MASK32
elif isinstance(val, float): lit32 = (_i64(val) >> 32) if is_64 else _i32(val) # f64: high 32 bits of f64 repr
else: return
self._literal = (lit32 << 32) if is_64 else lit32
def _encode_raw(self, name: str, val):
"""Encode a raw register field (vdst, vdata, etc.)."""
if isinstance(val, Reg):
encoded = _encode_reg(val)
if val.hi and 'opsel' not in self._fields: encoded |= 0x80
self._values[name] = encoded
if name == 'vdst' and val.hi and 'opsel' in self._fields: self._or_field('opsel', 8)
elif hasattr(val, 'value'): self._values[name] = val.value
def _validate(self, orig_args: dict):
"""Format-specific validation. Override in subclass or check by class name."""
cls_name, op = self.__class__.__name__, orig_args.get('op')
op_val = op.value if hasattr(op, 'value') else op
op_name = op.name if hasattr(op, 'name') else None
# SMEM: register count must match opcode (derive from name: b32=1, b64=2, b96=3, b128=4, b256=8, b512=16, i8/u8/i16/u16=1)
if cls_name == 'SMEM' and op_name:
expected = {'B32': 1, 'B64': 2, 'B96': 3, 'B128': 4, 'B256': 8, 'B512': 16, 'I8': 1, 'U8': 1, 'I16': 1, 'U16': 1}.get(op_name.split('_')[-1])
sdata = orig_args.get('sdata')
if expected and isinstance(sdata, Reg) and sdata.count != expected:
raise ValueError(f"SMEM op {op_name} expects {expected} registers, got {sdata.count}")
# SOP1: derive expected register sizes from op name (e.g., S_MOV_B64 -> dst=2, src=2; S_CTZ_I32_B64 -> dst=1, src=2)
if cls_name == 'SOP1' and hasattr(orig_args.get('op'), 'name'):
op_name = orig_args['op'].name
# Special cases: BITSET takes bit index (1 reg) regardless of dst size
if 'BITSET' in op_name:
dst_size = 2 if op_name.endswith('_B64') else 1
src_size = 1 # bit index is always 1 reg
else:
# Extract sizes from name: last suffix is src type, second-to-last (if exists) is dst type
sizes = {'B32': 1, 'I32': 1, 'U32': 1, 'B64': 2, 'I64': 2, 'U64': 2, 'B128': 4, 'B256': 8, 'B512': 16}
parts = op_name.split('_')
src_size = sizes.get(parts[-1], 1) if parts[-1] in sizes else 1
dst_size = sizes.get(parts[-2], src_size) if len(parts) >= 2 and parts[-2] in sizes else src_size
for fld, expected in [('sdst', dst_size), ('ssrc0', src_size)]:
if isinstance(orig_args.get(fld), Reg) and orig_args[fld].count != expected:
raise ValueError(f"SOP1 {op_name} expects {expected} register(s) for {fld}, got {orig_args[fld].count}")
def __init__(self, *args, literal: int | None = None, **kwargs):
self._values, self._literal = dict(self._defaults), literal
# Map positional args to field names
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))
orig_args.update(kwargs)
# Map Python-friendly names to actual field names (abs_ -> abs for Python reserved word)
if 'abs_' in kwargs: kwargs['abs'] = kwargs.pop('abs_')
# If more args than fields, treat extra arg as literal (for FMAAK/FMAMK style instructions)
# FMAMK has K in middle (vdst, src0, K, vsrc1), FMAAK has K at end (vdst, src0, vsrc1, K)
args = list(args)
if len(args) > len(field_names) and literal is None:
for i, a in enumerate(args):
if isinstance(a, int) and not isinstance(a, SrcEnum) and i < len(field_names) and field_names[i] in ('vsrc1',):
literal = args.pop(i)
break
else:
literal = args.pop() # fallback: last arg is literal
orig_args = dict(zip(field_names, args)) | kwargs
self._values.update(orig_args)
# 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}")
# FLAT: set sve=1 when addr is a VGPR for scratch only
# For scratch (seg=1), sve=1 means addr VGPR is used; sve=0 means addr is "off"
# For global (seg=2) and flat (seg=0), sve is always 0
if self.__class__.__name__ == 'FLAT' and 'sve' in self._fields:
seg_val = self._values.get('seg', 0)
if isinstance(seg_val, RawImm): seg_val = seg_val.val
addr_val = orig_args.get('addr')
if seg_val == 1 and isinstance(addr_val, VGPR): self._values['sve'] = 1
# VOP3P: v_fma_mix* instructions (opcodes 32-34) have opsel_hi default of 0, not 7
if self.__class__.__name__ == 'VOP3P':
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
if hasattr(op_val, 'value'): op_val = op_val.value
if op_val in (32, 33, 34) and 'opsel_hi' not in orig_args and 'opsel_hi2' not in orig_args:
self._values['opsel_hi'] = 0
self._values['opsel_hi2'] = 0
# Type check and encode values
self._precompute()
self._validate(orig_args)
# Pre-shift literal for 64-bit sources (literal param is always raw 32-bit value from user)
if literal is not None:
# Find which source uses the literal (255) and check its register count
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
v = orig_args.get(n)
if (isinstance(v, RawImm) and v.val == 255) or (isinstance(v, int) and v == 255):
self._literal = (literal << 32) if self.src_regs(idx) == 2 else literal
break
else:
self._literal = literal # fallback if no literal source found
cls_name = self.__class__.__name__
# Format-specific setup
if cls_name == 'FLAT' and 'sve' in self._fields:
seg = self._values.get('seg', 0)
if (seg.val if isinstance(seg, RawImm) else seg) == 1 and isinstance(orig_args.get('addr'), VGPR): self._values['sve'] = 1
if cls_name == 'VOP3P':
op = orig_args.get('op')
if hasattr(op, 'value'): op = op.value
# fma_mix ops (32-34) default to opsel_hi=0, WMMA ops (64-69) default to opsel_hi=7 to match LLVM
if op in (32, 33, 34) and 'opsel_hi' not in orig_args: self._values['opsel_hi'] = self._values['opsel_hi2'] = 0
if op in range(64, 70) and 'opsel_hi' not in orig_args: self._values['opsel_hi'], self._values['opsel_hi2'] = 3, 1
# Encode all fields
for name, val in list(self._values.items()):
if name == 'encoding': continue
# 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, SrcMod, 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 _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 source fields as RawImm for consistent disassembly
if name in SRC_FIELDS:
encoded = encode_src(val)
# For VOP1/VOP2/VOPC (no opsel field), encode hi bit in src value
if isinstance(val, Reg) and val.hi and 'opsel' not in self._fields:
encoded |= 0x80
self._values[name] = RawImm(encoded)
# Handle neg/abs/opsel modifiers for VOP3 instructions
if isinstance(val, SrcMod):
if val.neg and 'neg' in self._fields:
neg_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
cur_neg = self._values.get('neg', 0)
self._values['neg'] = (cur_neg.val if isinstance(cur_neg, RawImm) else cur_neg) | neg_bit
if val.abs_ and 'abs' in self._fields:
abs_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
cur_abs = self._values.get('abs', 0)
self._values['abs'] = (cur_abs.val if isinstance(cur_abs, RawImm) else cur_abs) | abs_bit
# Handle hi (opsel) for 16-bit ops - only for formats with opsel field
if isinstance(val, Reg) and val.hi and 'opsel' in self._fields:
opsel_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
cur_opsel = self._values.get('opsel', 0)
self._values['opsel'] = (cur_opsel.val if isinstance(cur_opsel, RawImm) else cur_opsel) | opsel_bit
# Track literal value if needed (encoded as 255)
# For 64-bit ops, store literal in high 32 bits (to match from_bytes decoding and to_bytes encoding)
if encoded == 255 and self._literal is None:
if isinstance(val, SrcMod) and not isinstance(val, Reg):
# SrcMod wrapping a literal value
self._literal = (val.val << 32) if self._is_64bit_op() else val.val
elif isinstance(val, int) and not isinstance(val, IntEnum):
self._literal = (val << 32) if self._is_64bit_op() else val
elif isinstance(val, float):
import struct
lit32 = struct.unpack('<I', struct.pack('<f', val))[0]
self._literal = (lit32 << 32) if self._is_64bit_op() else lit32
# Encode raw register fields for consistent repr
elif name in RAW_FIELDS:
if isinstance(val, Reg):
encoded = _encode_reg(val)
# For VOP1/VOP2/VOPC (no opsel field), encode hi bit in register value
if val.hi and 'opsel' not in self._fields:
encoded |= 0x80
self._values[name] = encoded
# Handle vdst hi (opsel bit 3) for 16-bit ops - only for formats with opsel field
if name == 'vdst' and val.hi and 'opsel' in self._fields:
cur_opsel = self._values.get('opsel', 0)
self._values['opsel'] = (cur_opsel.val if isinstance(cur_opsel, RawImm) else cur_opsel) | 8
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':
if isinstance(val, Reg): self._values[name] = val.idx // 2
elif isinstance(val, SrcMod): self._values[name] = val.val // 2 # Special regs like VCC_LO
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
# Encode by field type
if name in SRC_FIELDS: self._encode_src(name, val)
elif name in RAW_FIELDS: self._encode_raw(name, val)
elif name == 'sbase': self._values[name] = (val.idx if isinstance(val, Reg) else val.val if isinstance(val, SrcMod) else val * 2) // 2
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg): self._values[name] = _encode_reg(val) // 4
elif marker is _VDSTYEnc and isinstance(val, VGPR): self._values[name] = val.idx >> 1
self._precompute_fields()
def _encode_field(self, name: str, val) -> int:
if isinstance(val, RawImm): return val.val
if isinstance(val, SrcMod) and not isinstance(val, Reg): return val.val # Special regs like VCC_LO
if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
if name in {'srsrc', 'ssamp'}: return _encode_reg(val) // 4 if isinstance(val, Reg) else val
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val.val // 2 if isinstance(val, SrcMod) else val
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
@@ -287,35 +445,39 @@ class Inst:
return None
def _is_64bit_op(self) -> bool:
"""Check if this instruction uses 64-bit operands (and thus 64-bit literals).
Exception: V_LDEXP_F64 has 32-bit integer src1, so its literal is 32-bit."""
"""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 may be an enum (from __init__) or an int (from from_int)
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':
from extra.assembly.amd.autogen.rdna3 import VOP3Op
try: op_name = VOP3Op(op).name
except ValueError: pass
if op_name is None and self.__class__.__name__ == 'VOPC':
from extra.assembly.amd.autogen.rdna3 import VOPCOp
try: op_name = VOPCOp(op).name
except ValueError: pass
if op_name is None: return False
# V_LDEXP_F64 has 32-bit integer exponent in src1, so literal is 32-bit
if op_name == 'V_LDEXP_F64': return False
return op_name.endswith(('_F64', '_B64', '_I64', '_U64'))
# 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 ops, literal is stored in high 32 bits internally, but encoded as 4 bytes
lit32 = (lit >> 32) if self._is_64bit_op() else lit
return result + (lit32 & 0xffffffff).to_bytes(4, 'little')
# 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(cls) -> int: return cls._sz
def size(self) -> int:
# Literal is always 4 bytes in the binary (for 64-bit ops, it's in high 32 bits)
return self._size() + (4 if self._literal is not None else 0)
@@ -325,14 +487,22 @@ class Inst:
inst = object.__new__(cls)
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
inst._literal = None
inst._precompute()
inst._precompute_fields()
return inst
@classmethod
def from_bytes(cls, data: bytes):
import typing
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
op_val = inst._values.get('op', 0)
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))
# Check for instructions that always have a literal constant (FMAMK/FMAAK/MADMK/MADAK, SETREG_IMM32)
op_name = ''
if cls.__name__ in ('VOP2', 'SOP2', 'SOPK') and 'op' in (hints := typing.get_type_hints(cls, include_extras=True)):
if typing.get_origin(hints['op']) is typing.Annotated:
try: op_name = typing.get_args(hints['op'])[1](op_val).name
except (ValueError, TypeError): pass
has_literal = any(x in op_name for x in ('FMAMK', 'FMAAK', 'MADMK', 'MADAK', 'SETREG_IMM32'))
# VOPD fmaak/fmamk always have a literal (opx/opy value 1 or 2)
opx, opy = inst._values.get('opx', 0), inst._values.get('opy', 0)
has_literal = has_literal or (cls.__name__ == 'VOPD' and (opx in (1, 2) or opy in (1, 2)))
@@ -341,9 +511,16 @@ class Inst:
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')
inst._literal = (lit32 << 32) if inst._is_64bit_op() else lit32
# Find which source has literal (255) and check its register count
lit_src_is_64 = False
for n, idx in [('src0', 0), ('src1', 1), ('src2', 2), ('ssrc0', 0), ('ssrc1', 1)]:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255:
lit_src_is_64 = inst.src_regs(idx) == 2
break
inst._literal = (lit32 << 32) if lit_src_is_64 else lit32
return inst
def __repr__(self):
@@ -358,9 +535,14 @@ class Inst:
if name.startswith('_'): raise AttributeError(name)
return unwrap(self._values.get(name, 0))
def lit(self, v: int) -> str:
from extra.assembly.amd.asm import decode_src
return f"0x{self._literal:x}" if v == 255 and self._literal else decode_src(v)
def lit(self, v: int, neg: bool = False) -> str:
if v == 255 and self._literal is not None:
# For 64-bit sources, literal is stored shifted - extract the 32-bit value
lit32 = (self._literal >> 32) if self._literal > 0xffffffff else self._literal
s = f"0x{lit32:x}"
else:
s = decode_src(v, 'cdna' in self.__class__.__module__)
return f"-{s}" if neg else s
def __eq__(self, other):
if not isinstance(other, Inst): return NotImplemented
@@ -372,5 +554,60 @@ class Inst:
from extra.assembly.amd.asm import disasm
return disasm(self)
class Inst32(Inst): pass
class Inst64(Inst): pass
_enum_map = {'VOP1': VOP1Op, 'VOP2': VOP2Op, 'VOP3': VOP3Op, 'VOP3SD': VOP3SDOp, 'VOP3P': VOP3POp, 'VOPC': VOPCOp,
'SOP1': SOP1Op, 'SOP2': SOP2Op, 'SOPC': SOPCOp, 'SOPK': SOPKOp, 'SOPP': SOPPOp,
'SMEM': SMEMOp, 'DS': DSOp, 'FLAT': FLATOp, 'MUBUF': MUBUFOp, 'MTBUF': MTBUFOp, 'MIMG': MIMGOp,
'VOPD': VOPDOp, 'VINTERP': VINTERPOp}
_VOP3SD_OPS = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
def _precompute(self):
"""Precompute op, op_name, _spec_regs, _spec_dtype for fast access."""
val = self._values.get('op')
if val is None: self.op = None
elif hasattr(val, 'name'): self.op = val
else:
cls_name = self.__class__.__name__
is_cdna = cls_name in ('VOP3A', 'VOP3B')
# Try marker enum first (VOP3AOp, VOP3BOp, etc.)
marker = self._fields['op'].marker if 'op' in self._fields else None
if marker and issubclass(marker, IntEnum):
try: self.op = marker(val)
except ValueError: self.op = val
elif cls_name in self._enum_map:
try: self.op = self._enum_map[cls_name](val)
except ValueError: self.op = val
else: self.op = val
# Fallback for promoted instructions when marker lookup failed
if not hasattr(self.op, 'name') and cls_name in ('VOP3', 'VOP3A', 'VOP3B') and isinstance(val, int):
if val < 256:
try: self.op = VOPCOp(val)
except ValueError: pass
elif is_cdna and 256 <= val < 512:
try: self.op = (CDNA_VOP1Op(val - 320) if val >= 320 else CDNA_VOP2Op(val - 256))
except ValueError: pass
elif val in self._VOP3SD_OPS and not is_cdna:
try: self.op = VOP3SDOp(val)
except ValueError: pass
elif 256 <= val < 512 and not is_cdna:
try: self.op = VOP1Op(val - 384) if val >= 384 else VOP2Op(val - 256)
except ValueError: pass
self.op_name = self.op.name if hasattr(self.op, 'name') else ''
self._spec_regs = spec_regs(self.op_name)
self._spec_dtype = spec_dtype(self.op_name)
def _precompute_fields(self):
"""Unwrap all field values as direct attributes for fast access."""
for name, val in self._values.items():
if name != 'op': setattr(self, name, unwrap(val))
def dst_regs(self) -> int: return self._spec_regs[0]
def src_regs(self, n: int) -> int: return self._spec_regs[n + 1]
def num_srcs(self) -> int: return spec_num_srcs(self.op_name)
def dst_dtype(self) -> str | None: return self._spec_dtype[0]
def src_dtype(self, n: int) -> str | None: return self._spec_dtype[n + 1]
def is_src_16(self, n: int) -> bool: return self._spec_regs[n + 1] == 1 and is_dtype_16(self._spec_dtype[n + 1])
def is_src_64(self, n: int) -> bool: return self._spec_regs[n + 1] == 2
def is_16bit(self) -> bool: return spec_is_16bit(self.op_name)
def is_64bit(self) -> bool: return spec_is_64bit(self.op_name)
def is_dst_16(self) -> bool: return self._spec_regs[0] == 1 and is_dtype_16(self._spec_dtype[0])
+328 -705
View File
File diff suppressed because it is too large Load Diff
-772
View File
@@ -1,772 +0,0 @@
# PDF parsing and code generation for AMD ISA
# Generates both autogen/__init__.py (instruction formats) and gen_pcode.py (pseudocode functions)
# Usage: python -m extra.assembly.amd.pdf --arch rdna3
import re
PDF_URLS = {
"rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content", # RDNA3.5
"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_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'^([SV]_[A-Z0-9_]+)\s+(\d+)\s*$', re.M)
# ═══════════════════════════════════════════════════════════════════════════════
# SHARED PDF PARSING INFRASTRUCTURE
# ═══════════════════════════════════════════════════════════════════════════════
class _LazyPageCache:
"""Lazy page text/table extractor with caching to avoid redundant PDF parsing."""
__slots__ = ('_pdf', '_offset', '_text_cache', '_table_cache')
def __init__(self, pdf, offset: int = 0):
self._pdf, self._offset, self._text_cache, self._table_cache = pdf, offset, {}, {}
def text(self, idx: int) -> str:
if idx not in self._text_cache: self._text_cache[idx] = self._pdf.pages[self._offset + idx].extract_text() or ''
return self._text_cache[idx]
def tables(self, idx: int) -> list:
if idx not in self._table_cache: self._table_cache[idx] = [t.extract() for t in self._pdf.pages[self._offset + idx].find_tables()]
return self._table_cache[idx]
def texts_range(self, start: int, end: int) -> list[str]: return [self.text(i) for i in range(start, end)]
def _detect_doc_type(first_page_text: str) -> tuple[bool, str]:
"""Detect document type from first page text. Returns (is_cdna, doc_name)."""
is_cdna4 = 'CDNA4' in first_page_text or 'CDNA 4' in first_page_text
is_cdna3 = 'CDNA3' in first_page_text or 'CDNA 3' in first_page_text or 'MI300' in first_page_text
is_cdna = is_cdna3 or is_cdna4
is_rdna4 = 'RDNA4' in first_page_text or 'RDNA 4' in first_page_text
is_rdna35 = 'RDNA3.5' in first_page_text or 'RDNA 3.5' in first_page_text
is_rdna3 = not is_rdna35 and ('RDNA3' in first_page_text or 'RDNA 3' in first_page_text)
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"
return is_cdna, doc_name
def _find_chapter(cache: _LazyPageCache, total_pages: int, pattern: str, sample_pcts: list[float]) -> int | None:
"""Find chapter page by sampling at likely positions then searching nearby. Returns page index or None."""
for pct in sample_pcts:
idx = int(total_pages * pct)
if 0 <= idx < total_pages and re.search(pattern, cache.text(idx)): return idx
for pct in sample_pcts:
base = int(total_pages * pct)
for offset in range(-10, 11):
idx = base + offset
if 0 <= idx < total_pages and re.search(pattern, cache.text(idx)): return idx
return None
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 = 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
# ═══════════════════════════════════════════════════════════════════════════════
# PARSE SINGLE PDF - extracts both format definitions AND pseudocode
# ═══════════════════════════════════════════════════════════════════════════════
def _parse_single_pdf(url: str) -> dict:
"""Parse a single PDF and return raw data for both dsl and pcode generation."""
import pdfplumber
from tinygrad.helpers import fetch
pdf = pdfplumber.open(fetch(url))
total_pages = len(pdf.pages)
cache = _LazyPageCache(pdf)
# Auto-detect document type from first page
is_cdna, doc_name = _detect_doc_type(cache.text(0))
# Find chapter positions using sampling
instr_pattern = r'Chapter \d+\.\s+Instructions\b'
microcode_pattern = r'Chapter \d+\.\s+Microcode Formats'
def is_microcode_page(text):
return (re.search(r'\d+\.\d+\.\d+\.\s+SOP2\s*\n.*Description', text) or
re.search(r'Chapter \d+\.\s+Microcode Formats\s*\n.*This section', text))
# Find microcode section with sampling (RDNA ~23%, CDNA ~93%)
microcode_start = None
for pct in [0.233, 0.93, 0.24, 0.92, 0.25, 0.22]:
sample = int(total_pages * pct)
if is_microcode_page(cache.text(sample)):
microcode_start = sample
while microcode_start > 0 and is_microcode_page(cache.text(microcode_start - 1)):
microcode_start -= 1
break
if microcode_start is None:
for i in range(int(total_pages * 0.15), total_pages):
if is_microcode_page(cache.text(i)):
microcode_start = i
break
if microcode_start is None: microcode_start = int(total_pages * 0.9)
# Find instructions chapter for pseudocode extraction
if is_cdna:
instr_start = _find_chapter(cache, total_pages, instr_pattern, [0.17, 0.18, 0.16, 0.19, 0.15])
instr_end = _find_chapter(cache, total_pages, microcode_pattern, [0.93, 0.92, 0.94, 0.91, 0.95])
else:
instr_start = _find_chapter(cache, total_pages, instr_pattern, [0.30, 0.31, 0.29, 0.32, 0.28])
instr_end = None # RDNA: Instructions goes to end
if instr_start is None: instr_start = int(total_pages * (0.17 if is_cdna else 0.30))
if instr_end is None: instr_end = total_pages
# ─── Parse format definitions from Microcode Formats chapter ───
fmt_cache = _LazyPageCache(pdf, microcode_start)
page_count = min(45, total_pages - microcode_start)
for idx, text in cache._text_cache.items():
if microcode_start <= idx < microcode_start + page_count: fmt_cache._text_cache[idx - microcode_start] = text
# Parse SSRC encoding
src_enum = dict(SRC_EXTRAS)
for i in range(2, 12):
text = fmt_cache.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(fmt_cache.texts_range(2, page_count))
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 format field tables
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 (CDNA has formats from page 1, RDNA from page 2)
format_headers = []
for i in range(1 if is_cdna else 2, page_count):
text = fmt_cache.text(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 < page_count:
next_text = fmt_cache.text(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()))
# Parse instruction formats
formats: dict[str, list[tuple]] = {}
for fmt_name, page_idx, header_pos in format_headers:
if fmt_name in formats: continue
text, tables = fmt_cache.text(page_idx), fmt_cache.tables(page_idx)
field_pos = text.find('Field Name', header_pos)
fields = None
for offset in range(3):
if page_idx + offset >= page_count: break
if offset > 0 and has_header_before_fields(fmt_cache.text(page_idx + offset)): break
for t in fmt_cache.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 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}
# Look for continuation tables on subsequent pages
for pg_offset in range(1, 3):
if page_idx + pg_offset >= page_count or has_header_before_fields(fmt_cache.text(page_idx + pg_offset)): break
for t in fmt_cache.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']]
# ─── Parse pseudocode from Instructions chapter ───
instr_text = '\n'.join(cache.text(i) for i in range(instr_start, instr_end))
return {"formats": formats, "enums": enums, "src_enum": src_enum, "doc_name": doc_name, "is_cdna": is_cdna,
"instr_text": instr_text}
def _merge_results(results: list[dict]) -> dict:
"""Merge results from multiple PDFs (e.g., CDNA3 + CDNA4)."""
merged = {"formats": {}, "enums": {}, "src_enum": {}, "doc_names": [], "instr_texts": []}
for r in results:
merged["doc_names"].append(r["doc_name"])
merged["instr_texts"].append(r["instr_text"])
for val, name in r["src_enum"].items():
if val in merged["src_enum"]:
assert merged["src_enum"][val] == name, f"SrcEnum conflict: {val} = {merged['src_enum'][val]} vs {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, f"{enum_name} conflict: {val} = {merged['enums'][enum_name][val]} vs {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:
name, hi, lo = f[0], f[1], f[2]
if name in existing:
assert existing[name] == (hi, lo), f"Format {fmt_name} field {name} conflict: bits {existing[name]} vs ({hi}, {lo})"
else: merged["formats"][fmt_name].append(f)
return merged
# ═══════════════════════════════════════════════════════════════════════════════
# GENERATE __init__.py (instruction formats and enums)
# ═══════════════════════════════════════════════════════════════════════════════
def _generate_dsl(merged: dict, doc_name: str, output_path: str | None = None) -> str:
"""Generate instruction definitions code."""
formats, enums, src_enum = merged["formats"], merged["enums"], merged["src_enum"]
def enum_lines(name, items):
return [f"class {name}(IntEnum):"] + [f" {n} = {v}" for v, n in sorted(items.items())] + [""]
def field_key(f):
order = FIELD_ORDER.get(fmt_name, [])
return order.index(f[0].lower()) if f[0].lower() in order else 1000
lines = [f"# autogenerated from AMD {doc_name} ISA PDF by generate.py - do not edit", "from enum import IntEnum",
"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",
"import functools", ""]
lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
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"
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):
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=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"):
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})")
skip_exports = {'DPP8', 'DPP16'}
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 skip_exports]
if "NULL" in src_names: lines.append("OFF = NULL\n")
content = '\n'.join(lines)
if output_path is not None:
import pathlib
pathlib.Path(output_path).write_text(content)
return content
# ═══════════════════════════════════════════════════════════════════════════════
# PSEUDOCODE COMPILER: pseudocode -> Python
# ═══════════════════════════════════════════════════════════════════════════════
def _expr(e: str) -> str:
"""Expression transform: minimal - just fix syntax differences."""
e = e.strip()
e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
e = re.sub(r'!([^=])', r' not \1', e)
# Pack: { hi, lo } -> _pack(hi, lo)
e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
def pack(m):
hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
return f'_pack({hi}, {lo})'
e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
# Special constant: 1201'B(2.0 / PI) -> TWO_OVER_PI_1201 (precomputed 1201-bit 2/pi)
e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
# Literals: 1'0U -> 0, 32'I(x) -> (x), B(x) -> (x)
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)
# Constants
e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
# Verilog bit slice: [start +: width] -> [start + width - 1 : start]
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)
# Recursively process bracket contents
def process_brackets(s):
result, i = [], 0
while i < len(s):
if s[i] == '[':
depth, start = 1, i + 1
j = start
while j < len(s) and depth > 0:
if s[j] == '[': depth += 1
elif s[j] == ']': depth -= 1
j += 1
inner = _expr(s[start:j-1])
result.append('[' + inner + ']')
i = j
else:
result.append(s[i])
i += 1
return ''.join(result)
e = process_brackets(e)
# Ternary: a ? b : c -> (b if a else c)
while '?' in e:
depth, bracket, q = 0, 0, -1
for i, c in enumerate(e):
if c == '(': depth += 1
elif c == ')': depth -= 1
elif c == '[': bracket += 1
elif c == ']': bracket -= 1
elif c == '?' and depth == 0 and bracket == 0: q = i; break
if q < 0: break
depth, bracket, col = 0, 0, -1
for i in range(q + 1, len(e)):
if e[i] == '(': depth += 1
elif e[i] == ')': depth -= 1
elif e[i] == '[': bracket += 1
elif e[i] == ']': bracket -= 1
elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
if col < 0: break
cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
e = f'(({t}) if ({cond}) else ({f}))'
return e
def _assign(lhs: str, rhs: str) -> str:
"""Generate assignment. Bare tmp/SCC/etc get wrapped in Reg()."""
if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
return f"{lhs} = Reg({rhs})"
return f"{lhs} = {rhs}"
def compile_pseudocode(pseudocode: str) -> str:
"""Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
raw_lines = pseudocode.strip().split('\n')
joined_lines: list[str] = []
for line in raw_lines:
line = line.strip()
if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
(joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
else:
joined_lines.append(line)
lines = []
indent, need_pass, in_first_match_loop = 0, False, False
for line in joined_lines:
line = line.strip()
if not line or line.startswith('//'): continue
if line.startswith('if '):
lines.append(' ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
indent += 1
need_pass = True
elif line.startswith('elsif '):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
lines.append(' ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
indent += 1
need_pass = True
elif line == 'else':
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
lines.append(' ' * indent + "else:")
indent += 1
need_pass = True
elif line.startswith('endif'):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
need_pass = False
elif line.startswith('endfor'):
if need_pass: lines.append(' ' * indent + "pass")
indent -= 1
need_pass, in_first_match_loop = False, False
elif line.startswith('declare '):
pass
elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
start, end = _expr(m[2].strip()), _expr(m[3].strip())
lines.append(' ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
indent += 1
need_pass, in_first_match_loop = True, True
elif '=' in line and not line.startswith('=='):
need_pass = False
line = line.rstrip(';')
if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
rhs = _expr(m[1])
lines.append(' ' * indent + f"_full = {rhs}")
lines.append(' ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
lines.append(' ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
if op in line:
lhs, rhs = line.split(op, 1)
lines.append(' ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
break
else:
lhs, rhs = line.split('=', 1)
lhs_s, rhs_s = lhs.strip(), rhs.strip()
stmt = _assign(lhs_s, _expr(rhs_s))
if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
stmt += "; break"
lines.append(' ' * indent + stmt)
if need_pass: lines.append(' ' * indent + "pass")
return '\n'.join(lines)
# ═══════════════════════════════════════════════════════════════════════════════
# GENERATE gen_pcode.py (compiled pseudocode functions)
# ═══════════════════════════════════════════════════════════════════════════════
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: continue
if re.match(r'^\d+ of \d+$', s): continue
if re.match(r'^\d+\.\d+\..*Instructions', s): continue
if s.startswith('"RDNA') or s.startswith('AMD ') or s.startswith('CDNA'): continue
if s.startswith('Notes') or s.startswith('Functional examples'): break
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 =']) or
any(p in s for p in ['D0[', 'D1[', 'S0[', 'S1[', 'S2[']) 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 _parse_pseudocode(instr_texts: list[str], defined_ops: dict, OP_ENUMS: list) -> dict:
"""Parse pseudocode from instruction text(s). Returns {enum_cls: {op: pseudocode}}."""
instructions: dict = {cls: {} for cls in OP_ENUMS}
# Process in reverse order so newer PDFs take priority
for instr_text in reversed(instr_texts):
matches = list(INST_PATTERN.finditer(instr_text))
for i, match in enumerate(matches):
name, opcode = match.group(1), int(match.group(2))
key = (name, opcode)
if key not in defined_ops: continue
start = match.end()
end = matches[i + 1].start() if i + 1 < len(matches) else start + 2000
snippet = instr_text[start:end].strip()
if (pseudocode := _extract_pseudocode(snippet)):
for enum_cls, enum_val in defined_ops[key]:
if enum_val not in instructions[enum_cls]:
instructions[enum_cls][enum_val] = pseudocode
return instructions
UNSUPPORTED = ['SGPR[', 'V_SWAP', 'eval ', 'FATAL_HALT', 'HW_REGISTERS', 'vscnt', 'vmcnt', 'expcnt', 'lgkmcnt',
'CVT_OFF_TABLE', 'ThreadMask', 'S1[i', 'C.i32', 'S[i]', 'in[', 'if n.', 'DST.u32', 'addrd = DST', 'addr = DST']
def _generate_pcode(instr_texts: list[str], arch: str, output_path: str | None = None) -> tuple[int, int]:
"""Generate compiled pseudocode functions. Returns (compiled_count, skipped_count)."""
import importlib
# Load op enums from autogen module
autogen = importlib.import_module(f"extra.assembly.amd.autogen.{arch}")
OP_ENUMS = []
for name in ['SOP1Op', 'SOP2Op', 'SOPCOp', 'SOPKOp', 'SOPPOp', 'VOP1Op', 'VOP2Op', 'VOP3Op', 'VOP3SDOp', 'VOP3POp', 'VOPCOp', 'VOP3AOp', 'VOP3BOp']:
if hasattr(autogen, name): OP_ENUMS.append(getattr(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_')): defined_ops.setdefault((op.name, op.value), []).append((enum_cls, op))
by_cls = _parse_pseudocode(instr_texts, defined_ops, OP_ENUMS)
# Report coverage
total_found, total_ops = 0, 0
for enum_cls in OP_ENUMS:
total = sum(1 for op in enum_cls if op.name.startswith(('S_', 'V_')))
found = len(by_cls.get(enum_cls, {}))
total_found += found
total_ops += total
print(f"{enum_cls.__name__}: {found}/{total} ({100*found//total if total else 0}%)")
print(f"Total: {total_found}/{total_ops} ({100*total_found//total_ops}%)")
# Generate code
enum_names = [e.__name__ for e in OP_ENUMS]
lines = [f'''# autogenerated by generate.py - do not edit
# to regenerate: python -m extra.assembly.amd.generate --arch {arch}
# ruff: noqa: E501,F405,F403
# mypy: ignore-errors
from extra.assembly.amd.autogen.{arch} import {", ".join(enum_names)}
from extra.assembly.amd.pcode import *
''']
compiled_count, skipped_count = 0, 0
for enum_cls in OP_ENUMS:
cls_name = enum_cls.__name__
pseudocode_dict = by_cls.get(enum_cls, {})
if not pseudocode_dict: continue
fn_entries = []
for op, pc in pseudocode_dict.items():
if any(p in pc for p in UNSUPPORTED):
skipped_count += 1
continue
try:
code = compile_pseudocode(pc)
# Hardware behavior fixes (see pcode.py for detailed comments)
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 # denorm check moved to end')
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 # denorm check moved to end')
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)')
# Detect flags for result handling
is_64 = any(p in pc for p in ['D0.u64', 'D0.b64', 'D0.f64', 'D0.i64', 'D1.u64', 'D1.b64', 'D1.f64', 'D1.i64'])
has_d1 = '{ D1' in pc
if has_d1: is_64 = True
is_cmp = (cls_name == 'VOPCOp' or cls_name == 'VOP3Op') and 'D0.u64[laneId]' in pc
is_cmpx = (cls_name == 'VOPCOp' or cls_name == '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_pc = 'PC' in pc
fn_name = f"_{cls_name}_{op.name}"
lines.append(f"def {fn_name}(s0, s1, s2, d0, scc, vcc, lane, exec_mask, literal, VGPR, _vars, src0_idx=0, vdst_idx=0, pc=0):")
for pc_line in pc.split('\n'):
lines.append(f" # {pc_line}")
combined = code + pc
regs = [('S0', 'Reg(s0)'), ('S1', 'Reg(s1)'), ('S2', 'Reg(s2)'),
('D0', 'Reg(s0)' if is_div_scale else 'Reg(d0)'), ('D1', 'Reg(0)'),
('SCC', 'Reg(scc)'), ('VCC', 'Reg(vcc)'), ('EXEC', 'Reg(exec_mask)'),
('tmp', 'Reg(0)'), ('saveexec', 'Reg(exec_mask)'), ('laneId', 'lane'),
('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)'), ('PC', 'Reg(pc)')]
used = {name for name, _ in regs if name in combined}
if 'EXEC_LO' in combined or 'EXEC_HI' in combined: used.add('EXEC')
if 'VCCZ' in combined: used.add('VCC')
if 'EXECZ' in combined: used.add('EXEC')
for name, init in regs:
if name in used: lines.append(f" {name} = {init}")
if 'EXEC_LO' in combined: lines.append(" EXEC_LO = SliceProxy(EXEC, 31, 0)")
if 'EXEC_HI' in combined: lines.append(" EXEC_HI = SliceProxy(EXEC, 63, 32)")
if 'VCCZ' in combined: lines.append(" VCCZ = Reg(1 if VCC._val == 0 else 0)")
if 'EXECZ' in combined: lines.append(" EXECZ = Reg(1 if EXEC._val == 0 else 0)")
lines.append(" # --- compiled pseudocode ---")
for line in code.split('\n'):
lines.append(f" {line}")
lines.append(" # --- end pseudocode ---")
d0_val = "D0._val" if 'D0' in used else "d0"
scc_val = "SCC._val & 1" if 'SCC' in used else "scc & 1"
lines.append(f" result = {{'d0': {d0_val}, 'scc': {scc_val}}}")
if has_sdst:
lines.append(" result['vcc_lane'] = (VCC._val >> lane) & 1")
elif 'VCC' in used:
lines.append(" if VCC._val != vcc: result['vcc_lane'] = (VCC._val >> lane) & 1")
if is_cmpx:
lines.append(" result['exec_lane'] = (EXEC._val >> lane) & 1")
elif 'EXEC' in used:
lines.append(" if EXEC._val != exec_mask: result['exec'] = EXEC._val")
if is_cmp:
lines.append(" result['vcc_lane'] = (D0._val >> lane) & 1")
if is_64:
lines.append(" result['d0_64'] = True")
if has_d1:
lines.append(" result['d1'] = D1._val & 1")
if has_pc:
lines.append(" _pc = PC._val if PC._val < 0x8000000000000000 else PC._val - 0x10000000000000000")
lines.append(" result['new_pc'] = _pc # absolute byte address")
lines.append(" return result")
lines.append("")
fn_entries.append((op, fn_name))
compiled_count += 1
except Exception as e:
print(f" Warning: Failed to compile {op.name}: {e}")
skipped_count += 1
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('}')
lines.append('')
# Add V_WRITELANE_B32 for RDNA
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 # lane select (5 bits for wave32)
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:
cls_name = enum_cls.__name__
if by_cls.get(enum_cls): lines.append(f' {cls_name}: {cls_name}_FUNCTIONS,')
lines.append('}')
lines.append('')
lines.append('def get_compiled_functions(): return COMPILED_FUNCTIONS')
if output_path is not None:
from pathlib import Path
Path(output_path).write_text('\n'.join(lines))
return compiled_count, skipped_count
# ═══════════════════════════════════════════════════════════════════════════════
# MAIN ENTRY POINT
# ═══════════════════════════════════════════════════════════════════════════════
def generate(arch: str = "rdna3", output_dir: str = "extra/assembly/amd/autogen"):
"""Generate both __init__.py and gen_pcode.py for the given architecture."""
urls = PDF_URLS[arch]
if isinstance(urls, str): urls = [urls]
print(f"Parsing PDF(s) for {arch}...")
results = [_parse_single_pdf(url) for url in urls]
if len(results) == 1:
merged = results[0]
doc_name = merged["doc_name"]
instr_texts = [merged["instr_text"]]
else:
merged = _merge_results(results)
doc_name = "+".join(merged["doc_names"])
instr_texts = merged["instr_texts"]
# Generate __init__.py first (needed for pcode generation)
init_path = f"{output_dir}/{arch}/__init__.py"
_generate_dsl(merged, doc_name, init_path)
print(f"Generated {init_path}: SrcEnum ({len(merged['src_enum'])}) + {len(merged['enums'])} opcode enums + {len(merged['formats'])} format classes")
# Generate gen_pcode.py
print("\nCompiling pseudocode functions...")
pcode_path = f"{output_dir}/{arch}/gen_pcode.py"
compiled, skipped = _generate_pcode(instr_texts, arch, pcode_path)
print(f"Generated {pcode_path}: {compiled} compiled, {skipped} skipped")
if __name__ == "__main__":
import argparse, subprocess, sys
parser = argparse.ArgumentParser(description="Generate AMD ISA definitions from PDF")
parser.add_argument("--arch", choices=list(PDF_URLS.keys()) + ["all"], default="rdna3", help="Target architecture")
args = parser.parse_args()
if args.arch == "all":
procs = [subprocess.Popen([sys.executable, "-m", "extra.assembly.amd.generate", "--arch", arch]) for arch in PDF_URLS.keys()]
for p in procs: p.wait()
else:
generate(arch=args.arch)
File diff suppressed because it is too large Load Diff
+319
View File
@@ -0,0 +1,319 @@
# Generic PDF text extractor - no external dependencies
import re, zlib
from tinygrad.helpers import fetch, merge_dicts
PDF_URLS = {
"rdna3": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content",
"rdna4": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content",
"cdna": "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf",
}
# ═══════════════════════════════════════════════════════════════════════════════
# Generic PDF extraction tools
# ═══════════════════════════════════════════════════════════════════════════════
def extract(url: str) -> list[list[tuple[float, float, str, str]]]:
"""Extract positioned text from PDF. Returns list of text elements (x, y, text, font) per page."""
data = fetch(url).read_bytes()
# Parse xref table to locate objects
xref: dict[int, int] = {}
pos = int(re.search(rb'startxref\s+(\d+)', data).group(1)) + 4
while data[pos:pos+7] != b'trailer':
while data[pos:pos+1] in b' \r\n': pos += 1
line_end = data.find(b'\n', pos)
start_obj, count = map(int, data[pos:line_end].split()[:2])
pos = line_end + 1
for i in range(count):
if data[pos+17:pos+18] == b'n' and (off := int(data[pos:pos+10])) > 0: xref[start_obj + i] = off
pos += 20
def get_stream(n: int) -> bytes:
obj = data[xref[n]:data.find(b'endobj', xref[n])]
raw = obj[obj.find(b'stream\n') + 7:obj.find(b'\nendstream')]
return zlib.decompress(raw) if b'/FlateDecode' in obj else raw
# Find page content streams and extract text
pages = []
for n in sorted(xref):
if b'/Type /Page' not in data[xref[n]:xref[n]+500]: continue
if not (m := re.search(rb'/Contents (\d+) 0 R', data[xref[n]:xref[n]+500])): continue
stream = get_stream(int(m.group(1))).decode('latin-1')
elements, font = [], ''
for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
x, y = 0.0, 0.0
for m in re.finditer(r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ', bt.group(1)):
if m.group(1): font = m.group(1)
elif m.group(2): x, y = x + float(m.group(2)), y + float(m.group(3))
elif m.group(4): x, y = float(m.group(4)), float(m.group(5))
elif m.group(6) and (t := bytes.fromhex(m.group(6)).decode('latin-1')).strip(): elements.append((x, y, t, font))
elif m.group(7) and (t := ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', m.group(7)))).strip(): elements.append((x, y, t, font))
pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
return pages
def extract_tables(pages: list[list[tuple[float, float, str, str]]]) -> dict[int, tuple[str, list[list[str]]]]:
"""Extract numbered tables from PDF pages. Returns {table_num: (title, rows)} where rows is list of cells per row."""
def group_by_y(texts, key=lambda y: round(y)):
by_y: dict[int, list[tuple[float, float, str]]] = {}
for x, y, t, _ in texts:
by_y.setdefault(key(y), []).append((x, y, t))
return by_y
# Find all table headers by merging text on same line
table_positions = []
for page_idx, texts in enumerate(pages):
for items in group_by_y(texts).values():
line = ''.join(t for _, t in sorted((x, t) for x, _, t in items))
if m := re.search(r'Table (\d+)\. (.+)', line):
table_positions.append((int(m.group(1)), m.group(2).strip(), page_idx, items[0][1]))
table_positions.sort(key=lambda t: (t[2], -t[3]))
# For each table, find rows with matching X positions
result: dict[int, tuple[str, list[list[str]]]] = {}
for num, title, start_page, header_y in table_positions:
rows, col_xs = [], None
for page_idx in range(start_page, len(pages)):
page_texts = [(x, y, t) for x, y, t, _ in pages[page_idx] if 30 < y < 760 and (page_idx > start_page or y < header_y)]
for items in sorted(group_by_y([(x, y, t, '') for x, y, t in page_texts], key=lambda y: round(y / 5)).values(), key=lambda items: -items[0][1]):
xs = tuple(sorted(round(x) for x, _, _ in items))
if col_xs is None:
if len(xs) < 2: continue # Skip single-column rows before table starts
col_xs = xs
elif len(xs) == 1 and xs[0] in col_xs: continue # Skip continuation rows at known column positions
elif not any(c in xs for c in col_xs[:2]): break # Row missing first columns = end of table
rows.append([t for _, t in sorted((x, t) for x, _, t in items)])
else: continue
break
if rows: result[num] = (title, rows)
return result
# ═══════════════════════════════════════════════════════════════════════════════
# AMD specific extraction
# ═══════════════════════════════════════════════════════════════════════════════
def extract_enums(tables: dict[int, tuple[str, list[list[str]]]]) -> dict[str, dict[int, str]]:
"""Extract all enums from tables. Returns {enum_name: {value: name}}."""
enums: dict[str, dict[int, str]] = {}
for num, (title, rows) in tables.items():
# Opcode enums from "XXX Opcodes" tables
if m := re.match(r'(\w+) (?:Y-)?Opcodes', title):
fmt_name = 'VOPD' if 'Y-Opcodes' in title else m.group(1)
ops: dict[int, str] = {}
for row in rows:
for i in range(0, len(row) - 1, 2):
if row[i].isdigit() and re.match(r'^[A-Z][A-Z0-9_]+$', row[i + 1]):
ops[int(row[i])] = row[i + 1]
if ops: enums[fmt_name] = ops
# BufFmt from "Data Format" tables
if 'Data Format' in title:
for row in rows:
for i in range(0, len(row) - 1, 2):
if row[i].isdigit() and re.match(r'^[\dA-Z_]+$', row[i + 1]) and 'INVALID' not in row[i + 1]:
enums.setdefault('BufFmt', {})[int(row[i])] = row[i + 1]
return enums
def extract_ins(tables: dict[int, tuple[str, list[list[str]]]]) -> tuple[dict[str, list[tuple[str, int, int]]], dict[str, str]]:
"""Extract formats and encodings from 'XXX Fields' tables. Returns (formats, encodings)."""
formats: dict[str, list[tuple[str, int, int]]] = {}
encodings: dict[str, str] = {}
for num, (title, rows) in tables.items():
if not (m := re.match(r'(\w+) Fields$', title)): continue
fmt_name = m.group(1)
fields = []
for row in rows:
if len(row) < 2: continue
if (bits := re.match(r'\[?(\d+):(\d+)\]?$', row[1])) or (bits := re.match(r'\[(\d+)\]$', row[1])):
field_name = row[0].lower()
hi, lo = int(bits.group(1)), int(bits.group(2)) if bits.lastindex >= 2 else int(bits.group(1))
if field_name == 'encoding' and len(row) >= 3:
enc_bits = None
if "'b" in row[2]: enc_bits = row[2].split("'b")[-1].replace('_', '')
elif (enc := re.search(r':\s*([01_]+)', row[2])): enc_bits = enc.group(1).replace('_', '')
if enc_bits:
# If encoding bits exceed field width, extend field to match (AMD docs sometimes have this)
declared_width, actual_width = hi - lo + 1, len(enc_bits)
if actual_width > declared_width: lo = hi - actual_width + 1
encodings[fmt_name] = enc_bits
fields.append((field_name, hi, lo))
if fields: formats[fmt_name] = fields
return formats, encodings
def extract_pcode(pages: list[list[tuple[float, float, str, str]]], enums: dict[str, dict[int, str]]) -> dict[tuple[str, int], str]:
"""Extract pseudocode for instructions. Returns {(name, opcode): pseudocode}."""
# Build lookup from instruction name to opcode
name_to_op = {name: op for ops in enums.values() for op, name in ops.items()}
# First pass: find all instruction headers across all pages
all_instructions: list[tuple[int, float, str, int]] = [] # (page_idx, y, name, opcode)
for page_idx, page in enumerate(pages):
by_y: dict[int, list[tuple[float, str]]] = {}
for x, y, t, _ in page:
by_y.setdefault(round(y), []).append((x, t))
for y, items in sorted(by_y.items(), reverse=True):
left = [(x, t) for x, t in items if 55 < x < 65]
right = [(x, t) for x, t in items if 535 < x < 550]
if left and right and left[0][1] in name_to_op and right[0][1].isdigit():
all_instructions.append((page_idx, y, left[0][1], int(right[0][1])))
# Second pass: extract pseudocode between consecutive instructions
pcode: dict[tuple[str, int], str] = {}
for i, (page_idx, y, name, opcode) in enumerate(all_instructions):
# Get end boundary from next instruction
if i + 1 < len(all_instructions):
next_page, next_y = all_instructions[i + 1][0], all_instructions[i + 1][1]
else:
next_page, next_y = page_idx, 0
# Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
lines = []
for p in range(page_idx, next_page + 1):
start_y = y if p == page_idx else 800
end_y = next_y if p == next_page else 0
lines.extend((p, y2, t) for x, y2, t, f in pages[p] if f in ('/F6.0', '/F7.0') and end_y < y2 < start_y and 60 < x < 80)
if lines:
# Sort by page first, then by y descending within each page (higher y = earlier text in PDF)
sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
# Stop at large Y gaps (>30) - indicates section break (Notes, examples, etc)
filtered = [sorted_lines[0]]
for j in range(1, len(sorted_lines)):
prev_page, prev_y, _ = sorted_lines[j-1]
curr_page, curr_y, _ = sorted_lines[j]
if curr_page == prev_page and prev_y - curr_y > 30: break
if curr_page != prev_page and prev_y > 60 and curr_y < 730: break # examples spilled to next page (not at very top)
filtered.append(sorted_lines[j])
pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
return pcode
# ═══════════════════════════════════════════════════════════════════════════════
# Write autogen files
# ═══════════════════════════════════════════════════════════════════════════════
def write_enums(enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write enum.py file from extracted enums."""
lines = ["# autogenerated from AMD ISA PDF by pdf.py - do not edit", "from enum import IntEnum", ""]
for name, values in sorted(enums.items()):
suffix = "Op" if name not in ('Src', 'BufFmt') else ("Enum" if name == 'Src' else "")
prefix = "BUF_FMT_" if name == 'BufFmt' else ""
lines.append(f"class {name}{suffix}(IntEnum):")
for val, member in sorted(values.items()):
lines.append(f" {prefix}{member} = {val}")
lines.append("")
with open(path, "w") as f:
f.write("\n".join(lines))
def write_ins(formats: dict[str, list[tuple[str, int, int]]], encodings: dict[str, str], enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write ins.py file from extracted formats and enums."""
# Field types and ordering
def field_type(name, fmt):
if name == 'op' and fmt in enums: return f'Annotated[BitField, {fmt}Op]'
if name in ('opx', 'opy'): return 'Annotated[BitField, VOPDOp]'
if name == 'vdsty': return 'VDSTYEnc'
if name in ('vdst', 'vsrc1', 'vaddr', 'vdata', 'data', 'data0', 'data1', 'addr', 'vsrc0', 'vsrc2', 'vsrc3'): return 'VGPRField'
if name in ('sdst', 'sbase', 'sdata', 'srsrc', 'ssamp'): return 'SGPRField'
if name.startswith('ssrc') or name in ('saddr', 'soffset'): return 'SSrc'
if name in ('src0', 'srcx0', 'srcy0') or name.startswith('src') and name[3:].isdigit(): return 'Src'
if name.startswith('simm'): return 'SImm'
if name == 'offset' or name.startswith('imm'): return 'Imm'
return None
field_priority = ['encoding', 'op', 'opx', 'opy', 'vdst', 'vdstx', 'vdsty', 'sdst', 'vdata', 'sdata', 'addr', 'vaddr', 'data', 'data0', 'data1',
'src0', 'srcx0', 'srcy0', 'vsrc0', 'ssrc0', 'src1', 'vsrc1', 'vsrcx1', 'vsrcy1', 'ssrc1', 'src2', 'vsrc2', 'src3', 'vsrc3',
'saddr', 'sbase', 'srsrc', 'ssamp', 'soffset', 'offset', 'simm16', 'en', 'target', 'attr', 'attr_chan',
'omod', 'neg', 'neg_hi', 'abs', 'clmp', 'opsel', 'opsel_hi', 'waitexp', 'wait_va',
'dmask', 'dim', 'seg', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe', 'unrm', 'done', 'row']
def sort_fields(fields):
order = {name: i for i, name in enumerate(field_priority)}
return sorted(fields, key=lambda f: (order.get(f[0], 1000), f[2]))
# Generate format classes
lines = ["# autogenerated from AMD ISA PDF by pdf.py - do not edit", "# ruff: noqa: F401,F403",
"from typing import Annotated",
"from extra.assembly.amd.dsl import *",
f"from extra.assembly.amd.autogen.{arch}.enum import *", "import functools", ""]
for fmt_name, fields in sorted(formats.items()):
lines.append(f"class {fmt_name}(Inst):")
for name, hi, lo in sort_fields(fields):
bits_str = f"bits[{hi}:{lo}]" if hi != lo else f"bits[{hi}]"
if name == 'encoding' and fmt_name in encodings: lines.append(f" encoding = {bits_str} == 0b{encodings[fmt_name]}")
else:
ftype = field_type(name, fmt_name)
lines.append(f" {name}{f':{ftype}' if ftype else ''} = {bits_str}")
lines.append("")
# Generate instruction helpers
lines.append("# instruction helpers")
for fmt_name, ops in sorted(enums.items()):
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=1"}.get(fmt_name, "")
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt_name, f"{fmt_name}, {fmt_name}Op")
suffix = "_e32" if fmt_name in ("VOP1", "VOP2", "VOPC") else "_e64" if fmt_name == "VOP3" and len(ops) > 0 else ""
if fmt_name in formats or fmt_name in ("GLOBAL", "SCRATCH"):
for op_val, name in sorted(ops.items()):
fn_suffix = suffix if fmt_name != "VOP3" or op_val < 512 else ""
lines.append(f"{name.lower()}{fn_suffix} = functools.partial({tgt}.{name}{seg})")
with open(path, "w") as f:
f.write("\n".join(lines))
def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: str):
"""Write str_pcode.py file from extracted pseudocode."""
# Group pseudocode by enum class
by_enum: dict[str, list[tuple[str, int, str]]] = {}
for fmt_name, ops in enums.items():
for opcode, name in ops.items():
if (name, opcode) in pcode: by_enum.setdefault(f"{fmt_name}Op", []).append((name, opcode, pcode[(name, opcode)]))
# Generate file
enum_names = sorted(by_enum.keys())
lines = [f"# autogenerated by pdf.py - do not edit", f"# to regenerate: python -m extra.assembly.amd.pdf",
"# ruff: noqa: E501", f"from extra.assembly.amd.autogen.{arch}.enum import {', '.join(enum_names)}", ""]
for enum_name in enum_names:
lines.append(f"{enum_name}_PCODE = {{")
for name, opcode, code in sorted(by_enum[enum_name], key=lambda x: x[1]):
lines.append(f" {enum_name}.{name}: {code!r},")
lines.append("}\n")
lines.append(f"PSEUDOCODE_STRINGS = {{{', '.join(f'{e}: {e}_PCODE' for e in enum_names)}}}")
with open(path, "w") as f:
f.write("\n".join(lines))
if __name__ == "__main__":
import pathlib
for arch, url in PDF_URLS.items():
print(f"Processing {arch}...")
pages = extract(url)
tables = extract_tables(pages)
enums = extract_enums(tables)
formats, encodings = extract_ins(tables)
pcode = extract_pcode(pages, enums)
# Fix known PDF errors
if arch == 'rdna3':
fixes = {'SOPP': {8: 'S_WAITCNT_DEPCTR', 58: 'S_TTRACEDATA', 59: 'S_TTRACEDATA_IMM'},
'SOPK': {22: 'S_SUBVECTOR_LOOP_BEGIN', 23: 'S_SUBVECTOR_LOOP_END'},
'SMEM': {34: 'S_ATC_PROBE', 35: 'S_ATC_PROBE_BUFFER'},
'DS': {24: 'DS_GWS_SEMA_RELEASE_ALL', 25: 'DS_GWS_INIT', 26: 'DS_GWS_SEMA_V', 27: 'DS_GWS_SEMA_BR', 28: 'DS_GWS_SEMA_P', 29: 'DS_GWS_BARRIER'},
'FLAT': {40: 'GLOBAL_LOAD_ADDTID_B32', 41: 'GLOBAL_STORE_ADDTID_B32', 55: 'FLAT_ATOMIC_CSUB_U32'}}
for fmt, ops in fixes.items(): enums[fmt] = merge_dicts([enums[fmt], ops])
if arch == 'rdna4':
fixes = {'SMEM': {34: 'S_ATC_PROBE', 35: 'S_ATC_PROBE_BUFFER'},
'SOP1': {81: 'S_BARRIER_INIT', 82: 'S_BARRIER_JOIN'},
'SOPP': {21: 'S_BARRIER_LEAVE', 58: 'S_TTRACEDATA', 59: 'S_TTRACEDATA_IMM'}}
for fmt, ops in fixes.items(): enums[fmt] = merge_dicts([enums[fmt], ops])
if arch in ('rdna3', 'rdna4'):
# RDNA SMEM: PDF says DLC=[14], GLC=[16] but hardware uses DLC=[13], GLC=[14]
if 'SMEM' in formats:
formats['SMEM'] = [(n, 13 if n == 'dlc' else 14 if n == 'glc' else h, 13 if n == 'dlc' else 14 if n == 'glc' else l)
for n, h, l in formats['SMEM']]
if arch == 'cdna':
# CDNA DS: PDF is missing the GDS field (bit 16)
if 'DS' in formats and not any(n == 'gds' for n, _, _ in formats['DS']):
formats['DS'].append(('gds', 16, 16))
# CDNA DPP/SDWA: PDF only documents modifier fields (bits[63:32]), need to add VOP overlay fields (bits[31:0])
vop_overlay = [('encoding', 8, 0), ('vop_op', 16, 9), ('vdst', 24, 17), ('vop2_op', 31, 25)]
if 'DPP' in formats and not any(n == 'encoding' for n, _, _ in formats['DPP']):
formats['DPP'] = vop_overlay + [('bc' if n == 'bound_ctrl' else n, h, l) for n, h, l in formats['DPP']]
encodings['DPP'] = '11111010'
if 'SDWA' in formats and not any(n == 'encoding' for n, _, _ in formats['SDWA']):
formats['SDWA'] = vop_overlay + [(n, h, l) for n, h, l in formats['SDWA']]
encodings['SDWA'] = '11111001'
base = pathlib.Path(__file__).parent / "autogen" / arch
write_enums(enums, arch, base / "enum.py")
write_ins(formats, encodings, enums, arch, base / "ins.py")
write_pcode(pcode, enums, arch, base / "str_pcode.py")
print(f" {len(tables)} tables, {len(pcode)} pcode -> {base}")
+381
View File
@@ -0,0 +1,381 @@
"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
This module provides encoding and decoding of raw SQTT byte streams.
The format is nibble-based with variable-width packets determined by a state machine.
Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
"""
from __future__ import annotations
from enum import IntEnum
from typing import get_type_hints
from extra.assembly.amd.dsl import BitField, bits
# ═══════════════════════════════════════════════════════════════════════════════
# FIELD ENUMS
# ═══════════════════════════════════════════════════════════════════════════════
class MemSrc(IntEnum):
LDS = 0
LDS_ALT = 1
VMEM = 2
VMEM_ALT = 3
class AluSrc(IntEnum):
NONE = 0
SALU = 1
VALU = 2
VALU_ALT = 3
class InstOp(IntEnum):
"""SQTT instruction operation types.
Memory ops appear in two ranges depending on which SIMD executes them:
- 0x1x-0x2x range: ops on traced SIMD
- 0x5x range: ops on other SIMD (OTHER_ prefix)
GLOBAL memory ops encoding depends on addressing mode AND size:
- Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
- Stores: base + size_offset, where VADDR is shifted +1 from SADDR
SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
OTHER_ range follows same pattern but values overlap differently.
"""
SALU = 0x0
SMEM = 0x1
JUMP = 0x3 # branch taken
JUMP_NO = 0x4 # branch not taken
MESSAGE = 0x9
VALU_TRANS = 0xb # transcendental: exp, log, rcp, sqrt, sin, cos
VALU_64_SHIFT = 0xd # 64-bit shifts: lshl, lshr, ashr
VALU_MAD64 = 0xe # 64-bit multiply-add
VALU_64 = 0xf # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
VINTERP = 0x12 # interpolation: v_interp_p10_f32, v_interp_p2_f32
BARRIER = 0x13
# FLAT memory ops on traced SIMD (0x1x range)
FLAT_LOAD = 0x1c
FLAT_STORE = 0x1d
FLAT_STORE_64 = 0x1e
FLAT_STORE_96 = 0x1f
FLAT_STORE_128 = 0x20
# GLOBAL memory ops on traced SIMD (0x2x range)
GLOBAL_LOAD = 0x21 # saddr=SGPR, all sizes
GLOBAL_LOAD_VADDR = 0x22 # saddr=NULL, all sizes
GLOBAL_STORE = 0x24 # saddr=SGPR, 32-bit
GLOBAL_STORE_64 = 0x25 # saddr=SGPR 64 or saddr=NULL 32
GLOBAL_STORE_96 = 0x26 # saddr=SGPR 96 or saddr=NULL 64
GLOBAL_STORE_128 = 0x27 # saddr=SGPR 128 or saddr=NULL 96
GLOBAL_STORE_VADDR_128 = 0x28 # saddr=NULL, 128-bit
# LDS ops on traced SIMD
LDS_LOAD = 0x29
LDS_STORE = 0x2b
LDS_STORE_64 = 0x2c
LDS_STORE_128 = 0x2e
# Memory ops on other SIMD (0x5x range)
OTHER_LDS_LOAD = 0x50
OTHER_LDS_STORE = 0x51
OTHER_LDS_STORE_64 = 0x52
OTHER_LDS_STORE_128 = 0x54
OTHER_FLAT_LOAD = 0x55
OTHER_FLAT_STORE = 0x56
OTHER_FLAT_STORE_64 = 0x57
OTHER_FLAT_STORE_96 = 0x58
OTHER_FLAT_STORE_128 = 0x59
OTHER_GLOBAL_LOAD = 0x5a # saddr=SGPR, all sizes
OTHER_GLOBAL_LOAD_VADDR = 0x5b # saddr=NULL or saddr=SGPR store 32
OTHER_GLOBAL_STORE_64 = 0x5c # saddr=SGPR 64 or saddr=NULL 32
OTHER_GLOBAL_STORE_96 = 0x5d # saddr=SGPR 96 or saddr=NULL 64
OTHER_GLOBAL_STORE_128 = 0x5e # saddr=SGPR 128 or saddr=NULL 96
OTHER_GLOBAL_STORE_VADDR_128 = 0x5f # saddr=NULL, 128-bit
# EXEC-modifying ops (0x7x range)
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
VALU_CMPX = 0x73 # v_cmpx_*
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE BASE CLASS
# ═══════════════════════════════════════════════════════════════════════════════
class PacketType:
"""Base class for SQTT packet types."""
_encoding: tuple[BitField, int] | None = None
_field_types: dict[str, type] = {}
_values: dict[str, int]
_raw: int
_time: int
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
if 'encoding' in cls.__dict__ and isinstance(cls.__dict__['encoding'], tuple):
cls._encoding = cls.__dict__['encoding']
# Cache field type annotations for enum conversion
try: cls._field_types = {k: v for k, v in get_type_hints(cls).items() if isinstance(v, type) and issubclass(v, IntEnum)}
except Exception: cls._field_types = {}
# Cache fields and precompute extraction info: (name, lo, mask, enum_type)
cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField) and k != 'encoding'}
cls._extract_info = [(name, bf.lo, bf.mask(), cls._field_types.get(name)) for name, bf in cls._fields.items()]
cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4)
@classmethod
def from_raw(cls, raw: int, time: int = 0):
inst = object.__new__(cls)
inst._raw, inst._time, inst._values = raw, time, {}
for name, lo, mask, enum_type in cls._extract_info:
val = (raw >> lo) & mask
if enum_type is not None:
try: val = enum_type(val)
except ValueError: pass
inst._values[name] = val
return inst
def __getattr__(self, name: str):
if name.startswith('_'): raise AttributeError(name)
return self._values.get(name, 0)
def __repr__(self) -> str:
fields_str = ", ".join(f"{k}={v}" for k, v in self._values.items() if not k.startswith('_'))
return f"{self.__class__.__name__}({fields_str})"
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
class VALUINST(PacketType): # exclude: 1 << 2
encoding = bits[2:0] == 0b011
delta = bits[5:3]
flag = bits[6:6]
wave = bits[11:7]
class VMEMEXEC(PacketType): # exclude: 1 << 0
encoding = bits[3:0] == 0b1111
delta = bits[5:4]
src: MemSrc = bits[7:6]
class ALUEXEC(PacketType): # exclude: 1 << 1
encoding = bits[3:0] == 0b1110
delta = bits[5:4]
src: AluSrc = bits[7:6]
class IMMEDIATE(PacketType): # exclude: 1 << 5
encoding = bits[3:0] == 0b1101
delta = bits[6:4]
wave = bits[11:7]
class IMMEDIATE_MASK(PacketType): # exclude: 1 << 5
encoding = bits[4:0] == 0b00100
delta = bits[7:5]
mask = bits[23:8]
class WAVERDY(PacketType): # exclude: 1 << 3
encoding = bits[4:0] == 0b10100
delta = bits[7:5]
mask = bits[23:8]
class TS_DELTA_S8_W3(PacketType):
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[63:11]
class WAVEEND(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b10101
delta = bits[7:5]
flag7 = bits[8:8]
simd = bits[10:9]
cu_lo = bits[13:11]
wave = bits[19:15]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class WAVESTART(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b01100
delta = bits[6:5]
flag7 = bits[7:7]
simd = bits[9:8]
cu_lo = bits[12:10]
wave = bits[17:13]
id7 = bits[31:18]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[47:7]
class WAVEALLOC(PacketType): # exclude: 1 << 10
encoding = bits[4:0] == 0b00101
delta = bits[7:5]
_padding = bits[19:8]
class TS_DELTA_S5_W3(PacketType):
encoding = bits[4:0] == 0b00110
delta = bits[7:5]
_padding = bits[51:8]
class PERF(PacketType): # exclude: 1 << 11
encoding = bits[4:0] == 0b10110
delta = bits[7:5]
arg = bits[27:8]
class TS_DELTA_SHORT(PacketType):
encoding = bits[3:0] == 0b1000
delta = bits[7:4]
class NOP(PacketType):
encoding = bits[3:0] == 0b0000
delta = None # type: ignore
_padding = bits[3:0]
class TS_WAVE_STATE(PacketType):
encoding = bits[6:0] == 0b1010001
delta = bits[15:7]
coarse = bits[23:16]
@property
def wave_interest(self) -> bool: return bool(self.coarse & 1)
@property
def terminate_all(self) -> bool: return bool(self.coarse & 8)
class EVENT(PacketType): # exclude: 1 << 7
encoding = bits[7:0] == 0b01100001
delta = bits[10:8]
event = bits[23:11]
class EVENT_BIG(PacketType):
encoding = bits[7:0] == 0b11100001
delta = bits[10:8]
event = bits[31:11]
class REG(PacketType):
encoding = bits[3:0] == 0b1001
delta = bits[6:4]
slot = bits[9:7]
hi_byte = bits[15:8]
subop = bits[31:16]
val32 = bits[63:32]
@property
def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
class SNAPSHOT(PacketType):
encoding = bits[6:0] == 0b1110001
delta = bits[9:7]
snap = bits[63:10]
class TS_DELTA_OR_MARK(PacketType):
encoding = bits[6:0] == 0b0000001
delta = bits[47:12]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
class LAYOUT_HEADER(PacketType):
encoding = bits[6:0] == 0b0010001
delta = None # type: ignore
layout = bits[12:7]
simd = bits[14:13]
group = bits[17:15]
sel_a = bits[31:28]
sel_b = bits[36:33]
flag4 = bits[59:59]
_padding = bits[63:60]
class INST(PacketType):
encoding = bits[2:0] == 0b010
delta = bits[6:4]
flag1 = bits[3:3]
flag2 = bits[7:7]
wave = bits[12:8]
op: InstOp = bits[19:13]
class UTILCTR(PacketType):
encoding = bits[6:0] == 0b0110001
delta = bits[8:7]
ctr = bits[47:9]
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
PACKET_TYPES: list[type[PacketType]] = [
EVENT, EVENT_BIG,
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
VALUINST, INST,
NOP,
]
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
for byte_val in range(256):
for opcode, pkt_cls in enumerate(PACKET_TYPES):
if pkt_cls._encoding is None: continue
mask_bf, pattern = pkt_cls._encoding
if (byte_val & mask_bf.mask()) == pattern:
table[byte_val] = opcode
break
return bytes(table), opcode_to_class
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
# Precompute special case opcodes
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
_TS_DELTA_OR_MARK_BIT8 = (TS_DELTA_OR_MARK.bit8.lo, TS_DELTA_OR_MARK.bit8.mask())
_TS_DELTA_OR_MARK_BIT9 = (TS_DELTA_OR_MARK.bit9.lo, TS_DELTA_OR_MARK.bit9.mask())
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
_DECODE_INFO: dict[int, tuple] = {}
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
_delta_field = getattr(_pkt_cls, 'delta', None)
_delta_lo = _delta_field.lo if _delta_field else 0
_delta_mask = _delta_field.mask() if _delta_field else 0
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
# ═══════════════════════════════════════════════════════════════════════════════
# DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def decode(data: bytes) -> list[PacketType]:
"""Decode raw SQTT blob into list of packet instances."""
packets: list[PacketType] = []
packets_append = packets.append
n = len(data)
reg = 0
offset = 0
nib_count = 16
time = 0
state_to_opcode = STATE_TO_OPCODE
decode_info = _DECODE_INFO
mask64 = (1 << 64) - 1
while (offset >> 3) < n:
target = offset + nib_count * 4
while offset < target and (offset >> 3) < n:
byte = data[offset >> 3]
nib = (byte >> (offset & 4)) & 0xF
reg = ((reg >> 4) | (nib << 60)) & mask64
offset += 4
if offset < target: break
opcode = state_to_opcode[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
delta = (reg >> delta_lo) & delta_mask
if special == 1: # TS_DELTA_OR_MARK
bit8 = (reg >> _TS_DELTA_OR_MARK_BIT8[0]) & _TS_DELTA_OR_MARK_BIT8[1]
bit9 = (reg >> _TS_DELTA_OR_MARK_BIT9[0]) & _TS_DELTA_OR_MARK_BIT9[1]
if bit9 and not bit8: delta = 0
elif special == 2: # TS_DELTA_SHORT
delta = delta + 8
time += delta
packets_append(pkt_cls.from_raw(reg, time))
return packets
+46 -148
View File
@@ -1,13 +1,12 @@
#!/usr/bin/env python3
"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
import ctypes, time, os, struct, cProfile, pstats, io
"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
import ctypes, time, os
from pathlib import Path
from typing import Callable
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
@@ -42,7 +41,7 @@ def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = Non
ranges.add((args_ptr, ctypes.sizeof(args)))
return buffers, args, args_ptr, ranges
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
"""Benchmark an emulator and return average time."""
gx, gy, gz = global_size
lx, ly, lz = local_size
@@ -50,13 +49,13 @@ def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size
lib_ptr = ctypes.addressof(kernel_buf)
# Warmup
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
# Timed runs
times = []
for _ in range(iterations):
start = time.perf_counter()
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
end = time.perf_counter()
if result != 0:
print(f" {name} returned error: {result}")
@@ -65,27 +64,12 @@ def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size
return sum(times) / len(times)
def create_synthetic_kernel(n_ops: int) -> bytes:
"""Create a synthetic kernel with n_ops vector operations."""
instructions = []
# VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
ops = [
(0b0000011 << 25) | (1 << 17) | (0 << 9) | 256, # v_add_f32 v0, v0, v1
(0b0001000 << 25) | (1 << 17) | (0 << 9) | 256, # v_mul_f32 v0, v0, v1
(0b0010000 << 25) | (1 << 17) | (0 << 9) | 256, # v_max_f32 v0, v0, v1
(0b0001111 << 25) | (1 << 17) | (0 << 9) | 256, # v_min_f32 v0, v0, v1
]
for i in range(n_ops):
instructions.append(ops[i % len(ops)])
# S_ENDPGM
instructions.append((0b101111111 << 23) | (48 << 16) | 0)
return b''.join(struct.pack('<I', inst) for inst in instructions)
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
try:
from tinygrad import Tensor
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.autogen import hsa
import numpy as np
np.random.seed(42)
@@ -112,7 +96,9 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
lowered = ei.lower()
if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
lib = bytes(lowered.prg.p.lib)
image = memoryview(bytearray(lib))
_, sections, _ = elf_loader(lib)
rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
for sec in sections:
if sec.name == '.text':
buf_sizes = [b.nbytes for b in lowered.bufs]
@@ -122,67 +108,22 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
try: buf_data[i] = bytes(buf.base._buf)
except: pass
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
# Extract rsrc2 from ELF (same as ops_amd.py)
group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
lds_size = ((group_segment_size + 511) // 512) & 0x1FF
code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
return None
except Exception as e:
print(f" Error getting kernel: {e}")
return None
def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
"""Profile the Python emulator to find bottlenecks."""
gx, gy, gz = global_size
lx, ly, lz = local_size
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
lib_ptr = ctypes.addressof(kernel_buf)
pr = cProfile.Profile()
pr.enable()
for _ in range(n_runs):
python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
pr.disable()
s = io.StringIO()
ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
ps.print_stats(20)
return s.getvalue()
def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
"""Measure raw step_wave() performance (steps per second)."""
program = decode_program(kernel)
if not program: return 0.0
st = WaveState()
st.exec_mask = 0xffffffff
lds = bytearray(65536)
n_lanes = 32
# Reset PC for each measurement
start = time.perf_counter()
for _ in range(n_steps):
st.pc = 0
while st.pc in program:
result = step_wave(program, st, lds, n_lanes)
if result == -1: break
elapsed = time.perf_counter() - start
return n_steps / elapsed if elapsed > 0 else 0
# Test configurations
SYNTHETIC_TESTS = [
("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
]
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
def main():
import argparse
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
args = parser.parse_args()
@@ -197,98 +138,55 @@ def main():
results = []
# Synthetic workloads
if not args.tinygrad_only:
print("\n[SYNTHETIC WORKLOADS]")
print("-" * 90)
print("\n[TINYGRAD KERNELS]")
print("-" * 90)
for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
kernel = create_synthetic_kernel(n_ops)
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
for op_name in TINYGRAD_TESTS:
print(f"\n{op_name}:", end=" ", flush=True)
kernel_info = get_tinygrad_kernel(op_name)
if kernel_info is None:
print("failed to compile")
continue
print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
buf_sizes = [4096]
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
set_valid_mem_ranges(ranges)
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
# Benchmark
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
# Tinygrad kernels
if not args.synthetic_only:
print("\n[TINYGRAD KERNELS]")
print("-" * 90)
for op_name in TINYGRAD_TESTS:
print(f"\n{op_name}:", end=" ", flush=True)
kernel_info = get_tinygrad_kernel(op_name)
if kernel_info is None:
print("failed to compile")
continue
kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
# Optional profiling
if args.profile and py_time:
print("\n [PROFILE - Top 10 functions]")
profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
for line in profile_output.split('\n')[5:15]:
if line.strip(): print(f" {line}")
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
# Summary table
print("\n" + "=" * 90)
print("SUMMARY")
print("=" * 90)
print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print("-" * 90)
for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
for name, n_insts, n_wgs, py_time, rust_time in results:
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
if rust_time:
rust_ms = f"{rust_time*1000:.3f}"
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
else:
rust_ms, speedup = "N/A", "N/A"
print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
if __name__ == "__main__":
main()
@@ -3,7 +3,7 @@
# Currently many of these tests fail - they document desired behavior
import unittest
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
class TestRegisterSliceSyntax(unittest.TestCase):
+42
View File
@@ -22,3 +22,45 @@ def get_llvm_objdump():
for p in ['llvm-objdump', 'llvm-objdump-21', 'llvm-objdump-20']:
if shutil.which(p): return p
raise FileNotFoundError("llvm-objdump not found")
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION CONTEXT (for testing compiled pseudocode)
# ═══════════════════════════════════════════════════════════════════════════════
class ExecContext:
"""Context for running compiled pseudocode in tests."""
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
self.D0, self.D1 = Reg(d0), Reg(0)
self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
self.lane, self.laneId, self.literal = lane, lane, literal
self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
self.VGPR = vgprs if vgprs is not None else {}
self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
def run(self, code: str):
"""Execute compiled code."""
import extra.assembly.amd.pcode as pcode
ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
# Also include underscore-prefixed helpers that compiled pseudocode uses
for k in ['_pack', '_pack32']:
if hasattr(pcode, k): ns[k] = getattr(pcode, k)
ns.update({
'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
'tmp': self.tmp, 'saveexec': self.saveexec,
'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
})
exec(code, ns)
def _sync(ctx_reg, ns_val):
if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
else: ctx_reg._val = int(ns_val) & self._MASK64
for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
+1
View File
@@ -0,0 +1 @@
"""Hardware-validated emulator tests for RDNA3 instructions."""
+202
View File
@@ -0,0 +1,202 @@
"""Test infrastructure for hardware-validated RDNA3 emulator tests.
Uses run_asm() with memory output, so tests can run on both emulator and real hardware.
Set USE_HW=1 to run on both emulator and real hardware, comparing results.
"""
import ctypes, os, struct
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import RawImm
from extra.assembly.amd.emu import WaveState, run_asm, set_valid_mem_ranges
from extra.assembly.amd.pcode import _i32, _f32
VCC = SrcEnum.VCC_LO # For VOP3SD sdst field
USE_HW = os.environ.get("USE_HW", "0") == "1"
FLOAT_TOLERANCE = 1e-5
# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc
N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4 # 16 regs * 32 lanes * 4 bytes = 2048
SGPR_BYTES = N_SGPRS * 4 # 16 regs * 4 bytes = 64
OUT_BYTES = VGPR_BYTES + SGPR_BYTES + 8 # + vcc + scc
# Float conversion helpers
def f2i(f: float) -> int: return _i32(f)
def i2f(i: int) -> float: return _f32(i)
def f2i64(f: float) -> int: return struct.unpack('<Q', struct.pack('<d', f))[0]
def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
def assemble(instructions: list) -> bytes:
return b''.join(inst.to_bytes() for inst in instructions)
def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
"""Generate prologue and epilogue instructions for state capture."""
prologue = [
s_mov_b32(s[80], s[0]),
s_mov_b32(s[81], s[1]),
v_mov_b32_e32(v[255], v[0]),
]
for i in range(N_VGPRS):
prologue.append(v_mov_b32_e32(v[i], 0))
for i in range(N_SGPRS):
prologue.append(s_mov_b32(s[i], 0))
prologue.append(s_mov_b32(s[SrcEnum.VCC_LO - 128], 0))
epilogue = [
s_mov_b32(s[90], SrcEnum.VCC_LO),
s_cselect_b32(s[91], 1, 0),
s_load_b64(s[92:93], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_lshlrev_b32_e32(v[240], 2, v[255]),
]
for i in range(N_VGPRS):
epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92], offset=i * WAVE_SIZE * 4))
epilogue.append(v_mov_b32_e32(v[241], 0))
epilogue.append(v_cmp_eq_u32_e32(v[255], v[241]))
epilogue.append(s_and_saveexec_b32(s[94], SrcEnum.VCC_LO))
epilogue.append(v_mov_b32_e32(v[240], 0))
for i in range(N_SGPRS):
epilogue.append(v_mov_b32_e32(v[243], s[i]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92], offset=VGPR_BYTES + i * 4))
epilogue.append(v_mov_b32_e32(v[243], s[90]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92], offset=VGPR_BYTES + SGPR_BYTES))
epilogue.append(v_mov_b32_e32(v[243], s[91]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92], offset=VGPR_BYTES + SGPR_BYTES + 4))
epilogue.append(s_mov_b32(s[SrcEnum.EXEC_LO - 128], s[94]))
epilogue.append(s_endpgm())
return prologue, epilogue
def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
"""Parse output buffer into WaveState."""
st = WaveState()
for i in range(N_VGPRS):
for lane in range(n_lanes):
off = i * WAVE_SIZE * 4 + lane * 4
st.vgpr[lane][i] = struct.unpack_from('<I', out_buf, off)[0]
for i in range(N_SGPRS):
st.sgpr[i] = struct.unpack_from('<I', out_buf, VGPR_BYTES + i * 4)[0]
st.vcc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES)[0]
st.scc = struct.unpack_from('<I', out_buf, VGPR_BYTES + SGPR_BYTES + 4)[0]
return st
def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions via emulator run_asm, dump state to memory, return WaveState."""
out_buf = (ctypes.c_uint8 * OUT_BYTES)(*([0] * OUT_BYTES))
out_addr = ctypes.addressof(out_buf)
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
args = (ctypes.c_uint64 * 1)(out_addr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
lib_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(out_addr, OUT_BYTES), (args_ptr, 8)})
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
rsrc2 = 0x19c | (128 << 15)
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
assert result == 0, f"run_asm failed with {result}"
return parse_output(bytes(out_buf), n_lanes)
def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.helpers import flat_mv
dev = Device["AMD"]
compiler = HIPCompiler(dev.arch)
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
byte_str = ', '.join(f'0x{b:02x}' for b in code)
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 96
.amdhsa_wavefront_size32 1
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_kernarg_size 8
.amdhsa_group_segment_fixed_size 65536
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 8
.group_segment_fixed_size: 65536
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.wavefront_size: 32
.sgpr_count: 96
.vgpr_count: 256
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
lib = compiler.compile(asm_src)
prg = AMDProgram(dev, "test", lib)
out_gpu = dev.allocator.alloc(OUT_BYTES)
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
out_buf = bytearray(OUT_BYTES)
dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
return parse_output(bytes(out_buf), n_lanes)
def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
"""Compare two WaveStates and return list of differences."""
import math
diffs = []
for i in range(n_vgprs):
for lane in range(n_lanes):
emu_val = emu_st.vgpr[lane][i]
hw_val = hw_st.vgpr[lane][i]
if emu_val != hw_val:
emu_f, hw_f = _f32(emu_val), _f32(hw_val)
if math.isnan(emu_f) and math.isnan(hw_f):
continue
diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
for i in range(N_SGPRS):
emu_val = emu_st.sgpr[i]
hw_val = hw_st.sgpr[i]
if emu_val != hw_val:
diffs.append(f"s[{i}]: emu=0x{emu_val:08x} hw=0x{hw_val:08x}")
if emu_st.vcc != hw_st.vcc:
diffs.append(f"vcc: emu=0x{emu_st.vcc:08x} hw=0x{hw_st.vcc:08x}")
if emu_st.scc != hw_st.scc:
diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
return diffs
def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions and return WaveState.
If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
Otherwise, runs only on emulator.
"""
emu_st = run_program_emu(instructions, n_lanes)
if USE_HW:
hw_st = run_program_hw(instructions, n_lanes)
diffs = compare_wave_states(emu_st, hw_st, n_lanes)
if diffs:
raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
return hw_st
return emu_st
+629
View File
@@ -0,0 +1,629 @@
"""Tests for DS instructions - data share (LDS) operations.
Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
ds_add_*, ds_max_*, ds_min_*, ds_and_*, ds_or_*, ds_xor_*,
ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestDS2Addr(unittest.TestCase):
"""Tests for DS_*_2ADDR instructions."""
def test_ds_store_load_2addr_b32(self):
"""DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset * 4."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
def test_ds_store_load_2addr_b64(self):
"""DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=0, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
class TestDS2AddrMore(unittest.TestCase):
"""Additional DS_*_2ADDR tests."""
def test_ds_store_load_2addr_b32_nonzero_offsets(self):
"""DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=2, offset1=5),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
def test_ds_2addr_b64_no_overlap(self):
"""DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_mov_b32(s[2], 0x33333333),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
s_mov_b32(s[2], 0x44444444),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")
def test_ds_load_2addr_b32_no_overwrite(self):
"""DS_LOAD_2ADDR_B32 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[4], s[2]), # Sentinel
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
def test_ds_load_b64_no_overwrite(self):
"""DS_LOAD_B64 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[2]),
ds_store_b64(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0x12345678),
v_mov_b32_e32(v[4], s[2]), # Sentinel
ds_load_b64(addr=v[10], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][3], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSAtomic(unittest.TestCase):
"""Tests for DS atomic operations."""
def test_ds_max_rtn_u32(self):
"""DS_MAX_RTN_U32: atomically store max and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_min_rtn_u32(self):
"""DS_MIN_RTN_U32: atomically store min and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[1], s[2]),
ds_min_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 200)
self.assertEqual(st.vgpr[0][3], 100)
def test_ds_and_rtn_b32(self):
"""DS_AND_RTN_B32: atomically AND and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xFF00FF00),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0xFFFF0000),
v_mov_b32_e32(v[1], s[2]),
ds_and_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFF00FF00)
self.assertEqual(st.vgpr[0][3], 0xFF000000)
def test_ds_or_rtn_b32(self):
"""DS_OR_RTN_B32: atomically OR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x00FF0000),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0x000000FF),
v_mov_b32_e32(v[1], s[2]),
ds_or_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x00FF0000)
self.assertEqual(st.vgpr[0][3], 0x00FF00FF)
def test_ds_xor_rtn_b32(self):
"""DS_XOR_RTN_B32: atomically XOR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0xFFFFFFFF),
v_mov_b32_e32(v[1], s[2]),
ds_xor_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0x55555555)
def test_ds_inc_rtn_u32(self):
"""DS_INC_RTN_U32: increment with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_inc_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 6)
def test_ds_dec_rtn_u32(self):
"""DS_DEC_RTN_U32: decrement with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 4)
def test_ds_cmpstore_b32_match(self):
"""DS_CMPSTORE_B32: conditional store when compare matches."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[2], s[2]), # compare = 100 (matches)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 200)
def test_ds_cmpstore_b32_no_match(self):
"""DS_CMPSTORE_B32: no store when compare doesn't match."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[2], s[2]), # compare = 50 (doesn't match)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 100)
def test_ds_max_u32_no_rtn(self):
"""DS_MAX_U32 (no RTN): atomically store max, no return value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_add_u32_no_rtn_preserves_vdst(self):
"""DS_ADD_U32 (no RTN) should NOT write to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF, "v2 should preserve sentinel")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_add_rtn_u32_writes_vdst(self):
"""DS_ADD_RTN_U32 should write old value to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_dec_rtn_u32_wrap(self):
"""DS_DEC_RTN_U32: decrement wraps when value is 0 or > limit."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0), # Start at 0
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, "v2 should have old value (0)")
# When mem == 0 or mem > limit, result = limit
self.assertEqual(st.vgpr[0][3], 10, "v3 should wrap to limit (10)")
class TestDSStorexchg(unittest.TestCase):
"""Tests for DS_STOREXCHG instructions."""
def test_ds_storexchg_rtn_b32(self):
"""DS_STOREXCHG_RTN_B32: exchange value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STOREXCHG_RTN_B32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
class TestDSRegisterWidth(unittest.TestCase):
"""Regression tests: DS loads should only write correct number of VGPRs."""
def test_ds_load_b32_no_overwrite(self):
"""DS_LOAD_B32 should only write 1 VGPR."""
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]), # sentinel
ds_store_b32(addr=v[0], data0=v[1], offset0=0),
s_waitcnt(lgkmcnt=0),
ds_load_b32(addr=v[0], vdst=v[1], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
class TestDS2AddrStride64(unittest.TestCase):
"""Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
def test_ds_store_load_2addr_stride64_b32(self):
"""DS_STORE_2ADDR_STRIDE64_B32: stores at ADDR + offset*256."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[2], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 from addr 256")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 from addr 512")
def test_ds_store_load_2addr_stride64_b64(self):
"""DS_STORE_2ADDR_STRIDE64_B64: stores at ADDR + offset*512."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[2], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
def test_ds_storexchg_2addr_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_RTN_B32: exchange at two addresses."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6], offset0=0, offset1=1),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "old val 0")
self.assertEqual(st.vgpr[0][5], 0x22222222, "old val 1")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
def test_ds_storexchg_rtn_b64(self):
"""DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]), # initial low
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]), # initial high
DS(DSOp.DS_STORE_B64, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]), # new low
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]), # new high
DS(DSOp.DS_STOREXCHG_RTN_B64, addr=v[10], data0=v[2], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should have old low dword")
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE, "v5 should have old high dword")
self.assertEqual(st.vgpr[0][6], 0x12345678, "v6 should have new low dword")
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0, "v7 should have new high dword")
def test_ds_store_load_2addr_stride64_b64_roundtrip(self):
"""DS_STORE_2ADDR_STRIDE64_B64 followed by DS_LOAD_2ADDR_STRIDE64_B64 works correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[0], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have val1 low")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have val1 high")
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have val2 low")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have val2 high")
def test_ds_storexchg_2addr_stride64_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32: exchange at two addresses (offset*256)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have old value")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have old value")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have new value")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have new value")
def test_ds_storexchg_2addr_stride64_rtn_b64_returns_old(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64: returns old values correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0], data1=v[0], vdst=v[0], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[6], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[7], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64, addr=v[10], data0=v[6], data1=v[6], vdst=v[8], offset0=1, offset1=2),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][8], 0x11111111, "v8 should have old val1 low")
self.assertEqual(st.vgpr[0][9], 0x22222222, "v9 should have old val1 high")
self.assertEqual(st.vgpr[0][10], 0x11111111, "v10 should have old val2 low")
self.assertEqual(st.vgpr[0][11], 0x22222222, "v11 should have old val2 high")
class TestAtomicOrdering(unittest.TestCase):
"""Tests for atomic operation return values and ordering."""
def test_ds_add_rtn_sequence(self):
"""DS_ADD_RTN returns correct old values in sequence."""
instructions = [
v_mov_b32_e32(v[10], 0),
v_mov_b32_e32(v[0], 100),
DS(DSOp.DS_STORE_B32, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[1], 25),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[3], offset0=0),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B32, addr=v[10], vdst=v[4], offset0=0),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "First add should return 100")
self.assertEqual(st.vgpr[0][3], 125, "Second add should return 125")
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
if __name__ == '__main__':
unittest.main()
+363
View File
@@ -0,0 +1,363 @@
"""Tests for FLAT instructions - flat memory operations.
Includes: flat_load_*, flat_store_*, flat_atomic_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestFlatAtomic(unittest.TestCase):
"""Tests for FLAT atomic instructions."""
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
"""Helper to create atomic test instructions."""
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
check_fn(st)
def test_flat_atomic_add_u32(self):
"""FLAT_ATOMIC_ADD_U32 adds to memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_swap_b32(self):
"""FLAT_ATOMIC_SWAP_B32 swaps memory value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_and_b32(self):
"""FLAT_ATOMIC_AND_B32 ANDs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xFF00FF00),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xFFFF0000),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_AND_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xFF00FF00)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_or_b32(self):
"""FLAT_ATOMIC_OR_B32 ORs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0x00FF0000),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0x0000FF00),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_OR_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0x00FF0000)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_inc_u32(self):
"""FLAT_ATOMIC_INC_U32 increments and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 10),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 100), # threshold
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 10)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_dec_u32(self):
"""FLAT_ATOMIC_DEC_U32 decrements and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 10),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_DEC_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 10)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_sub_u32(self):
"""FLAT_ATOMIC_SUB_U32 subtracts from memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 30),
v_mov_b32_e32(v[3], s[0]), # sub 30
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SUB_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_xor_b32(self):
"""FLAT_ATOMIC_XOR_B32 XORs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[3], s[0]), # XOR mask
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_XOR_B32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 should have old value")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_min_u32(self):
"""FLAT_ATOMIC_MIN_U32 stores min and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]), # compare value (smaller)
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_MIN_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_max_u32(self):
"""FLAT_ATOMIC_MAX_U32 stores max and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[3], s[0]), # compare value (larger)
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_MAX_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 50, "v4 should have old value (50)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_inc_u64_returns_old_value(self):
"""FLAT_ATOMIC_INC_U64 should return full 64-bit old value."""
TEST_OFFSET = 2000
setup = [
# Store initial 64-bit value: 0xCAFEBABE_DEADBEEF
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Threshold: 0xFFFFFFFF_FFFFFFFF
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[4], s[0]),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF, "v6 should have old value low dword")
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE, "v7 should have old value high dword")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_add_u64(self):
"""FLAT_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0x00000001), # add 1
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0x11111111, "v6 should have old value low")
self.assertEqual(st.vgpr[0][7], 0x22222222, "v7 should have old value high")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_swap_b64(self):
"""FLAT_ATOMIC_SWAP_B64 swaps 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0xDDDDDDDD),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have old value low")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have old value high")
self._make_test(setup, atomic, check, TEST_OFFSET)
class TestFlatLoad(unittest.TestCase):
"""Tests for FLAT load instructions."""
def test_flat_load_b32(self):
"""FLAT_LOAD_B32 loads 32-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B32, addr=v[0], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
def test_flat_load_b64(self):
"""FLAT_LOAD_B64 loads 64-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B64, addr=v[0], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
def test_flat_load_b96(self):
"""FLAT_LOAD_B96 loads 96-bit (3 dword) value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[4], s[0]),
global_store_b96(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B96, addr=v[0], vdst=v[5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][5], 0x11111111)
self.assertEqual(st.vgpr[0][6], 0x22222222)
self.assertEqual(st.vgpr[0][7], 0x33333333)
def test_flat_load_b128(self):
"""FLAT_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x44444444),
v_mov_b32_e32(v[5], s[0]),
global_store_b128(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(FLATOp.FLAT_LOAD_B128, addr=v[0], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][6], 0x11111111)
self.assertEqual(st.vgpr[0][7], 0x22222222)
self.assertEqual(st.vgpr[0][8], 0x33333333)
self.assertEqual(st.vgpr[0][9], 0x44444444)
if __name__ == '__main__':
unittest.main()
+364
View File
@@ -0,0 +1,364 @@
"""Tests for GLOBAL instructions - global memory operations.
Includes: global_load_*, global_store_*, global_atomic_*, global_load_d16_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestGlobalAtomic(unittest.TestCase):
"""Tests for GLOBAL atomic instructions."""
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
"""Helper to create atomic test instructions."""
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
] + setup_instrs + [atomic_instr, s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
check_fn(st)
def test_global_atomic_add_u32(self):
"""GLOBAL_ATOMIC_ADD_U32 adds to memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(GLOBALOp.GLOBAL_ATOMIC_ADD_U32, addr=v[0], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1, seg=2)
def check(st):
self.assertEqual(st.vgpr[0][4], 100)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_global_atomic_add_u64(self):
"""GLOBAL_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[0], 0x00000001),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(GLOBALOp.GLOBAL_ATOMIC_ADD_U64, addr=v[0], data=v[4], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1, seg=2)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][7], 0x00000000)
self._make_test(setup, atomic, check, TEST_OFFSET)
class TestGlobalLoad(unittest.TestCase):
"""Tests for GLOBAL load instructions."""
def test_global_load_b96(self):
"""GLOBAL_LOAD_B96 loads 96-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[0]),
global_store_b96(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B96, addr=v[0], vdst=v[5], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][5], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][6], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][7], 0xCCCCCCCC)
def test_global_load_b128(self):
"""GLOBAL_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[5], s[0]),
global_store_b128(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B128, addr=v[0], vdst=v[6], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][8], 0x12345678)
self.assertEqual(st.vgpr[0][9], 0x9ABCDEF0)
class TestGlobalStore(unittest.TestCase):
"""Tests for GLOBAL store instructions."""
def test_global_store_b64_basic(self):
"""GLOBAL_STORE_B64 stores 8 bytes from v[n:n+1] to memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
s_mov_b32(s[5], 0xCAFEBABE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b64(addr=v[0], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B64, addr=v[0], vdst=v[4], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
class TestD16HiLoads(unittest.TestCase):
"""Tests for D16_HI load instructions that load into high 16 bits."""
def test_global_load_d16_hi_b16_preserves_low_bits(self):
"""GLOBAL_LOAD_D16_HI_B16 must preserve low 16 bits of destination."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[3], s[4]),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[0], vdst=v[3], data=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFEBEEF, f"Expected 0xCAFEBEEF, got 0x{result:08x}")
def test_global_load_d16_hi_b16_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_B16 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b16(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[0], s[4]), # data field - should NOT affect result
v_mov_b32_e32(v[1], 0), # vdst - low bits should be preserved
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
def test_global_load_d16_hi_u8_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_U8 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xAB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b8(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[4], s[4]), # data field
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[5], s[4]), # vdst
v_mov_b32_e32(v[3], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_U8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0x00ABBEEF, f"Expected 0x00ABBEEF, got 0x{result:08x}")
def test_global_load_d16_hi_b16_same_addr_and_dst_zero_addr(self):
"""GLOBAL_LOAD_D16_HI_B16 with same register for addr and vdst, addr value=0."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b16(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[1], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
def test_global_load_d16_hi_b16_tril_exact_pattern(self):
"""Exact pattern from tril() failure: data=v0 differs from vdst=v1."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0x01010101),
v_mov_b32_e32(v[10], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b32(addr=v[3], data=v[10], saddr=s[2], offset=TEST_OFFSET),
global_store_b32(addr=v[3], data=v[10], saddr=s[2], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
# Set v[0] to 0x0101 (simulating prior u16 load result)
s_mov_b32(s[4], 0x0101),
v_mov_b32_e32(v[0], s[4]),
# Set v[1] to 0
v_mov_b32_e32(v[1], 0),
# Load using v[1] as addr AND vdst, but v[0] as data
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2], offset=TEST_OFFSET+6, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
# Expected: hi=0x0101 (loaded), lo=0x0000 (from v1) -> 0x01010000
self.assertEqual(result, 0x01010000, f"Expected 0x01010000, got 0x{result:08x}")
def test_global_load_d16_hi_i8_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_I8 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0x80), # negative signed byte = -128
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b8(addr=v[3], data=v[2], saddr=s[2], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[4], s[4]), # data field
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[5], s[4]), # vdst
v_mov_b32_e32(v[3], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_I8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2], offset=TEST_OFFSET, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
# 0x80 sign-extended = 0xFF80, lo=0xBEEF -> 0xFF80BEEF
self.assertEqual(result, 0xFF80BEEF, f"Expected 0xFF80BEEF, got 0x{result:08x}")
def test_global_store_b64_tril_pattern(self):
"""Test the exact pattern from tril() kernel that was failing."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0x01010101),
v_mov_b32_e32(v[10], s[4]),
v_mov_b32_e32(v[11], s[4]),
s_mov_b32(s[4], 0x01),
v_mov_b32_e32(v[12], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b64(addr=v[0], data=v[10], saddr=s[2], offset=TEST_OFFSET),
global_store_b8(addr=v[0], data=v[12], saddr=s[2], offset=TEST_OFFSET+8),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[2], 0),
v_mov_b32_e32(v[1], 0),
FLAT(GLOBALOp.GLOBAL_LOAD_U16, addr=v[2], vdst=v[0], data=v[0], saddr=s[2], offset=TEST_OFFSET+3, seg=2),
FLAT(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2], offset=TEST_OFFSET+6, seg=2),
FLAT(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[3], data=v[3], saddr=s[2], offset=TEST_OFFSET, seg=2),
FLAT(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[4], data=v[4], saddr=s[2], offset=TEST_OFFSET+8, seg=2),
s_waitcnt(vmcnt=0),
v_and_b32_e32(v[5], 0xffff, v[0]),
v_lshlrev_b32_e32(v[0], 24, v[0]),
v_lshrrev_b32_e32(v[5], 8, v[5]),
v_or_b32_e32(v[0], v[3], v[0]),
v_or_b32_e32(v[1], v[5], v[1]),
global_store_b64(addr=v[2], data=v[0], saddr=s[2], offset=TEST_OFFSET+16),
s_waitcnt(vmcnt=0),
FLAT(GLOBALOp.GLOBAL_LOAD_B64, addr=v[2], vdst=v[6], data=v[6], saddr=s[2], offset=TEST_OFFSET+16, seg=2),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
v0 = st.vgpr[0][0]
v1 = st.vgpr[0][1]
self.assertEqual(v0, 0x01000001, f"v0: expected 0x01000001, got 0x{v0:08x}")
self.assertEqual(v1, 0x01010001, f"v1: expected 0x01010001, got 0x{v1:08x}")
byte5 = (v1 >> 8) & 0xff
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
if __name__ == '__main__':
unittest.main()
+342
View File
@@ -0,0 +1,342 @@
"""Tests for SOP instructions - scalar operations.
Includes: s_add_u32, s_mov_b32, s_and_b32, s_or_b32, s_quadmask_b32, s_wqm_b32,
s_cbranch_vccnz, s_cbranch_vccz
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestBasicScalar(unittest.TestCase):
"""Tests for basic scalar operations."""
def test_s_add_u32(self):
"""S_ADD_U32 adds two scalar values."""
instructions = [
s_mov_b32(s[0], 100),
s_mov_b32(s[1], 200),
s_add_u32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 300)
def test_s_add_u32_carry(self):
"""S_ADD_U32 sets SCC on overflow."""
instructions = [
s_mov_b32(s[0], 64),
s_not_b32(s[0], s[0]), # ~64 = 0xffffffbf
s_mov_b32(s[1], 64),
s_add_u32(s[2], s[0], s[1]), # 0xffffffbf + 64 = 0xffffffff
s_mov_b32(s[3], 1),
s_add_u32(s[4], s[2], s[3]), # 0xffffffff + 1 = overflow
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0)
self.assertEqual(st.scc, 1)
def test_s_brev_b32(self):
"""S_BREV_B32 reverses bits of a 32-bit value."""
# 10 = 0b00000000000000000000000000001010
# reversed = 0b01010000000000000000000000000000 = 0x50000000
instructions = [
s_mov_b32(s[0], 10),
s_brev_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x50000000)
def test_s_brev_b32_all_ones(self):
"""S_BREV_B32 with all ones stays all ones."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_brev_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
def test_s_brev_b32_single_bit(self):
"""S_BREV_B32 with bit 0 set becomes bit 31."""
instructions = [
s_mov_b32(s[0], 1),
s_brev_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x80000000)
class TestQuadmaskWqm(unittest.TestCase):
"""Tests for S_QUADMASK_B32 and S_WQM_B32."""
def test_s_quadmask_b32_all_quads_active(self):
"""S_QUADMASK_B32 with all quads active."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF), # All lanes active
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# Each quad (4 lanes) with any bit set -> 1 bit in result
# 32 lanes = 8 quads, all active -> 0xFF
self.assertEqual(st.sgpr[1], 0xFF)
def test_s_quadmask_b32_alternating_quads(self):
"""S_QUADMASK_B32 with alternating quads active."""
instructions = [
s_mov_b32(s[0], 0x0F0F0F0F), # Quads 0,2,4,6 active
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# Quads 0,2,4,6 have at least one bit -> 0b01010101 = 0x55
self.assertEqual(st.sgpr[1], 0x55)
def test_s_quadmask_b32_no_quads_active(self):
"""S_QUADMASK_B32 with no quads active."""
instructions = [
s_mov_b32(s[0], 0),
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0)
def test_s_quadmask_b32_single_lane_per_quad(self):
"""S_QUADMASK_B32 with single lane active in each quad."""
instructions = [
s_mov_b32(s[0], 0x11111111), # Bit 0 of each nibble
s_quadmask_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# All 8 quads have at least one lane -> 0xFF
self.assertEqual(st.sgpr[1], 0xFF)
def test_s_wqm_b32_all_active(self):
"""S_WQM_B32 with all lanes active returns all 1s."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_wqm_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
def test_s_wqm_b32_alternating_quads(self):
"""S_WQM_B32 with single lane per quad expands to full quads."""
instructions = [
s_mov_b32(s[0], 0x11111111), # One lane per quad
s_wqm_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
# Each quad with any bit expands to all 4 bits
self.assertEqual(st.sgpr[1], 0xFFFFFFFF)
def test_s_wqm_b32_zero(self):
"""S_WQM_B32 with zero input returns zero."""
instructions = [
s_mov_b32(s[0], 0),
s_wqm_b32(s[1], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0)
class TestBranch(unittest.TestCase):
"""Tests for branch instructions."""
def test_cbranch_vccnz_ignores_vcc_hi(self):
"""S_CBRANCH_VCCNZ should only check VCC_LO in wave32."""
instructions = [
# Set VCC_LO = 0, VCC_HI = 1
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0),
s_mov_b32(s[SrcEnum.VCC_HI - 128], 1),
v_mov_b32_e32(v[0], 0),
# If VCC_HI is incorrectly used, branch will be taken
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
v_mov_b32_e32(v[0], 42), # This should execute
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 0)")
def test_cbranch_vccz_ignores_vcc_hi(self):
"""S_CBRANCH_VCCZ should only check VCC_LO in wave32."""
instructions = [
# Set VCC_LO = 1, VCC_HI = 0
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),
s_mov_b32(s[SrcEnum.VCC_HI - 128], 0),
v_mov_b32_e32(v[0], 0),
# If VCC_HI is incorrectly used, branch will be taken
s_cbranch_vccz(1), # Skip next instruction if VCC == 0
v_mov_b32_e32(v[0], 42), # This should execute
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 1)")
def test_cbranch_vccnz_branches_on_vcc_lo(self):
"""S_CBRANCH_VCCNZ branches when VCC_LO is non-zero."""
instructions = [
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1),
v_mov_b32_e32(v[0], 0),
s_cbranch_vccnz(1), # Skip next instruction if VCC != 0
v_mov_b32_e32(v[0], 42), # This should be skipped
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0, "Branch should be taken (VCC_LO is 1)")
class Test64BitLiterals(unittest.TestCase):
"""Tests for 64-bit literal encoding in instructions."""
def test_64bit_literal_negative_encoding(self):
"""64-bit literal -2^32 encodes correctly."""
lit = -4294967296.0 # -2^32
lit_bits = f2i64(lit)
instructions = [
s_mov_b32(s[0], lit_bits & 0xffffffff),
s_mov_b32(s[1], lit_bits >> 32),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
result = i642f(st.vgpr[0][0] | (st.vgpr[0][1] << 32))
self.assertAlmostEqual(result, -4294967296.0, places=5)
def test_64bit_literal_positive_encoding(self):
"""64-bit instruction encodes large positive literals correctly."""
large_val = 0x12345678
inst = v_add_f64(v[2], v[0], large_val)
self.assertIsNotNone(inst._literal, "Literal should be set")
actual_lit = (inst._literal >> 32) & 0xffffffff
self.assertEqual(actual_lit, large_val, f"Literal should be {large_val:#x}, got {actual_lit:#x}")
class TestSCCBehavior(unittest.TestCase):
"""Tests for SCC condition code behavior."""
def test_scc_from_s_cmp(self):
"""SCC should be set by scalar compare."""
instructions = [
s_mov_b32(s[0], 10),
s_cmp_eq_u32(s[0], 10),
s_cselect_b32(s[1], 1, 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 1, "SCC should be true")
self.assertEqual(st.scc, 1)
def test_scc_clear(self):
"""SCC should be cleared by failing compare."""
instructions = [
s_mov_b32(s[0], 10),
s_cmp_eq_u32(s[0], 20),
s_cselect_b32(s[1], 1, 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0, "SCC should be false")
self.assertEqual(st.scc, 0)
class TestSignedArithmetic(unittest.TestCase):
"""Tests for S_ADD_I32, S_SUB_I32 and their SCC overflow behavior."""
def test_s_add_i32_no_overflow(self):
"""S_ADD_I32: 1 + 1 = 2, no overflow, SCC=0."""
instructions = [
s_mov_b32(s[0], 1),
s_add_i32(s[1], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 2)
self.assertEqual(st.scc, 0, "No overflow, SCC should be 0")
def test_s_add_i32_positive_overflow(self):
"""S_ADD_I32: MAX_INT + 1 overflows, SCC=1."""
instructions = [
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
s_add_i32(s[1], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x80000000) # Wraps to MIN_INT
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
def test_s_add_i32_negative_no_overflow(self):
"""S_ADD_I32: -10 + 20 = 10, no overflow."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFF6), # -10 in two's complement
s_mov_b32(s[1], 20),
s_add_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 10)
self.assertEqual(st.scc, 0)
def test_s_add_i32_negative_overflow(self):
"""S_ADD_I32: MIN_INT + (-1) underflows, SCC=1."""
instructions = [
s_mov_b32(s[0], 0x80000000), # MIN_INT
s_mov_b32(s[1], 0xFFFFFFFF), # -1
s_add_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0x7FFFFFFF) # Wraps to MAX_INT
self.assertEqual(st.scc, 1, "Underflow, SCC should be 1")
def test_s_sub_i32_no_overflow(self):
"""S_SUB_I32: 10 - 5 = 5, no overflow."""
instructions = [
s_mov_b32(s[0], 10),
s_mov_b32(s[1], 5),
s_sub_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 5)
self.assertEqual(st.scc, 0)
def test_s_sub_i32_overflow(self):
"""S_SUB_I32: MAX_INT - (-1) overflows, SCC=1."""
instructions = [
s_mov_b32(s[0], 0x7FFFFFFF), # MAX_INT
s_mov_b32(s[1], 0xFFFFFFFF), # -1
s_sub_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0x80000000) # Wraps to MIN_INT
self.assertEqual(st.scc, 1, "Overflow, SCC should be 1")
def test_s_mul_hi_u32(self):
"""S_MUL_HI_U32: high 32 bits of u32 * u32."""
instructions = [
s_mov_b32(s[0], 0x80000000), # 2^31
s_mov_b32(s[1], 4),
s_mul_hi_u32(s[2], s[0], s[1]), # (2^31 * 4) >> 32 = 2
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 2)
def test_s_mul_i32(self):
"""S_MUL_I32: signed multiply low 32 bits."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF), # -1
s_mov_b32(s[1], 10),
s_mul_i32(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xFFFFFFF6) # -10
def test_division_sequence_from_llvm(self):
"""Test the division sequence pattern from LLVM-generated code."""
# This sequence is from the sin kernel and computes integer division
# s10 = dividend, s18 = divisor, result in s6/s14
dividend = 0x28BE60DB # Some value from the sin kernel
divisor = 3 # Simplified divisor
instructions = [
s_mov_b32(s[10], dividend),
s_mov_b32(s[18], divisor),
# Compute reciprocal approximation: s6 = ~0 / divisor (approx)
s_mov_b32(s[11], 0),
s_sub_i32(s[11], s[11], s[18]), # s11 = -divisor
# For testing, just verify basic arithmetic works
s_mul_i32(s[6], s[10], 2),
s_add_i32(s[7], s[6], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], (dividend * 2) & 0xFFFFFFFF)
self.assertEqual(st.sgpr[7], ((dividend * 2) + 1) & 0xFFFFFFFF)
if __name__ == '__main__':
unittest.main()
File diff suppressed because it is too large Load Diff
+451
View File
@@ -0,0 +1,451 @@
"""Tests for VOP2 instructions - two operand vector operations.
Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
v_lshrrev_b32, v_lshlrev_b32, v_fmac_f32, v_fmaak_f32, v_fmamk_f32,
v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestBasicArithmetic(unittest.TestCase):
"""Tests for basic arithmetic VOP2 instructions."""
def test_v_add_f32(self):
"""V_ADD_F32 adds two floats."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], 2.0),
v_add_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
def test_v_mul_f32(self):
"""V_MUL_F32 multiplies two floats."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_mul_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 8.0, places=5)
def test_v_fmac_f32(self):
"""V_FMAC_F32: d = d + a*b using inline constants."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_mov_b32_e32(v[2], 1.0),
v_fmac_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmaak_f32(self):
"""V_FMAAK_F32: d = a * b + K using inline constants."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_fmaak_f32_e32(v[2], v[0], v[1], 0x3f800000),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmamk_f32_basic(self):
"""V_FMAMK_F32: d = a * K + b."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 1.0),
v_fmamk_f32_e32(v[2], v[0], 0x40800000, v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmamk_f32_small_constant(self):
"""V_FMAMK_F32 with small constant."""
instructions = [
v_mov_b32_e32(v[0], 4.0),
v_mov_b32_e32(v[1], 1.0),
v_fmamk_f32_e32(v[2], v[0], f2i(0.5), v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
class TestBitManipulation(unittest.TestCase):
"""Tests for bit manipulation VOP2 instructions."""
def test_v_and_b32(self):
"""V_AND_B32 bitwise and."""
instructions = [
s_mov_b32(s[0], 0xff),
s_mov_b32(s[1], 0x0f),
v_mov_b32_e32(v[0], s[0]),
v_and_b32_e32(v[1], s[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x0f)
def test_v_and_b32_quadrant(self):
"""V_AND_B32 for quadrant extraction (n & 3)."""
instructions = [
s_mov_b32(s[0], 15915),
v_mov_b32_e32(v[0], s[0]),
v_and_b32_e32(v[1], 3, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 15915 & 3)
def test_v_lshrrev_b32(self):
"""V_LSHRREV_B32 logical shift right."""
instructions = [
s_mov_b32(s[0], 0xff00),
v_mov_b32_e32(v[0], s[0]),
v_lshrrev_b32_e32(v[1], 8, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xff)
def test_v_lshlrev_b32(self):
"""V_LSHLREV_B32 logical shift left."""
instructions = [
s_mov_b32(s[0], 0xff),
v_mov_b32_e32(v[0], s[0]),
v_lshlrev_b32_e32(v[1], 8, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xff00)
def test_v_xor_b32(self):
"""V_XOR_B32 bitwise xor (used in sin for sign)."""
instructions = [
s_mov_b32(s[0], 0x80000000),
s_mov_b32(s[1], f2i(1.0)),
v_mov_b32_e32(v[0], s[1]),
v_xor_b32_e32(v[1], s[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), -1.0, places=5)
def test_v_xor_b32_sign_flip(self):
"""V_XOR_B32 for sign flip pattern."""
instructions = [
s_mov_b32(s[0], 0x80000000),
v_mov_b32_e32(v[0], -2.0),
v_xor_b32_e32(v[1], s[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.0, places=5)
class TestSpecialValues(unittest.TestCase):
"""Tests for special float values - inf, nan, zero handling."""
def test_v_mul_f32_zero_times_inf(self):
"""V_MUL_F32: 0 * inf = NaN."""
import math
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0x7f800000),
v_mov_b32_e32(v[1], s[0]),
v_mul_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
def test_v_add_f32_inf_minus_inf(self):
"""V_ADD_F32: inf + (-inf) = NaN."""
import math
instructions = [
s_mov_b32(s[0], 0x7f800000),
s_mov_b32(s[1], 0xff800000),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
class TestF16Ops(unittest.TestCase):
"""Tests for 16-bit VOP2 operations."""
def test_v_add_f16_basic(self):
"""V_ADD_F16 adds two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4200, f"Expected 0x4200 (f16 3.0), got 0x{result:04x}")
def test_v_add_f16_negative(self):
"""V_ADD_F16 with negative values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0xc000), # f16 -2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0xbc00, f"Expected 0xbc00 (f16 -1.0), got 0x{result:04x}")
def test_v_mul_f16_basic(self):
"""V_MUL_F16 multiplies two f16 values."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mul_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4600, f"Expected 0x4600 (f16 6.0), got 0x{result:04x}")
def test_v_mul_f16_by_zero(self):
"""V_MUL_F16 by zero."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0),
v_mul_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x0000, f"Expected 0x0000 (f16 0.0), got 0x{result:04x}")
def test_v_fmac_f16_basic(self):
"""V_FMAC_F16: d = d + a*b."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
s_mov_b32(s[2], 0x3c00), # f16 1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
v_fmac_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
def test_v_fmaak_f16_basic(self):
"""V_FMAAK_F16: d = a * b + K."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_fmaak_f16_e32(v[2], v[0], v[1], 0x3c00), # + f16 1.0
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
class TestHiHalfOps(unittest.TestCase):
"""Tests for VOP2 16-bit operations with hi-half operands."""
def test_v_add_f16_src0_hi_fold(self):
"""V_ADD_F16 with src0 hi-half fold (same register, different halves)."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # lo=f16(1.0), hi=f16(2.0)
v_mov_b32_e32(v[0], s[0]),
VOP3(VOP3Op.V_ADD_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_add_f16_src0_hi_different_reg(self):
"""V_ADD_F16 with src0 hi-half from different register."""
instructions = [
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
s_mov_b32(s[1], 0x00003c00), # hi=0, lo=f16(1.0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
VOP3(VOP3Op.V_ADD_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_mul_f16_src0_hi(self):
"""V_MUL_F16 with src0 from high half."""
instructions = [
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
s_mov_b32(s[1], 0x00004200), # hi=0, lo=f16(3.0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
VOP3(VOP3Op.V_MUL_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
def test_v_mul_f16_hi_half(self):
"""V_MUL_F16 reading from high half."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # lo=1.0, hi=2.0
v_mov_b32_e32(v[0], s[0]),
VOP3(VOP3Op.V_MUL_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0011),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x4400, f"Expected f16(4.0)=0x4400, got 0x{result:04x}")
def test_v_fma_f16_hi_dest(self):
"""V_FMA_F16 writing to high half with opsel.
Uses V_FMA_F16 (not V_FMAC_F16) because it has explicit src2 operand
which makes opsel handling clearer.
"""
instructions = [
s_mov_b32(s[0], 0x3c000000), # hi=f16(1.0), lo=0
s_mov_b32(s[1], 0x4000), # f16(2.0) in lo
s_mov_b32(s[2], 0x4200), # f16(3.0) in lo
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
# V_FMA_F16: dst = src0 * src1 + src2
# opsel=0b1100: bit2=src2 hi, bit3=dst hi
# So: v[0].hi = v[1].lo * v[2].lo + v[0].hi = 2.0 * 3.0 + 1.0 = 7.0
VOP3(VOP3Op.V_FMA_F16, vdst=v[0], src0=v[1], src1=v[2], src2=v[0], opsel=0b1100),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][0] >> 16) & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(hi, 0x4700, f"Expected f16(7.0)=0x4700 in hi, got 0x{hi:04x}")
def test_v_add_f16_multilane(self):
"""V_ADD_F16 with multiple lanes."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
for lane in range(4):
result = st.vgpr[lane][2] & 0xffff
self.assertEqual(result, 0x4200, f"Lane {lane}: expected 0x4200, got 0x{result:04x}")
class TestCndmask(unittest.TestCase):
"""Tests for V_CNDMASK_B32 and V_CNDMASK_B16."""
def test_v_cndmask_b16_select_src0(self):
"""V_CNDMASK_B16 selects src0 when VCC bit is 0."""
instructions = [
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0), # VCC = 0
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cndmask_b16(v[2], v[0], v[1], VCC),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected src0=0x3c00, got 0x{result:04x}")
def test_v_cndmask_b16_select_src1(self):
"""V_CNDMASK_B16 selects src1 when VCC bit is 1."""
instructions = [
s_mov_b32(s[SrcEnum.VCC_LO - 128], 1), # VCC = 1
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cndmask_b16(v[2], v[0], v[1], VCC),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected src1=0x4000, got 0x{result:04x}")
def test_v_cndmask_b16_write_hi(self):
"""V_CNDMASK_B16 can write to high 16 bits with opsel."""
instructions = [
s_mov_b32(s[0], 0x3c003800), # src0: hi=1.0, lo=0.5
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x4000c000), # src1: hi=2.0, lo=-2.0
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], 0xDEAD0000), # v2 initial: hi=0xDEAD, lo=0
v_mov_b32_e32(v[2], s[2]),
s_mov_b32(s[SrcEnum.VCC_LO - 128], 0), # vcc = 0, select src0
# opsel=0b1011: bit0=src0 hi, bit1=src1 hi, bit3=dst hi
VOP3(VOP3Op.V_CNDMASK_B16, vdst=v[2], src0=v[0], src1=v[1], src2=SrcEnum.VCC_LO, opsel=0b1011),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][2] >> 16) & 0xffff
lo = st.vgpr[0][2] & 0xffff
# vcc=0 selects src0.h = 1.0 = 0x3c00, writes to hi
self.assertEqual(hi, 0x3c00, f"Expected hi=0x3c00 (1.0), got 0x{hi:04x}")
self.assertEqual(lo, 0x0000, f"Expected lo preserved as 0, got 0x{lo:04x}")
class TestSpecialFloatValues(unittest.TestCase):
"""Tests for special float value handling in VOP2 instructions."""
def test_neg_zero_add(self):
"""-0.0 + 0.0 = +0.0 (IEEE 754)."""
neg_zero = 0x80000000
instructions = [
s_mov_b32(s[0], neg_zero),
v_mov_b32_e32(v[0], s[0]),
v_add_f32_e32(v[1], 0.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
def test_neg_zero_mul(self):
"""-0.0 * -1.0 = +0.0."""
neg_zero = 0x80000000
instructions = [
s_mov_b32(s[0], neg_zero),
v_mov_b32_e32(v[0], s[0]),
v_mul_f32_e32(v[1], -1.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
def test_inf_minus_inf(self):
"""+inf - inf = NaN."""
import math
pos_inf = 0x7f800000
neg_inf = 0xff800000
instructions = [
s_mov_b32(s[0], pos_inf),
s_mov_b32(s[1], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_sub_f32_e32(v[2], v[0], v[1]), # inf - (-inf) = inf
v_add_f32_e32(v[3], v[0], v[1]), # inf + (-inf) = NaN
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], pos_inf, "inf - (-inf) = inf")
self.assertTrue(math.isnan(i2f(st.vgpr[0][3])), "inf + (-inf) = NaN")
def test_denormal_f32_mul_ftz(self):
"""Denormal * normal - RDNA3 flushes denormals to zero (FTZ mode)."""
smallest_denorm = 0x00000001 # Smallest positive denormal
instructions = [
s_mov_b32(s[0], smallest_denorm),
v_mov_b32_e32(v[0], s[0]),
v_mul_f32_e32(v[1], 2.0, v[0]), # Denormal input gets flushed to 0
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000)
if __name__ == '__main__':
unittest.main()
File diff suppressed because it is too large Load Diff
+538
View File
@@ -0,0 +1,538 @@
"""Tests for VOP3P instructions - packed 16-bit vector operations.
Includes: v_pk_add_f16, v_pk_mul_f16, v_pk_fma_f16, v_pack_b32_f16, v_wmma_*, v_dot2_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestPackInstructions(unittest.TestCase):
"""Tests for pack instructions."""
def test_v_pack_b32_f16(self):
"""V_PACK_B32_F16 packs two f16 values into one 32-bit register."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_opsel_hi_hi(self):
"""V_PACK_B32_F16 with opsel to read high halves."""
inst = v_pack_b32_f16(v[2], v[0], v[1])
inst._values['opsel'] = 0b0011
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
inst,
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x44004000, f"Expected 0x44004000, got 0x{result:08x}")
class TestPackMore(unittest.TestCase):
"""Additional pack instruction tests."""
def test_v_pack_b32_f16_basic(self):
"""V_PACK_B32_F16 packs two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x40003c00, f"Expected 0x40003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_with_cvt(self):
"""V_PACK_B32_F16 after V_CVT_F16_F32 conversions."""
instructions = [
s_mov_b32(s[0], 0x3f800000),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cvt_f16_f32_e32(v[2], v[0]),
v_cvt_f16_f32_e32(v[3], v[1]),
v_pack_b32_f16(v[4], v[2], v[3]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][4]
self.assertEqual(result, 0x3c003c00, f"Expected 0x3c003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_packed_sources(self):
"""V_PACK_B32_F16 with packed f16 sources (reads lo halves)."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
# Expected: hi=v1.lo=0x4200 (3.0), lo=v0.lo=0x3c00 (1.0) -> 0x42003c00
self.assertEqual(result, 0x42003c00, f"Expected 0x42003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_opsel_lo_hi(self):
"""V_PACK_B32_F16 with opsel=0b0010 to read lo from src0, hi from src1."""
inst = v_pack_b32_f16(v[2], v[0], v[1])
inst._values['opsel'] = 0b0010
instructions = [
s_mov_b32(s[0], 0x40003c00),
s_mov_b32(s[1], 0x44004200),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
inst,
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x44003c00, f"Expected 0x44003c00, got 0x{result:08x}")
def test_v_pack_b32_f16_opsel_hi_lo(self):
"""V_PACK_B32_F16 with opsel=0b0001 to read hi from src0, lo from src1."""
inst = v_pack_b32_f16(v[2], v[0], v[1])
inst._values['opsel'] = 0b0001
instructions = [
s_mov_b32(s[0], 0x40003c00),
s_mov_b32(s[1], 0x44004200),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
inst,
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x42004000, f"Expected 0x42004000, got 0x{result:08x}")
def test_v_pack_b32_f16_zeros(self):
"""V_PACK_B32_F16 with zero values."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0)
def test_v_pack_b32_f16_both_positive(self):
"""V_PACK_B32_F16 with positive f16 values."""
instructions = [
s_mov_b32(s[0], 0x4200), # f16 3.0
s_mov_b32(s[1], 0x4400), # f16 4.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pack_b32_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x44004200, f"Expected 0x44004200, got 0x{result:08x}")
class TestFmaMix(unittest.TestCase):
"""Tests for V_FMA_MIX_F32 and V_FMA_MIXLO_F16."""
def test_v_fma_mix_f32_all_f32_sources(self):
"""V_FMA_MIX_F32 with all f32 sources."""
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(1.0)),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 7.0, places=5)
def test_v_fma_mix_f32_src2_f16_lo(self):
"""V_FMA_MIX_F32 with src2 as f16 from lo bits."""
from extra.assembly.amd.pcode import f32_to_f16
f16_2 = f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f16_2),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 5.0, places=5)
def test_v_fma_mix_f32_src2_f16_hi(self):
"""V_FMA_MIX_F32 with src2 as f16 from hi bits."""
from extra.assembly.amd.pcode import f32_to_f16
f16_2 = f32_to_f16(2.0)
val = (f16_2 << 16) | 0
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], val),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=4, opsel_hi=0, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 5.0, places=5)
def test_v_fma_mix_f32_with_abs(self):
"""V_FMA_MIX_F32 with abs modifier on src2."""
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(-1.0)),
v_mov_b32_e32(v[2], s[2]),
VOP3P(VOP3POp.V_FMA_MIX_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0, neg_hi=4),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 7.0, places=5)
def test_v_fma_mixlo_f16(self):
"""V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(3.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(1.0)),
v_mov_b32_e32(v[2], s[2]),
s_mov_b32(s[3], 0xdead0000),
v_mov_b32_e32(v[3], s[3]),
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
hi = (st.vgpr[0][3] >> 16) & 0xffff
self.assertAlmostEqual(lo, 7.0, places=1)
self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
def test_v_fma_mixlo_f16_all_f32_sources(self):
"""V_FMA_MIXLO_F16 with all f32 sources."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], f2i(2.0)),
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], f2i(3.0)),
v_mov_b32_e32(v[2], s[2]),
v_mov_b32_e32(v[3], 0),
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
# 1*2+3 = 5
self.assertAlmostEqual(lo, 5.0, places=1)
def test_v_fma_mixlo_f16_sin_case(self):
"""V_FMA_MIXLO_F16 case from sin kernel."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x3f800000), # f32 1.0
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[1], 0xaf05a309), # f32 tiny negative
s_mov_b32(s[6], s[1]),
s_mov_b32(s[2], 0xc0490fdb), # f32 -π
v_mov_b32_e32(v[5], s[2]),
s_mov_b32(s[3], 0x3f800000),
v_mov_b32_e32(v[3], s[3]),
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(lo, -3.14159, delta=0.01)
class TestVOP3P(unittest.TestCase):
"""Tests for VOP3P packed 16-bit operations."""
def test_v_pk_add_f16_basic(self):
"""V_PK_ADD_F16 adds two packed f16 values."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 4.0, places=2)
self.assertAlmostEqual(hi, 6.0, places=2)
def test_v_pk_mul_f16_basic(self):
"""V_PK_MUL_F16 multiplies two packed f16 values."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x42004000), # hi=3.0, lo=2.0
s_mov_b32(s[1], 0x45004400), # hi=5.0, lo=4.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_mul_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 8.0, places=1)
self.assertAlmostEqual(hi, 15.0, places=1)
def test_v_pk_fma_f16_basic(self):
"""V_PK_FMA_F16: D = A * B + C for packed f16."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x42004000), # A: hi=3.0, lo=2.0
s_mov_b32(s[1], 0x45004400), # B: hi=5.0, lo=4.0
s_mov_b32(s[2], 0x3c003c00), # C: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
v_pk_fma_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 9.0, places=1) # 2*4+1
self.assertAlmostEqual(hi, 16.0, places=0) # 3*5+1
def test_v_pk_add_f16_with_inline_constant(self):
"""V_PK_ADD_F16 with inline constant POS_ONE (1.0).
Inline constants for VOP3P are f16 values in the low 16 bits only.
hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
"""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x3c003c00), # packed f16: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
v_pk_add_f16(v[1], v[0], SrcEnum.POS_ONE, opsel_hi=3, opsel_hi2=1), # Add inline constant 1.0
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
# lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
self.assertAlmostEqual(lo, 2.0, places=2)
self.assertAlmostEqual(hi, 1.0, places=2)
def test_v_pk_mul_f16_with_inline_constant(self):
"""V_PK_MUL_F16 with inline constant POS_TWO (2.0).
Inline constant has value only in low 16 bits, hi is 0.
"""
from extra.assembly.amd.pcode import _f16
# v0 = packed (3.0, 4.0), multiply by POS_TWO
# lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
instructions = [
s_mov_b32(s[0], 0x44004200), # packed f16: hi=4.0, lo=3.0
v_mov_b32_e32(v[0], s[0]),
v_pk_mul_f16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 6.0, places=1)
self.assertAlmostEqual(hi, 0.0, places=1)
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
instructions.append(v_wmma_f32_16x16x16_f16(v[0], v[16], v[24], v[0]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_f16_with_accumulator(self):
"""V_WMMA_F32_16X16X16_F16 with non-zero accumulator."""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00))
instructions.append(s_mov_b32(s[1], f2i(5.0)))
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
instructions.append(v_wmma_f32_16x16x16_f16(v[0], v[16], v[24], v[0]))
st = run_program(instructions, n_lanes=32)
expected = f2i(21.0) # 16 + 5
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
class TestSpecialOps(unittest.TestCase):
"""Tests for special operations (SAD, PERM, DOT2)."""
def test_v_sad_u8_basic(self):
"""V_SAD_U8 computes sum of absolute differences."""
instructions = [
s_mov_b32(s[0], 0x04030201), # bytes: 1, 2, 3, 4
s_mov_b32(s[1], 0x05040302), # bytes: 2, 3, 4, 5
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_sad_u8(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# |1-2| + |2-3| + |3-4| + |4-5| = 1 + 1 + 1 + 1 = 4
self.assertEqual(st.vgpr[0][3], 4)
def test_v_sad_u8_identical_bytes(self):
"""V_SAD_U8 with identical inputs returns accumulator."""
instructions = [
s_mov_b32(s[0], 0x04030201),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 10),
v_mov_b32_e32(v[2], s[1]),
v_sad_u8(v[3], v[0], v[0], v[2]),
]
st = run_program(instructions, n_lanes=1)
# Same inputs -> SAD = 0, result = accumulator = 10
self.assertEqual(st.vgpr[0][3], 10)
def test_v_sad_u16_basic(self):
"""V_SAD_U16 computes sum of absolute differences of u16 pairs."""
instructions = [
s_mov_b32(s[0], 0x00030001), # hi=3, lo=1
s_mov_b32(s[1], 0x00050002), # hi=5, lo=2
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_sad_u16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# |1-2| + |3-5| = 1 + 2 = 3
self.assertEqual(st.vgpr[0][3], 3)
def test_v_sad_u32_basic(self):
"""V_SAD_U32 computes absolute difference of u32 values."""
instructions = [
s_mov_b32(s[0], 100),
s_mov_b32(s[1], 70),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_sad_u32(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# |100-70| = 30
self.assertEqual(st.vgpr[0][3], 30)
def test_v_msad_u8_masked(self):
"""V_MSAD_U8 masked SAD operation."""
instructions = [
s_mov_b32(s[0], 0x04030201),
s_mov_b32(s[1], 0x05040302),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_msad_u8(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
# V_MSAD_U8 skips bytes where src0 is 0
# Since no bytes are 0, result same as V_SAD_U8 = 4
self.assertEqual(st.vgpr[0][3], 4)
def test_v_perm_b32_select_bytes(self):
"""V_PERM_B32 selects bytes from two sources.
V_PERM_B32 concatenates {S1, S0} as a 64-bit value with S1 in low 32 bits.
Selector byte values 0-3 select from S1, values 4-7 select from S0.
"""
instructions = [
s_mov_b32(s[0], 0x44332211), # src0: bytes 4-7 in 64-bit view
s_mov_b32(s[1], 0x88776655), # src1: bytes 0-3 in 64-bit view
s_mov_b32(s[2], 0x07060504), # select bytes 4,5,6,7 (from src0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_perm_b32(v[2], v[0], v[1], s[2]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x44332211)
def test_v_dot2_f32_bf16_basic(self):
"""V_DOT2_F32_BF16 computes dot product of bf16 pairs."""
# bf16 1.0 = 0x3f80, bf16 2.0 = 0x4000
instructions = [
s_mov_b32(s[0], 0x3f803f80), # packed bf16: lo=1.0, hi=1.0
s_mov_b32(s[1], 0x40003f80), # packed bf16: lo=1.0, hi=2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot2_f32_bf16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
# 1.0*1.0 + 1.0*2.0 + 0 = 3.0
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 3.0, places=4)
class TestPackedMixedSigns(unittest.TestCase):
"""Tests for packed operations with mixed sign values."""
def test_pk_add_f16_mixed_signs(self):
"""V_PK_ADD_F16 with mixed positive/negative values."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0xc0003c00), # packed: hi=-2.0, lo=1.0
s_mov_b32(s[1], 0x3c003c00), # packed: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_add_f16(v[2], v[0], v[1], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 2.0, places=2) # 1.0 + 1.0
self.assertAlmostEqual(hi, -1.0, places=2) # -2.0 + 1.0
def test_pk_mul_f16_zero(self):
"""V_PK_MUL_F16 with zero."""
from extra.assembly.amd.pcode import _f16
instructions = [
s_mov_b32(s[0], 0x40004000), # packed: 2.0, 2.0
s_mov_b32(s[1], 0x00000000), # packed: 0.0, 0.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_mul_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
if __name__ == '__main__':
unittest.main()
+486
View File
@@ -0,0 +1,486 @@
"""Tests for VOPC instructions - vector compare operations.
Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
VCC = 106 # SGPR index for VCC_LO
class TestCmpClass(unittest.TestCase):
"""Tests for V_CMP_CLASS_F32 float classification."""
def test_cmp_class_quiet_nan(self):
"""V_CMP_CLASS_F32 detects quiet NaN."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
def test_cmp_class_signaling_nan(self):
"""V_CMP_CLASS_F32 detects signaling NaN."""
signal_nan = 0x7f800001
instructions = [
s_mov_b32(s[0], signal_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
def test_cmp_class_positive_inf(self):
"""V_CMP_CLASS_F32 detects +inf."""
pos_inf = 0x7f800000
instructions = [
s_mov_b32(s[0], pos_inf),
s_mov_b32(s[1], 0b1000000000), # bit 9 = +inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect +inf")
def test_cmp_class_negative_inf(self):
"""V_CMP_CLASS_F32 detects -inf."""
neg_inf = 0xff800000
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000100), # bit 2 = -inf
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect -inf")
def test_cmp_class_normal_positive(self):
"""V_CMP_CLASS_F32 detects positive normal."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
s_mov_b32(s[1], 0b0100000000), # bit 8 = positive normal
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
def test_cmp_class_normal_negative(self):
"""V_CMP_CLASS_F32 detects negative normal."""
instructions = [
v_mov_b32_e32(v[0], -1.0),
v_mov_b32_e32(v[1], 0b0000001000), # bit 3 = negative normal
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
def test_cmp_class_quiet_nan_not_signaling(self):
"""Quiet NaN does not match signaling NaN mask."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN only
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Quiet NaN should not match signaling mask")
def test_cmp_class_signaling_nan_not_quiet(self):
"""Signaling NaN does not match quiet NaN mask."""
signal_nan = 0x7f800001
instructions = [
s_mov_b32(s[0], signal_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN only
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
def test_v_cmp_sets_vcc_bits(self):
"""V_CMP_EQ sets VCC bits based on per-lane comparison."""
instructions = [
s_mov_b32(s[0], 5),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
class TestCmpClassF16(unittest.TestCase):
"""Tests for V_CMP_CLASS_F16 float classification.
Class bit mapping:
bit 0 = signaling NaN
bit 1 = quiet NaN
bit 2 = -infinity
bit 3 = -normal
bit 4 = -denormal
bit 5 = -zero
bit 6 = +zero
bit 7 = +denormal
bit 8 = +normal
bit 9 = +infinity
"""
def test_cmp_class_f16_positive_zero(self):
"""V_CMP_CLASS_F16: +zero matches bit 6."""
instructions = [
v_mov_b32_e32(v[0], 0x0000), # f16 +0.0
v_mov_b32_e32(v[1], 0x40), # bit 6 = +zero
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive zero")
def test_cmp_class_f16_negative_zero(self):
"""V_CMP_CLASS_F16: -zero matches bit 5."""
instructions = [
s_mov_b32(s[0], 0x8000), # f16 -0.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x20), # bit 5 = -zero
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative zero")
def test_cmp_class_f16_positive_normal(self):
"""V_CMP_CLASS_F16: +1.0 (normal) matches bit 8."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 +1.0
s_mov_b32(s[1], 0x100), # bit 8 = +normal
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
def test_cmp_class_f16_negative_normal(self):
"""V_CMP_CLASS_F16: -1.0 (normal) matches bit 3."""
instructions = [
s_mov_b32(s[0], 0xbc00), # f16 -1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x08), # bit 3 = -normal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
def test_cmp_class_f16_positive_infinity(self):
"""V_CMP_CLASS_F16: +inf matches bit 9."""
instructions = [
s_mov_b32(s[0], 0x7c00), # f16 +inf
s_mov_b32(s[1], 0x200), # bit 9 = +inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive infinity")
def test_cmp_class_f16_negative_infinity(self):
"""V_CMP_CLASS_F16: -inf matches bit 2."""
instructions = [
s_mov_b32(s[0], 0xfc00), # f16 -inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x04), # bit 2 = -inf
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative infinity")
def test_cmp_class_f16_quiet_nan(self):
"""V_CMP_CLASS_F16: quiet NaN matches bit 1."""
instructions = [
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x02), # bit 1 = quiet NaN
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
def test_cmp_class_f16_signaling_nan(self):
"""V_CMP_CLASS_F16: signaling NaN matches bit 0."""
instructions = [
s_mov_b32(s[0], 0x7c01), # f16 signaling NaN
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x01), # bit 0 = signaling NaN
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
def test_cmp_class_f16_positive_denormal(self):
"""V_CMP_CLASS_F16: positive denormal matches bit 7."""
instructions = [
v_mov_b32_e32(v[0], 1), # f16 +denormal (0x0001)
v_mov_b32_e32(v[1], 0x80), # bit 7 = +denormal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive denormal")
def test_cmp_class_f16_negative_denormal(self):
"""V_CMP_CLASS_F16: negative denormal matches bit 4."""
instructions = [
s_mov_b32(s[0], 0x8001), # f16 -denormal
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x10), # bit 4 = -denormal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative denormal")
def test_cmp_class_f16_combined_mask_zeros(self):
"""V_CMP_CLASS_F16: mask 0x60 covers both +zero and -zero."""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
v_mov_b32_e32(v[1], 0x60), # bits 5 and 6 (+-zero)
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x60")
def test_cmp_class_f16_combined_mask_1f8(self):
"""V_CMP_CLASS_F16: mask 0x1f8 covers -normal,-denorm,-zero,+zero,+denorm,+normal.
This is the exact mask used in the f16 sin kernel at PC=46.
"""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
s_mov_b32(s[0], 0x1f8),
v_mov_b32_e32(v[1], s[0]), # mask 0x1f8
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x1f8")
def test_cmp_class_f16_vop3_encoding(self):
"""V_CMP_CLASS_F16 in VOP3 encoding (v_cmp_class_f16_e64)."""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
s_mov_b32(s[0], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with VOP3 encoding")
def test_cmp_class_f16_vop3_normal_positive(self):
"""V_CMP_CLASS_F16 VOP3 encoding with +1.0 (normal)."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 +1.0
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +1.0 (normal) with mask 0x1f8")
def test_cmp_class_f16_vop3_nan_fails_mask(self):
"""V_CMP_CLASS_F16 VOP3: NaN should NOT match mask 0x1f8 (no NaN bits set)."""
instructions = [
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for NaN with mask 0x1f8 (no NaN bits)")
def test_cmp_class_f16_vop3_inf_fails_mask(self):
"""V_CMP_CLASS_F16 VOP3: +inf should NOT match mask 0x1f8 (no inf bits set)."""
instructions = [
s_mov_b32(s[0], 0x7c00), # f16 +inf
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
VOP3(VOP3Op.V_CMP_CLASS_F16, vdst=RawImm(VCC), src0=v[0], src1=s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for +inf with mask 0x1f8 (no inf bits)")
class TestCmpInt(unittest.TestCase):
"""Tests for integer comparison operations."""
def test_v_cmp_eq_u32(self):
"""V_CMP_EQ_U32 sets VCC bits based on per-lane comparison."""
instructions = [
s_mov_b32(s[0], 5),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
def test_cmp_eq_u16_opsel_lo_lo(self):
"""V_CMP_EQ_U16 comparing lo halves."""
instructions = [
s_mov_b32(s[0], 0x12340005), # lo=5, hi=0x1234
s_mov_b32(s[1], 0xABCD0005), # lo=5, hi=0xABCD
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_eq_u16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Lo halves should be equal")
def test_cmp_eq_u16_opsel_hi_hi(self):
"""V_CMP_EQ_U16 comparing hi halves with VOP3 opsel.
VOPC doesn't have opsel, so we use VOP3 form for hi-half comparisons.
VOP3 compares write result to SGPR via vdst field.
"""
instructions = [
s_mov_b32(s[2], 0x00051234), # hi=5, lo=0x1234
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x0005ABCD), # hi=5, lo=0xABCD
v_mov_b32_e32(v[1], s[2]),
# opsel=3 means compare hi halves, vdst=v[0] actually writes to s[0]
VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
# Result is in sgpr[0], not vcc
self.assertEqual(st.sgpr[0] & 1, 1, "Hi halves should be equal: 5==5")
def test_cmp_eq_u16_opsel_hi_hi_equal(self):
"""V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
instructions = [
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x12340009), # lo=9, hi=0x1234
v_mov_b32_e32(v[1], s[2]),
VOP3(VOP3Op.V_CMP_EQ_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
def test_cmp_gt_u16_opsel_hi(self):
"""V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
instructions = [
s_mov_b32(s[2], 0x99990005), # lo=5, hi=0x9999
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
v_mov_b32_e32(v[1], s[2]),
VOP3(VOP3Op.V_CMP_GT_U16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
class TestCmpFloat(unittest.TestCase):
"""Tests for float comparison operations."""
def test_v_cmp_lt_f16_vsrc1_hi(self):
"""V_CMP_LT_F16 with both operands from high half using VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
v_mov_b32_e32(v[1], s[2]),
# opsel=3 means read hi halves for both src0 and src1
VOP3(VOP3Op.V_CMP_LT_F16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
# Result is in sgpr[0]
self.assertEqual(st.sgpr[0] & 1, 1, "1.0 < 2.0 should be true")
def test_v_cmp_gt_f16_vsrc1_hi(self):
"""V_CMP_GT_F16 with both operands from high half using VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
v_mov_b32_e32(v[1], s[2]),
# opsel=3 means read hi halves for both src0 and src1
VOP3(VOP3Op.V_CMP_GT_F16, vdst=v[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
# Result is in sgpr[0]
self.assertEqual(st.sgpr[0] & 1, 1, "2.0 > 1.0 should be true")
def test_v_cmp_eq_f16_vsrc1_hi_equal(self):
"""v_cmp_eq_f16 with equal low and high halves."""
instructions = [
s_mov_b32(s[0], 0x42004200), # hi=3.0 (0x4200), lo=3.0 (0x4200)
v_mov_b32_e32(v[0], s[0]),
v_cmp_eq_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (3.0 == 3.0)")
def test_v_cmp_neq_f16_vsrc1_hi(self):
"""v_cmp_neq_f16 with different low and high halves."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0 (0x4000), lo=1.0 (0x3c00)
v_mov_b32_e32(v[0], s[0]),
v_cmp_lg_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")
def test_v_cmp_nge_f16_inf_self(self):
"""v_cmp_nge_f16 comparing -inf with itself (unordered less than).
Regression test: -inf < -inf should be false (IEEE 754).
"""
instructions = [
s_mov_b32(s[0], 0xFC00FC00), # both halves = -inf (0xFC00)
v_mov_b32_e32(v[0], s[0]),
v_cmp_nge_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Expected vcc=0 (-inf >= -inf)")
def test_v_cmp_f16_multilane(self):
"""v_cmp_lt_f16 with vsrc1=v128 across multiple lanes."""
instructions = [
# Lane 0: v0 = 0x40003c00 (hi=2.0, lo=1.0) -> 1.0 < 2.0 = true
# Lane 1: v0 = 0x3c004000 (hi=1.0, lo=2.0) -> 2.0 < 1.0 = false
v_mov_b32_e32(v[0], 0x40003c00), # default
v_cmp_eq_u32_e32(1, v[255]), # vcc = (lane == 1)
v_cndmask_b32_e64(v[0], v[0], 0x3c004000, SrcEnum.VCC_LO),
v_cmp_lt_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vcc & 1, 1, "Lane 0: expected vcc=1 (1.0 < 2.0)")
self.assertEqual((st.vcc >> 1) & 1, 0, "Lane 1: expected vcc=0 (2.0 < 1.0)")
class TestVCCBehavior(unittest.TestCase):
"""Tests for VCC condition code behavior."""
def test_vcc_all_lanes_true(self):
"""VCC should have all bits set when all lanes compare true."""
instructions = [
v_mov_b32_e32(v[0], 5),
v_mov_b32_e32(v[1], 5),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=32)
self.assertEqual(st.vcc, 0xFFFFFFFF, "All 32 lanes should be true")
def test_vcc_lane_dependent(self):
"""VCC should differ per lane based on lane_id comparison."""
instructions = [
v_mov_b32_e32(v[0], 16),
v_cmp_lt_u32_e32(v[255], v[0]), # lanes 0-15 are < 16
]
st = run_program(instructions, n_lanes=32)
self.assertEqual(st.vcc & 0xFFFF, 0xFFFF, "Lanes 0-15 should be true")
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
if __name__ == '__main__':
unittest.main()
@@ -9,10 +9,9 @@ os.environ["AMD"] = "1"
os.environ["MOCKGPU"] = "1"
os.environ["PYTHON_REMU"] = "1"
from extra.assembly.amd.emu import WaveState, decode_program, step_wave, WAVE_SIZE, set_valid_mem_ranges
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
from extra.assembly.amd.test.helpers import KernelInfo
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
from extra.assembly.amd.test.bench_emu import REMU_PATH
def _is_f32_nan(bits: int) -> bool:
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
@@ -92,19 +91,15 @@ class PythonEmulator:
def __init__(self):
self.state: WaveState | None = None
self.program: dict | None = None
self.lds: bytearray | None = None
self.n_lanes = 0
def create(self, kernel: bytes, n_lanes: int):
self.program = decode_program(kernel)
self.state = WaveState()
self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
self.state.exec_mask = (1 << n_lanes) - 1
self.lds = bytearray(65536)
self.n_lanes = n_lanes
def step(self) -> int:
assert self.program is not None and self.state is not None and self.lds is not None
return step_wave(self.program, self.state, self.lds, self.n_lanes)
assert self.program is not None and self.state is not None
return self.program[self.state.pc]._dispatch(self.state, self.program[self.state.pc])
def set_sgpr(self, idx: int, val: int):
assert self.state is not None
self.state.sgpr[idx] = val & 0xffffffff
@@ -163,8 +158,9 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
# v_div_scale/v_div_fixup: Rust has different VCC handling
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
# s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
'v_cvt_f16_f32'))
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
diffs = rust_before.diff(python_before, n_lanes)
if diffs:
trace_lines = []
@@ -397,6 +393,9 @@ class TestTinygradKernels(unittest.TestCase):
x_np = np.random.randn(16, 10).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
def test_sin_f64(self):
from tinygrad import dtypes
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
if __name__ == "__main__":
unittest.main()
File diff suppressed because it is too large Load Diff
+77 -2
View File
@@ -1,8 +1,9 @@
#!/usr/bin/env python3
"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
import unittest
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.dsl import encode_src
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import encode_src, RawImm
from extra.assembly.amd.decode import detect_format
class TestMUBUF(unittest.TestCase):
"""Test MUBUF (buffer) instructions."""
@@ -328,5 +329,79 @@ class TestVOP3Literal(unittest.TestCase):
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()
+4 -1
View File
@@ -2,7 +2,7 @@
# the Inst constructor should be looking at the types of the fields to correctly set the value
import unittest, struct
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.asm import asm
from extra.assembly.amd.test.test_roundtrip import compile_asm
@@ -21,6 +21,9 @@ class TestIntegration(unittest.TestCase):
self.assertEqual(repr(self.inst), repr(reasm))
print(desc)
def test_wmma(self):
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[189:192], v[140:143], v[0:7])
def test_load_b128(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
+1 -1
View File
@@ -1,7 +1,7 @@
#!/usr/bin/env python3
"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
import unittest, re, io, sys, subprocess
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.asm import waitcnt, asm
from extra.assembly.amd.test.helpers import get_llvm_mc
+104 -183
View File
@@ -1,201 +1,122 @@
#!/usr/bin/env python3
"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
import unittest, re, subprocess
"""Test AMD assembler/disassembler against LLVM test vectors."""
import unittest, re, subprocess, functools
from tinygrad.helpers import fetch
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.asm import asm
from extra.assembly.amd.asm import asm, disasm
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
# Format info: (filename, format_class, op_enum)
LLVM_TEST_FILES = {
# Scalar ALU
'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
# Vector ALU
'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp), # VOP3SD shares file with VOP3
'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp), # VOPCX uses VOPC format
# VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
# Memory
'ds': ('gfx11_asm_ds.s', DS, DSOp),
'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
# WMMA (matrix multiply)
'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
# Additional features
'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
# Alias files (alternative mnemonics)
'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
}
RDNA_FILES = ['gfx11_asm_sop1.s', 'gfx11_asm_sop2.s', 'gfx11_asm_sopp.s', 'gfx11_asm_sopk.s', 'gfx11_asm_sopc.s',
'gfx11_asm_vop1.s', 'gfx11_asm_vop2.s', 'gfx11_asm_vopc.s', 'gfx11_asm_vop3.s', 'gfx11_asm_vop3p.s', 'gfx11_asm_vinterp.s',
'gfx11_asm_vopd.s', 'gfx11_asm_vopcx.s', 'gfx11_asm_vop3_from_vop1.s', 'gfx11_asm_vop3_from_vop2.s', 'gfx11_asm_vop3_from_vopc.s',
'gfx11_asm_vop3_from_vopcx.s', 'gfx11_asm_ds.s', 'gfx11_asm_smem.s', 'gfx11_asm_flat.s', 'gfx11_asm_mubuf.s', 'gfx11_asm_mtbuf.s',
'gfx11_asm_mimg.s', 'gfx11_asm_wmma.s', 'gfx11_asm_vop3_features.s', 'gfx11_asm_vop3p_features.s', 'gfx11_asm_vopd_features.s',
'gfx11_asm_vop3_alias.s', 'gfx11_asm_vop3p_alias.s', 'gfx11_asm_vopc_alias.s', 'gfx11_asm_vopcx_alias.s', 'gfx11_asm_vinterp_alias.s',
'gfx11_asm_smem_alias.s', 'gfx11_asm_mubuf_alias.s', 'gfx11_asm_mtbuf_alias.s']
# CDNA test files - includes gfx9 files for shared instructions, plus gfx90a/gfx942 specific files
# gfx90a_ldst_acc.s has MIMG mixed in, filtered via is_mimg check
CDNA_FILES = ['gfx9_asm_sop1.s', 'gfx9_asm_sop2.s', 'gfx9_asm_sopp.s', 'gfx9_asm_sopk.s', 'gfx9_asm_sopc.s',
'gfx9_asm_vop1.s', 'gfx9_asm_vop2.s', 'gfx9_asm_vopc.s', 'gfx9_asm_vop3.s', 'gfx9_asm_vop3p.s',
'gfx9_asm_ds.s', 'gfx9_asm_flat.s', 'gfx9_asm_smem.s', 'gfx9_asm_mubuf.s', 'gfx9_asm_mtbuf.s',
'gfx90a_ldst_acc.s', 'gfx90a_asm_features.s', 'flat-scratch-gfx942.s', 'gfx942_asm_features.s',
'mai-gfx90a.s', 'mai-gfx942.s']
# RDNA4 (gfx12) test files - excludes alias/err/fake16/dpp files, and vimage/vsample (not supported)
# NOTE: vflat/vdsdir excluded - not implemented; features.s has mixed formats
RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx12_asm_sopk.s', 'gfx12_asm_sopc.s',
'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
'gfx12_asm_vbuffer_mubuf.s', 'gfx12_asm_vbuffer_mtbuf.s', 'gfx12_asm_wmma_w32.s', 'gfx12_asm_exp.s']
def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
"""Parse LLVM test format into (asm, expected_bytes) pairs."""
tests, lines = [], text.split('\n')
for i, line in enumerate(lines):
line = line.strip()
if not line or line.startswith(('//', '.', ';')): continue
asm_text = line.split('//')[0].strip()
if not asm_text: continue
for j in range(i, min(i + 3, len(lines))):
# Match GFX11, W32, or W64 encodings (all valid for gfx11)
# Format 1: "// GFX11: v_foo ... ; encoding: [0x01,0x02,...]"
# Format 2: "// GFX11: [0x01,0x02,...]" (used by DS, older files)
if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
elif m := re.search(r'(?:GFX11|W32|W64)[^:]*:\s*\[(0x[0-9a-fA-F,x\s]+)\]', lines[j]):
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
else:
continue
if hex_bytes:
try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
except ValueError: pass
break
def _is_mimg(data: bytes) -> bool: return (int.from_bytes(data[:4], 'little') >> 26) & 0x3f == 0b111100
def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
tests = []
for block in text.split('\n\n'):
asm_text, encoding = None, None
for line in block.split('\n'):
line = line.strip()
if not line or line.startswith(('.', ';')): continue
if not line.startswith('//'):
asm_text = line.split('//')[0].strip() or asm_text
if m := re.search(pattern + r'[^:]*:.*?(?:encoding:\s*)?\[(0x[0-9a-f,x\s]+)\]', line, re.I):
encoding = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
if asm_text and encoding:
try: tests.append((asm_text, bytes.fromhex(encoding)))
except ValueError: pass
return tests
def try_assemble(text: str):
"""Try to assemble instruction text, return bytes or None on failure."""
try: return asm(text).to_bytes()
except: return None
@functools.cache
def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]:
text = fetch(f"{LLVM_BASE}/{f}").read_bytes().decode('utf-8', errors='ignore')
if arch == "rdna3":
tests = _parse_llvm_tests(text, r'(?:GFX11|W32|W64)')
elif arch == "rdna4":
# Match GFX12 but not GFX1250 (which has different lit64 encoding)
tests = _parse_llvm_tests(text, r'(?:GFX12(?!50)|W32|W64)')
elif 'gfx90a' in f or 'gfx942' in f:
tests = _parse_llvm_tests(text, r'(?:GFX90A|GFX942)')
else:
tests = _parse_llvm_tests(text, r'(?:VI9|GFX9|CHECK)')
return [(a, d) for a, d in tests if not _is_mimg(d)] if arch == "cdna" else tests
def compile_asm_batch(instrs: list[str]) -> list[bytes]:
"""Compile multiple instructions with a single llvm-mc call."""
def _compile_asm_batch(instrs: list[str], arch: str = "rdna3") -> list[bytes]:
if not instrs: return []
asm_text = ".text\n" + "\n".join(instrs) + "\n"
result = subprocess.run(
[get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=asm_text, capture_output=True, text=True, timeout=30)
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
# Parse all encodings from output
results = []
for line in result.stdout.split('\n'):
if 'encoding:' not in line: continue
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
results.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
if len(results) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(results)}")
return results
mcpu = {'rdna3': 'gfx1100', 'rdna4': 'gfx1200'}.get(arch, 'gfx1100')
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True, timeout=30)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
return [bytes.fromhex(line.split('encoding:')[1].strip()[1:-1].replace('0x', '').replace(',', '').replace(' ', ''))
for line in result.stdout.split('\n') if 'encoding:' in line]
class TestLLVM(unittest.TestCase):
"""Test assembler and disassembler against all LLVM test vectors."""
tests: dict[str, list[tuple[str, bytes]]] = {}
@classmethod
def setUpClass(cls):
for name, (filename, _, _) in LLVM_TEST_FILES.items():
try:
data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
except Exception as e:
print(f"Warning: couldn't fetch {filename}: {e}")
cls.tests[name] = []
# Generate test methods dynamically for each format
def _make_asm_test(name):
def _make_test(f: str, arch: str, test_type: str):
def test(self):
passed, failed, skipped = 0, 0, 0
for asm_text, expected in self.tests.get(name, []):
result = try_assemble(asm_text)
if result is None: skipped += 1
elif result == expected: passed += 1
else: failed += 1
print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
self.assertEqual(failed, 0)
tests = _get_tests(f, arch)
name = f"{arch}_{test_type}_{f}"
if test_type == "roundtrip":
for _, data in tests:
decoded = detect_format(data, arch).from_bytes(data)
self.assertEqual(decoded.to_bytes()[:len(data)], data)
print(f"{name}: {len(tests)} passed")
elif test_type == "asm":
passed, skipped = 0, 0
for asm_text, expected in tests:
try:
self.assertEqual(asm(asm_text, arch).to_bytes(), expected)
passed += 1
except: skipped += 1
print(f"{name}: {passed} passed, {skipped} skipped")
elif test_type == "disasm":
to_test = []
for _, data in tests:
try:
decoded = decode_inst(data, arch)
# Skip if roundtrip fails, disasm fails, or op_name is missing (disasm starts with space)
if decoded.to_bytes()[:len(data)] == data and (d := disasm(decoded)) and not d.startswith(' '): to_test.append((data, d))
except: pass
skipped = len(tests) - len(to_test)
print(f"{name}: {len(to_test)} passed, {skipped} skipped")
if arch in ("rdna3", "rdna4"):
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
for (data, _), llvm in zip(to_test, _compile_asm_batch([t[1] for t in to_test], arch)): self.assertEqual(llvm, data)
return test
def _make_disasm_test(name):
def test(self):
_, fmt_cls, op_enum = LLVM_TEST_FILES[name]
# VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}
class TestLLVM(unittest.TestCase): pass
# First pass: decode all instructions and collect disasm strings
to_test: list[tuple[str, bytes, str | None, str | None]] = [] # (asm_text, data, disasm_str, error)
skipped = 0
for asm_text, data in self.tests.get(name, []):
if len(data) > fmt_cls._size(): continue
temp_inst = fmt_cls.from_bytes(data)
temp_op = temp_inst._values.get('op', 0)
temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
if temp_op in undocumented.get(name, set()): skipped += 1; continue
if name == 'sopp':
simm16 = temp_inst._values.get('simm16', 0)
simm16 = simm16.val if hasattr(simm16, 'val') else simm16
sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}
if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
try:
if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
temp = VOP3.from_bytes(data)
op_val = temp._values.get('op', 0)
op_val = op_val.val if hasattr(op_val, 'val') else op_val
is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
if is_vop3sd: VOP3SDOp(op_val)
else: VOP3Op(op_val)
else:
decoded = fmt_cls.from_bytes(data)
op_val = decoded._values.get('op', 0)
op_val = op_val.val if hasattr(op_val, 'val') else op_val
op_enum(op_val)
if decoded.to_bytes()[:len(data)] != data:
to_test.append((asm_text, data, None, "decode roundtrip failed"))
continue
to_test.append((asm_text, data, decoded.disasm(), None))
except Exception as e:
to_test.append((asm_text, data, None, f"exception: {e}"))
# Batch compile all disasm strings with single llvm-mc call
disasm_strs = [(i, t[2]) for i, t in enumerate(to_test) if t[2] is not None]
llvm_results = compile_asm_batch([s for _, s in disasm_strs]) if disasm_strs else []
llvm_map = {i: llvm_results[j] for j, (i, _) in enumerate(disasm_strs)}
# Match results back
passed, failed = 0, 0
failures: list[str] = []
for idx, (asm_text, data, disasm_str, error) in enumerate(to_test):
if error:
failed += 1; failures.append(f"{error} for {data.hex()}")
elif disasm_str is not None and idx in llvm_map:
llvm_bytes = llvm_map[idx]
if llvm_bytes is not None and llvm_bytes == data: passed += 1
elif llvm_bytes is not None: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
if failures[:10]: print(" " + "\n ".join(failures[:10]))
self.assertEqual(failed, 0)
return test
for name in LLVM_TEST_FILES:
setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
for f in RDNA_FILES:
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
setattr(TestLLVM, f"test_rdna3_asm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "asm"))
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
for f in CDNA_FILES:
setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
for f in RDNA4_FILES:
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
setattr(TestLLVM, f"test_rdna4_asm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "asm"))
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
if __name__ == "__main__":
unittest.main()
@@ -20,11 +20,12 @@ runner = get_runner(dev.device, si.ast)
prg = runner._prg
lib = bytearray(prg.lib)
# Find s_endpgm (0xBFB00000) and replace with invalid SOPP op=127 (0xBFFF0000)
# Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
# VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
found = False
for i in range(0, len(lib) - 4, 4):
if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
lib[i:i+4] = struct.pack("<I", 0xBFFF0000)
lib[i:i+4] = struct.pack("<I", 0x7E008400)
found = True
break
assert found, "s_endpgm not found"
@@ -46,8 +47,7 @@ dev.synchronize()
elapsed = time.perf_counter() - st
self.assertNotEqual(result.returncode, 0, "should have raised")
self.assertTrue("NotImplementedError" in result.stderr or "ValueError" in result.stderr,
f"expected NotImplementedError or ValueError in stderr")
self.assertTrue("Error" in result.stderr, f"expected an error in stderr, got: {result.stderr[:500]}")
# Should exit immediately, not wait for the full timeout
self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
+39 -76
View File
@@ -1,53 +1,16 @@
#!/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.generate import compile_pseudocode, _expr
from extra.assembly.amd.autogen.rdna3.gen_pcode import _VOP3SDOp_V_DIV_SCALE_F32, _VOPCOp_V_CMP_CLASS_F32
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
from extra.assembly.amd.test.helpers import ExecContext
from extra.assembly.amd.autogen.rdna3.str_pcode import VOP3SDOp_PCODE, VOPCOp_PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
class ExecContext:
"""Context for running compiled pseudocode (test-only)."""
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=MASK32, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
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_mod
ns = {k: v for k, v in vars(pcode_mod).items() if not k.startswith('_') or k in ('_f32', '_i32', '_f16', '_i16', '_f64', '_i64', '_bf16', '_ibf16',
'_div', '_sext', '_isnan', '_isquietnan', '_issignalnan', '_fma',
'_gt_neg_zero', '_lt_neg_zero', '_signext', '_to_f16_bits', '_pack')}
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': SliceProxy(self.EXEC, 31, 0), 'EXEC_HI': 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, Reg): ctx_reg._val = ns_val._val
else: ctx_reg._val = int(ns_val) & MASK64
if ns.get('SCC') is not self.SCC: _sync(self.SCC, ns['SCC'])
if ns.get('VCC') is not self.VCC: _sync(self.VCC, ns['VCC'])
if ns.get('EXEC') is not self.EXEC: _sync(self.EXEC, ns['EXEC'])
if ns.get('D0') is not self.D0: _sync(self.D0, ns['D0'])
if ns.get('D1') is not self.D1: _sync(self.D1, ns['D1'])
if ns.get('tmp') is not self.tmp: _sync(self.tmp, ns['tmp'])
if ns.get('saveexec') is not self.saveexec: _sync(self.saveexec, ns['saveexec'])
def result(self) -> dict:
return {"d0": self.D0._val, "scc": self.SCC._val & 1}
# Compile pseudocode functions on demand for regression tests
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', VOP3SDOp_PCODE[VOP3SDOp.V_DIV_SCALE_F32])
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', VOPCOp_PCODE[VOPCOp.V_CMP_CLASS_F32])
class TestReg(unittest.TestCase):
def test_u32_read(self):
@@ -83,7 +46,7 @@ class TestReg(unittest.TestCase):
class TestTypedView(unittest.TestCase):
def test_bit_slice(self):
r = Reg(0xDEADBEEF)
# Slices return SliceProxy which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
# Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
self.assertEqual(r.u32[7:0].u32, 0xEF)
self.assertEqual(r.u32[15:8].u32, 0xBE)
self.assertEqual(r.u32[23:16].u32, 0xAD)
@@ -108,7 +71,7 @@ class TestTypedView(unittest.TestCase):
# S0.u32[S1.u32[4:0]] - access bit at position from another register
s0 = Reg(0b11010101)
s1 = Reg(3)
bit_pos = s1.u32[4:0] # SliceProxy, int value = 3
bit_pos = s1.u32[4:0] # TypedView, int value = 3
bit_val = s0.u32[int(bit_pos)] # bit 3 of s0 = 0
self.assertEqual(int(bit_pos), 3)
self.assertEqual(bit_val, 0)
@@ -126,7 +89,7 @@ class TestTypedView(unittest.TestCase):
self.assertFalse(r1.u32 < r2.u32)
self.assertTrue(r1.u32 != r2.u32)
class TestSliceProxy(unittest.TestCase):
class TestTypedView(unittest.TestCase):
def test_slice_read(self):
r = Reg(0x56781234)
self.assertEqual(r[15:0].u16, 0x1234)
@@ -195,19 +158,19 @@ class TestExecContext(unittest.TestCase):
self.assertEqual(ctx.SCC._val, 0)
def test_ternary(self):
code = compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
ctx = ExecContext(s0=5, s1=3)
ctx.run(code)
self.assertEqual(ctx.D0._val, 1)
def test_pack(self):
code = compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
ctx = ExecContext(s0=0x1234, s1=0x5678)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0x56781234)
def test_tmp_with_typed_access(self):
code = compile_pseudocode("""tmp = S0.u32 + S1.u32
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
@@ -215,7 +178,7 @@ D0.u32 = tmp.u32""")
def test_s_add_u32_pattern(self):
# Real pseudocode pattern from S_ADD_U32
code = compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
# Test overflow case
@@ -225,7 +188,7 @@ D0.u32 = tmp.u32""")
self.assertEqual(ctx.SCC._val, 1) # Carry set
def test_s_add_u32_no_overflow(self):
code = compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
@@ -247,7 +210,7 @@ D0.u32 = tmp.u32""")
def test_for_loop(self):
# CTZ pattern - find first set bit
code = compile_pseudocode("""tmp = -1
code = _compile_pseudocode("""tmp = -1
for i in 0 : 31 do
if S0.u32[i] == 1 then
tmp = i
@@ -270,17 +233,17 @@ class TestPseudocodeRegressions(unittest.TestCase):
"""Regression tests for pseudocode instruction emulation bugs."""
def test_v_div_scale_f32_vcc_always_returned(self):
"""V_DIV_SCALE_F32 must always return vcc_lane, even when VCC=0 (no scaling needed).
Bug: when VCC._val == vcc (both 0), vcc_lane wasn't returned, so VCC bits weren't written.
"""V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
This caused division to produce wrong results for multiple lanes."""
# Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
s0 = 0x3f800000 # 1.0
s1 = 0x40400000 # 3.0
s2 = 0x3f800000 # 1.0 (numerator)
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None, {})
# Must always have vcc_lane in result
self.assertIn('vcc_lane', result, "V_DIV_SCALE_F32 must always return vcc_lane")
self.assertEqual(result['vcc_lane'], 0, "vcc_lane should be 0 when no scaling needed")
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
# Must always have VCC in result
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
def test_v_cmp_class_f32_detects_quiet_nan(self):
"""V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
@@ -289,28 +252,28 @@ class TestPseudocodeRegressions(unittest.TestCase):
signal_nan = 0x7f800001 # signaling NaN: exponent=255, bit22=0
# Test quiet NaN detection (bit 1 in mask)
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
self.assertEqual(result['vcc_lane'], 1, "Should detect quiet NaN with quiet NaN mask")
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
# Test signaling NaN detection (bit 0 in mask)
s1_signal = 0b0000000001 # bit 0 = signaling NaN
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
self.assertEqual(result['vcc_lane'], 1, "Should detect signaling NaN with signaling NaN mask")
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
# Test that quiet NaN doesn't match signaling NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
self.assertEqual(result['vcc_lane'], 0, "Quiet NaN should not match signaling NaN mask")
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
# Test that signaling NaN doesn't match quiet NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None, {})
self.assertEqual(result['vcc_lane'], 0, "Signaling NaN should not match quiet NaN mask")
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
def test_isnan_with_typed_view(self):
"""_isnan must work with TypedView objects, not just Python floats.
Bug: _isnan checked isinstance(x, float) which returned False for TypedView."""
def testisNAN_with_typed_view(self):
"""isNAN must work with TypedView objects, not just Python floats.
Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
nan_reg = Reg(0x7fc00000) # quiet NaN
normal_reg = Reg(0x3f800000) # 1.0
inf_reg = Reg(0x7f800000) # +inf
self.assertTrue(_isnan(nan_reg.f32), "_isnan should return True for NaN TypedView")
self.assertFalse(_isnan(normal_reg.f32), "_isnan should return False for normal TypedView")
self.assertFalse(_isnan(inf_reg.f32), "_isnan should return False for inf TypedView")
self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
class TestBF16(unittest.TestCase):
"""Tests for BF16 (bfloat16) support."""
@@ -349,7 +312,7 @@ class TestBF16(unittest.TestCase):
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
def test_bf16_slice_property(self):
"""Test SliceProxy.bf16 property."""
"""Test TypedView.bf16 property."""
r = Reg(0x40404040) # Two bf16 3.0 values
self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
+71
View File
@@ -0,0 +1,71 @@
#!/usr/bin/env python3
"""Test pdf.py PDF parser and enum generation."""
import unittest, tempfile, importlib.util
from extra.assembly.amd.pdf import extract, extract_tables, extract_enums, extract_pcode, write_enums, PDF_URLS
EXPECTED = {
"rdna3": {"pages": 655, "tables": 115, "sop2_ops": 67, "sop2_first": "S_ADD_U32"},
"rdna4": {"pages": 711, "tables": 125, "sop2_ops": 74, "sop2_first": "S_ADD_CO_U32"},
"cdna": {"pages": 610, "tables": 104, "sop2_ops": 52, "sop2_first": "S_ADD_U32"},
}
class TestPDF2(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.data = {name: extract(url) for name, url in PDF_URLS.items()}
cls.tables = {name: extract_tables(pages) for name, pages in cls.data.items()}
cls.enums = {name: extract_enums(cls.tables[name]) for name in PDF_URLS}
cls.pcode = {name: extract_pcode(cls.data[name], cls.enums[name]) for name in PDF_URLS}
def test_page_counts(self):
for name, exp in EXPECTED.items():
self.assertEqual(len(self.data[name]), exp["pages"], f"{name} page count")
def test_table_counts(self):
for name, exp in EXPECTED.items():
self.assertEqual(len(self.tables[name]), exp["tables"], f"{name} table count")
def test_tables_sequential(self):
for name in PDF_URLS:
nums = sorted(self.tables[name].keys())
missing = set(range(1, max(nums) + 1)) - set(nums)
self.assertEqual(missing, set(), f"{name} missing tables: {missing}")
def test_generate_enums(self):
for name, exp in EXPECTED.items():
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
write_enums(self.enums[name], name, f.name)
spec = importlib.util.spec_from_file_location("enum", f.name)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
# Check SOP2Op
self.assertTrue(hasattr(mod, 'SOP2Op'), f"{name} missing SOP2Op")
self.assertEqual(len(mod.SOP2Op), exp["sop2_ops"], f"{name} SOP2Op count")
self.assertEqual(mod.SOP2Op(0).name, exp["sop2_first"], f"{name} SOP2Op first")
# Check all enums have at least 2 ops
for attr in dir(mod):
if attr.endswith('Op'):
self.assertGreaterEqual(len(getattr(mod, attr)), 2, f"{name} {attr} has too few ops")
def test_pcode_rdna3_tricky(self):
"""Test specific pseudocode patterns that are tricky to extract correctly."""
pcode = self.pcode['rdna3']
# BUFFER_ATOMIC_MAX_U64: should have 4 statements (not truncated)
self.assertEqual(pcode[('BUFFER_ATOMIC_MAX_U64', 72)],
'tmp = MEM[ADDR].u64;\nsrc = DATA.u64;\nMEM[ADDR].u64 = src >= tmp ? src : tmp;\nRETURN_DATA.u64 = tmp')
# GLOBAL_STORE_B128: should have 4 MEM stores (not truncated)
self.assertEqual(pcode[('GLOBAL_STORE_B128', 29)],
'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\nMEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
# S_CMOVK_I32: should have full if/endif block
self.assertEqual(pcode[('S_CMOVK_I32', 2)],
"if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
def test_pcode_no_examples(self):
"""Pseudocode should not contain example lines with '=>'."""
for name in PDF_URLS:
for (op_name, opcode), code in self.pcode[name].items():
# No example lines (test vectors like "S_CTZ_I32_B32(0xaaaaaaaa) => 1")
self.assertNotIn('=>', code, f"{name} {op_name} contains example line with '=>'")
if __name__ == "__main__":
unittest.main()
-153
View File
@@ -1,153 +0,0 @@
#!/usr/bin/env python3
"""Test that PDF parser correctly extracts format fields."""
import unittest, os
from extra.assembly.amd.autogen.rdna3 import (
SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
SOP1Op, SOP2Op, VOP1Op, VOP3Op
)
# expected formats with key fields and whether they have ENCODING
EXPECTED_FORMATS = {
'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
'LDSDIR': (['VDST', 'OP'], True),
'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
'SOP1': (['OP', 'SDST', 'SSRC0'], True),
'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
'SOPK': (['OP', 'SDST', 'SIMM16'], True),
'SOPP': (['OP', 'SIMM16'], True),
'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOP1': (['OP', 'VDST', 'SRC0'], True),
'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
}
# Skip PDF parsing tests by default - only run with TEST_PDF_PARSER=1
# These are slow (~5s) and only needed when regenerating autogen/
@unittest.skipUnless(os.environ.get("TEST_PDF_PARSER"), "set TEST_PDF_PARSER=1 to run PDF parser tests")
class TestPDFParserGenerate(unittest.TestCase):
"""Test the PDF parser by running generate() and checking results."""
def test_pdf_parser(self):
"""Single test that validates all PDF parser outputs."""
from extra.assembly.amd.dsl import generate
result = generate()
# test_all_formats_present
for fmt_name in EXPECTED_FORMATS:
self.assertIn(fmt_name, result["formats"], f"missing format {fmt_name}")
# test_format_count
self.assertEqual(len(result["formats"]), 23)
# test_no_duplicate_fields
for fmt_name, fields in result["formats"].items():
field_names = [f[0] for f in fields]
self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
# test_expected_fields
for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
fields = {f[0] for f in result["formats"].get(fmt_name, [])}
for field in expected_fields:
self.assertIn(field, fields, f"{fmt_name} missing {field}")
if has_encoding:
self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
else:
self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
# test_vopd_no_dpp16_fields
vopd_fields = {f[0] for f in result["formats"].get("VOPD", [])}
for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
# test_dpp16_no_vinterp_fields
dpp16_fields = {f[0] for f in result["formats"].get("DPP16", [])}
for field in ['VDST', 'WAITEXP']:
self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
# test_sopp_no_smem_fields
sopp_fields = {f[0] for f in result["formats"].get("SOPP", [])}
for field in ['SBASE', 'SDATA']:
self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
class TestPDFParser(unittest.TestCase):
"""Verify format classes have correct fields from PDF parsing."""
def test_sop2_fields(self):
"""SOP2 should have op, sdst, ssrc0, ssrc1."""
for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
self.assertIn(field, SOP2._fields)
self.assertEqual(SOP2._fields['op'].hi, 29)
self.assertEqual(SOP2._fields['op'].lo, 23)
def test_sop1_fields(self):
"""SOP1 should have op, sdst, ssrc0 with correct bit positions."""
for field in ['op', 'sdst', 'ssrc0']:
self.assertIn(field, SOP1._fields)
self.assertNotIn('simm16', SOP1._fields)
self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
assert SOP1._encoding is not None
self.assertEqual(SOP1._encoding[0].hi, 31)
self.assertEqual(SOP1._encoding[1], 0b101111101)
def test_vop3sd_fields(self):
"""VOP3SD should have all fields including src0/src1/src2 from page continuation."""
for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
self.assertIn(field, VOP3SD._fields)
self.assertEqual(VOP3SD._fields['src0'].hi, 40)
self.assertEqual(VOP3SD._fields['src0'].lo, 32)
self.assertEqual(VOP3SD._size(), 8)
def test_flat_has_vdst(self):
"""FLAT should have vdst field."""
self.assertIn('vdst', FLAT._fields)
self.assertEqual(FLAT._fields['vdst'].hi, 63)
self.assertEqual(FLAT._fields['vdst'].lo, 56)
def test_encoding_bits(self):
"""Verify encoding bits are correct for major formats."""
tests = [
(SOP2, 31, 30, 0b10),
(SOPK, 31, 28, 0b1011),
(SOPP, 31, 23, 0b101111111),
(VOP1, 31, 25, 0b0111111),
(VOP2, 31, 31, 0b0),
(VOPC, 31, 25, 0b0111110),
(FLAT, 31, 26, 0b110111),
]
for cls, hi, lo, val in tests:
assert cls._encoding is not None
self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
def test_opcode_enums_exist(self):
"""Verify opcode enums are generated with expected counts."""
self.assertGreater(len(SOP1Op), 50)
self.assertGreater(len(SOP2Op), 50)
self.assertGreater(len(VOP1Op), 50)
self.assertGreater(len(VOP3Op), 200)
def test_vopd_no_duplicate_fields(self):
"""VOPD should not have duplicate fields and should not include DPP16 fields."""
field_names = list(VOPD._fields.keys())
self.assertEqual(len(field_names), len(set(field_names)))
for field in ['srcx0', 'srcy0', 'opx', 'opy']:
self.assertIn(field, VOPD._fields)
for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
self.assertNotIn(field, VOPD._fields)
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
import unittest, subprocess
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.test.helpers import get_llvm_mc
def llvm_assemble(asm: str) -> bytes:
+40 -45
View File
@@ -1,12 +1,18 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re, subprocess, os
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.asm import asm
from extra.assembly.amd.asm import detect_format
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump
# arch: (mcpu, mattr)
ARCH_CONFIG = {
'rdna3': ('gfx1100', '+real-true16,+wavefrontsize32'),
'rdna4': ('gfx1200', '+real-true16,+wavefrontsize32'),
'cdna': ('gfx942', '+wavefrontsize64'),
}
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
old_stdout = sys.stdout
@@ -30,61 +36,40 @@ def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
continue
return results
def compile_asm(instr: str, compiler=None) -> bytes:
"""Compile a single instruction with llvm-mc and return the machine code bytes."""
llvm_mc = get_llvm_mc()
result = subprocess.run(
[llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=f".text\n{instr}\n", capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed for '{instr}': {result.stderr.strip()}")
# Parse encoding: [0x01,0x39,0x0a,0x7e]
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
return bytes.fromhex(hex_vals)
raise RuntimeError(f"no encoding found in llvm-mc output for: {instr}")
def compile_asm(instr: str, arch: str = 'rdna3') -> bytes:
"""Compile a single instruction using LLVM."""
return compile_asm_batch([instr], arch)[0]
def compile_asm_batch(instrs: list[str]) -> list[bytes]:
def compile_asm_batch(instrs: list[str], arch: str = 'rdna3') -> list[bytes]:
"""Compile multiple instructions with a single llvm-mc call."""
if not instrs: return []
llvm_mc = get_llvm_mc()
src = ".text\n" + "\n".join(instrs) + "\n"
result = subprocess.run(
[llvm_mc, '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=src, capture_output=True, text=True)
mcpu, mattr = ARCH_CONFIG[arch]
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-show-encoding'],
input=".text\n" + "\n".join(instrs) + "\n", capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc batch failed: {result.stderr.strip()}")
# Parse all encodings in order
encodings = []
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
encodings.append(bytes.fromhex(hex_vals))
encodings.append(bytes.fromhex(enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')))
if len(encodings) != len(instrs): raise RuntimeError(f"expected {len(instrs)} encodings, got {len(encodings)}")
return encodings
def compile_and_disasm_batch(instrs: list[str], compiler) -> list[str]:
def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str]:
"""Compile instructions with LLVM and get LLVM's disassembly."""
import tempfile, os
import tempfile
if not instrs: return []
# Build assembly source with all instructions
src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n"
src += "\n".join(f" {instr}" for instr in instrs) + "\n"
# Use llvm-mc to assemble to object file
mcpu, mattr = ARCH_CONFIG[arch]
src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n" + "\n".join(f" {instr}" for instr in instrs) + "\n"
with tempfile.NamedTemporaryFile(suffix='.o', delete=False) as f:
obj_path = f.name
try:
result = subprocess.run(
[get_llvm_mc(), '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-filetype=obj', '-o', obj_path],
input=src, capture_output=True, text=True)
result = subprocess.run([get_llvm_mc(), '-triple=amdgcn', f'-mcpu={mcpu}', f'-mattr={mattr}', '-filetype=obj', '-o', obj_path],
input=src, capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-mc failed: {result.stderr.strip()}")
# Disassemble with llvm-objdump
result = subprocess.run([get_llvm_objdump(), '-d', '--mcpu=gfx1100', obj_path], capture_output=True, text=True)
result = subprocess.run([get_llvm_objdump(), '-d', f'--mcpu={mcpu}', obj_path], capture_output=True, text=True)
if result.returncode != 0: raise RuntimeError(f"llvm-objdump failed: {result.stderr.strip()}")
# Parse disassembly output
results: list[str] = []
for line in result.stdout.splitlines():
if '//' not in line: continue
@@ -96,6 +81,7 @@ def compile_and_disasm_batch(instrs: list[str], compiler) -> list[str]:
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
arch = 'rdna3'
def _test_kernel_roundtrip(self, op_fn):
"""Generate kernel from op_fn, test:
@@ -103,11 +89,14 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
2. asm(disasm()) matches LLVM output
3. our disasm() matches LLVM's disassembly string exactly
"""
arch = self.arch
mcpu, mattr = ARCH_CONFIG[arch]
from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.compiler_amd import HIPCompiler
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
compiler = HIPCompiler('gfx1100')
compiler = HIPCompiler(mcpu)
# First pass: decode all instructions and collect info
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
@@ -115,7 +104,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
offset = 0
while offset < len(kernel.code):
remaining = kernel.code[offset:]
fmt = detect_format(remaining)
fmt = detect_format(remaining, arch)
if fmt is None:
decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
offset += 4
@@ -152,11 +141,11 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
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_results = compile_asm_batch([d for _, d in asm_test_instrs], arch)
asm_llvm_map = {idx: result for (idx, _), result in zip(asm_test_instrs, asm_llvm_results)}
# Batch compile+disasm for disasm comparison test
disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], compiler)
disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], arch)
disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
# Now evaluate results
@@ -207,9 +196,9 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
else:
disasm_skipped += 1
print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"[{arch}] asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"[{arch}] disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
# Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
@@ -259,5 +248,11 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
@unittest.skip("no asm support for RDNA4")
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
@unittest.skip("no asm support for CDNA")
class TestTinygradKernelRoundtripCDNA(TestTinygradKernelRoundtrip): arch = 'cdna'
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,206 @@
#!/usr/bin/env python3
"""Tests for SQTT packet decoding using real captured examples."""
import pickle, unittest, ctypes
from pathlib import Path
from tinygrad.helpers import DEBUG, colored
from tinygrad.runtime.autogen import rocprof
from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES, InstOp, AluSrc, MemSrc)
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
# INST ops for non-traced SIMDs (excluded from instruction count)
OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
InstOp.OTHER_GLOBAL_STORE_VADDR_128}
PACKET_COLORS = {
"INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
"IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
"WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
"REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
}
def format_packet(p, time_offset: int = 0) -> str:
name, cycle = type(p).__name__, p._time - time_offset
if isinstance(p, INST):
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
elif hasattr(p, '_values'):
fields = " ".join(f"{k}=0x{v:x}" if k in {'snap', 'val32'} else f"{k}={v}"
for k, v in p._values.items() if not k.startswith('_') and k != 'delta')
else: fields = ""
return f"{cycle:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
def print_packets(packets: list) -> None:
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
time_offset = packets[0]._time if packets else 0
for p in packets:
if type(p).__name__ not in skip: print(format_packet(p, time_offset))
# ═══════════════════════════════════════════════════════════════════════════════
# ROCPROF DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int):
"""Run rocprof decoder on SQTT blobs, returning raw occupancy and instruction records."""
image, sections, _ = elf_loader(lib)
text = next((sh for sh in sections if sh.name == ".text"), None)
assert text is not None, "no .text section found"
text_off, text_size = text.header.sh_addr, text.header.sh_size
blob_iter, current_blob = iter(blobs), [None]
occupancy_records: list[tuple[int, int, int, int, bool]] = [] # (wave_id, simd, cu, time, is_start)
wave_insts: list[list[tuple[int, int]]] = [] # per-wave list of (time, stall)
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, _):
blob = next(blob_iter, None)
if blob is None: return 0
current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob)
buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte))
buf_size[0] = len(current_blob[0])
return len(current_blob[0])
@rocprof.rocprof_trace_decoder_trace_callback_t
def trace_cb(record_type, events_ptr, n, _):
if record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr):
occupancy_records.append((ev.wave_id, ev.simd, ev.cu, ev.time, ev.start))
elif record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr):
if ev.instructions_size > 0:
sz = ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t)
insts_blob = bytearray(sz)
ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
insts = list((rocprof.rocprofiler_thread_trace_decoder_inst_t * ev.instructions_size).from_buffer(insts_blob))
wave_insts.append([(inst.time, inst.stall) for inst in insts])
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
@rocprof.rocprof_trace_decoder_isa_callback_t
def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
offset = pc.address - base
if offset < text_off or offset >= text_off + text_size:
mem_size_ptr[0] = 0
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
try:
inst = decode_inst(data := image[offset:])
mem_size_ptr[0] = inst._size()
except (ValueError, AssertionError):
mem_size_ptr[0] = 0
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: mem_size_ptr[0] = 0
# rocprof parses instruction string to determine type; v_nop works for all
if (max_sz := size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
ctypes.memmove(instr_ptr, b"v_nop", min(5, max_sz - 1))
size_ptr[0] = min(5, max_sz - 1)
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
return occupancy_records, wave_insts
class TestSQTTExamples(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.examples = {}
for pkl_path in sorted(EXAMPLES_DIR.glob("*.pkl")):
with open(pkl_path, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
if sqtt_events and prg:
cls.examples[pkl_path.stem] = (sqtt_events, prg.lib, prg.base)
def test_examples_loaded(self):
self.assertGreater(len(self.examples), 0, "no example files found")
def test_decode_all_examples(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
packets = decode(event.blob)
if DEBUG >= 2: print(f"\n=== {name} event {i} ==="); print_packets(packets)
self.assertGreater(len(packets), 0, f"no packets decoded from {name} event {i}")
self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
def test_packet_types_valid(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
for pkt in decode(event.blob):
self.assertIn(type(pkt), PACKET_TYPES, f"unknown packet type {type(pkt)} in {name}")
def test_wave_lifecycle(self):
for name, (events, *_) in self.examples.items():
if "empty" in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, WAVESTART)]), 0, f"no WAVESTART in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
def test_time_monotonic(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
times = [p._time for p in decode(event.blob)]
self.assertEqual(times, sorted(times), f"timestamps not monotonic in {name}")
def test_gemm_has_instructions(self):
for name, (events, *_) in self.examples.items():
if "gemm" not in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, INST)]), 0, f"no INST packets in {name}")
def test_rocprof_wave_times_match(self):
"""Wave start/end times must match rocprof exactly."""
for name, (events, lib, base) in self.examples.items():
with self.subTest(example=name):
occupancy, _ = run_rocprof_decoder([e.blob for e in events], lib, base)
# extract from rocprof occupancy records
roc_starts: dict[tuple[int, int, int], int] = {}
roc_waves: list[tuple[int, int]] = []
for wave_id, simd, cu, time, is_start in occupancy:
key = (wave_id, simd, cu)
if is_start: roc_starts[key] = time
elif key in roc_starts: roc_waves.append((roc_starts.pop(key), time))
# extract from our decoder
our_waves: list[tuple[int, int]] = []
for event in events:
packets = decode(event.blob)
wave_starts: dict[tuple[int, int, int], int] = {}
for p in packets:
if isinstance(p, WAVESTART): wave_starts[(p.wave, p.simd, p.cu)] = p._time
elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
our_waves.append((wave_starts[key], p._time))
self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
def test_rocprof_inst_times_match(self):
"""Instruction times must match rocprof exactly (excluding s_endpgm)."""
for name, (events, lib, base) in self.examples.items():
with self.subTest(example=name):
_, wave_insts = run_rocprof_decoder([e.blob for e in events], lib, base)
# skip last inst per wave (s_endpgm) - it needs special handling (time + duration instead of time + stall)
roc_insts = [time + stall for insts in wave_insts for time, stall in insts[:-1]]
# extract from our decoder
our_insts: list[int] = []
for event in events:
for p in decode(event.blob):
if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
elif isinstance(p, VALUINST): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE_MASK):
for _ in range(bin(p.mask).count('1')): our_insts.append(p._time)
self.assertEqual(sorted(our_insts), sorted(roc_insts), f"instruction times mismatch in {name}")
if __name__ == "__main__":
unittest.main()
+101
View File
@@ -0,0 +1,101 @@
from typing import Callable, Any
from tinygrad import Tensor, dtypes, nn, UOp
from tinygrad.uop.ops import KernelInfo, AxisType, Ops
def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
x_abs_max = x.abs().max().detach()
scale = fp8_max / (x_abs_max + 1e-8)
x_scaled = x * scale
x_det = x_scaled.detach()
x_clamped = x_det.clamp(fp8_min, fp8_max)
x_clamped_ste = x_scaled + (x_clamped - x_det)
res = x_clamped_ste.cast(dtype)
return res, scale.float().reciprocal()
def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
SEQ = inp.shape[1]
OUT = weight.shape[0]
IN = weight.shape[-1]
seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
out_idx = UOp.range(OUT, 3, AxisType.LOOP)
batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
reduced = product.reduce(reduce_idx, arg=Ops.ADD)
store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
_, input_uop, weight_uop = kernel.src
input_tensor = Tensor(input_uop, device=input_uop.device)
grad_tensor = Tensor(gradient, device=gradient.device)
weight_tensor = Tensor(weight_uop, device=weight_uop.device)
grad_quantized, scale = quantize_to_fp8(grad_tensor)
scale_scalar = scale.reshape(())
grad_weight = Tensor.einsum("bso,bsi->oi", grad_quantized, input_tensor, dtype=dtypes.float)
grad_weight = grad_weight * scale_scalar
grad_2d = grad_quantized.reshape(grad_tensor.shape[0] * grad_tensor.shape[1], grad_tensor.shape[-1])
grad_input = (grad_2d.dot(weight_tensor, dtype=dtypes.float)).contiguous().reshape(input_tensor.shape) * scale
return (None, grad_input.uop, grad_weight.uop)
class FP8Linear:
def __init__(self, in_features:int, out_features:int, bias:bool=True):
self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.float32)
self.bias = Tensor.empty(out_features, dtype=dtypes.float32) if bias else None
def __call__(self, x: Tensor) -> Tensor:
original_ndim = len(x.shape)
if original_ndim == 2: x = x.reshape(x.shape[0], 1, x.shape[1])
batch, seq, _ = x.shape
w_fp8, w_scale = quantize_to_fp8(self.weight)
x_fp8, x_scale = quantize_to_fp8(x)
GPUS = self.weight.device
if isinstance(GPUS, tuple) and len(GPUS) > 1:
y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
else:
y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
y = y * w_scale * x_scale
if self.bias is not None: y = y + self.bias
if original_ndim == 2: y = y.reshape(batch, self.weight.shape[0])
return y.cast(x.dtype)
def _replace_linear(layer: nn.Linear):
fp8_linear = FP8Linear(layer.weight.shape[1], layer.weight.shape[0], layer.bias is not None)
fp8_linear.weight = layer.weight
if layer.bias is not None: fp8_linear.bias = layer.bias
return fp8_linear
def _swap_linear_with_fp8(model, module_filter_fn:Callable[[Any, str],bool]|None=None, fqn:str="", parent:Any|None=None,
attr_name:str="", visited:set|None=None):
if visited is None: visited = set()
if id(model) in visited: return
visited.add(id(model))
if isinstance(model, (str, int, float, bool, type(None), Tensor, UOp)): return
elif isinstance(model, nn.Linear):
if module_filter_fn is not None and not module_filter_fn(model, fqn): return
fp8_linear = _replace_linear(model)
if parent is not None and attr_name:
setattr(parent, attr_name, fp8_linear)
elif isinstance(model, list):
for i, item in enumerate(model):
child_fqn = f"{fqn}.{i}" if fqn else str(i)
if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[i] = _replace_linear(item)
else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
elif isinstance(model, dict):
for key, item in list(model.items()):
child_fqn = f"{fqn}.{key}" if fqn else str(key)
if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[key] = _replace_linear(item)
else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
elif hasattr(model, "__dict__"):
for attr_key in list(vars(model).keys()):
try: attr = getattr(model, attr_key)
except Exception: continue
child_fqn = f"{fqn}.{attr_key}" if fqn else attr_key
_swap_linear_with_fp8(attr, module_filter_fn, child_fqn, model, attr_key, visited)
def convert_to_float8_training(model, module_filter_fn:Callable[[Any,str],bool]|None=None):
_swap_linear_with_fp8(model, module_filter_fn, "", None, "")
return model
+655
View File
@@ -0,0 +1,655 @@
# RDNA3 128x128 tiled GEMM kernel - DSL version
# Computes C = A @ B for 4096x4096 float32 matrices using 128x128 tiles
#
# Architecture: RDNA3 (gfx1100)
# Tile size: 128x128 (each workgroup computes one tile of C)
# Workgroup: 128 threads (arranged as 32x4 for coalesced memory access)
# Inner loop: 8 iterations per K-block, processing 8 columns of A and 8 rows of B
#
# Accumulators: 128 vgprs (v[2-117], v[120-124], v[126-129], v[131-133])
import numpy as np
from pathlib import Path
from tinygrad import Tensor, Device, Context, GlobalCounters
from tinygrad.helpers import getenv, colored
from tinygrad.engine.realize import Runner, Estimates, ExecItem
from extra.assembly.amd.dsl import s, v, VCC_LO, RawImm, EXEC_LO
from extra.assembly.amd.autogen.rdna3.ins import *
# =============================================================================
# Kernel constants
# =============================================================================
LDS_SIZE = 8320 # Local data share size in bytes
MATRIX_DIM = 4096 # Matrix dimension N (assumes square NxN matrices)
LDS_A_STRIDE = 0x210 # LDS stride for A tile (528 bytes)
LDS_B_STRIDE = 0x200 # LDS stride for B tile (512 bytes)
LDS_BASE_OFFSET = 0x1080 # Base LDS offset for tiles
ADDR_MASK = 0x3fffff80 # Address alignment mask
# s_waitcnt encodings: wait for memory operations to complete
WAIT_LGKM = 64519 # wait for LDS/GDS/KMEM (lgkm_cnt=0)
WAIT_ALL = 0 # wait for everything
WAIT_VMEM = 1015 # wait for VMEM only (vm_cnt=0, lgkm_cnt=63)
# =============================================================================
# Named register assignments (VGPRs) - COMPACT LAYOUT
# =============================================================================
V_LANE_ID_MOD8 = 214 # lane_id & 7 (column within 8-wide tile chunk)
V_OUTPUT_ROW = 131 # output row coordinate
V_LANE_MOD8_X4 = 134 # V_LANE_ID_MOD8 << 2 (byte offset)
V_LANE_DIV8_X4 = 135 # (lane_id >> 3) << 2
V_ADDR_HI_ZERO = 136 # always 0 (for 64-bit address high bits)
V_LDS_A_BASE = 133 # LDS A-tile base address for inner loop (in ACC_RESERVED gap)
V_LDS_B_BASE = 130 # LDS B-tile base address for inner loop (in ACC_RESERVED gap)
V_GLOBAL_A_ADDR = 131 # global memory A prefetch address (reuses V_OUTPUT_ROW slot during main loop)
V_GLOBAL_B_ADDR = 154 # global memory B prefetch address
# LDS tile register destinations - SEPARATE from DATA to avoid overlap
# DATA regs (v155-170) receive global prefetch
# A on banks 2-3, B on banks 0-1 to avoid bank conflicts in VOPD
# This layout matches kernel8's optimization for VGPR cache utilization
V_A_TILE_REGS = [186, 190, 194, 198] # A tile: banks 2,2,2,2 (186%4=2, 190%4=2, etc.)
V_B_TILE_REGS = [184, 188, 192, 196, 200, 204, 208, 212] # B tile: banks 0,0,0,0,0,0,0,0
# =============================================================================
# Named register assignments (SGPRs)
# =============================================================================
S_OUT_PTR = (0, 1) # output C matrix base pointer
S_TILE_X = 2 # workgroup_x << 7
S_TILE_Y = 3 # workgroup_y << 7
S_DIM_N = 4 # matrix dimension N
S_LOOP_BOUND = 7 # K-8 (loop termination bound)
S_A_PTR = (8, 9) # A matrix base pointer
S_B_PTR = (10, 11) # B matrix base pointer
S_LOOP_CTR = 12 # loop counter (increments by 8)
S_PREFETCH_FLAG = 13 # prefetch condition flag / row stride in epilogue
S_WORKGROUP_X = 14 # workgroup_id_x
S_WORKGROUP_Y = 15 # workgroup_id_y
# Kernarg load destinations (before copy to working regs)
S_KERNARG_OUT = (16, 17) # output pointer from kernarg
S_KERNARG_A = (20, 21) # A pointer from kernarg
S_KERNARG_B = (22, 23) # B pointer from kernarg
# Prefetch base pointers (8 pairs each, 16KB/256KB apart)
S_PREFETCH_B = 24 # s[24:39] - 8 B tile pointers
S_PREFETCH_A = 40 # s[40:55] - 8 A tile pointers
# =============================================================================
# Data tables
# =============================================================================
# Accumulator grid: ACC_GRID[a_idx][b_idx] = vgpr for C[a,b]
# a_idx: which A value (0-7), b_idx: which B value (0-15)
# Scattered due to VOPD bank constraints (vdst_x % 4 != vdst_y % 4)
# Range is from v2 - v129
ACC_GRID = [
[ 5, 3, 9, 8, 37, 35, 41, 40, 69, 67, 73, 72, 101, 99,105,104], # a0
[ 4, 2, 7, 6, 36, 34, 39, 38, 68, 66, 71, 70, 100, 98,103,102], # a1
[ 17, 16, 13, 11, 49, 48, 45, 43, 81, 80, 77, 75, 113,112,109,107], # a2
[ 15, 14, 12, 10, 47, 46, 44, 42, 79, 78, 76, 74, 111,110,108,106], # a3
[ 21, 19, 25, 24, 53, 51, 57, 56, 85, 83, 89, 88, 117,115,121,120], # a4
[ 20, 18, 23, 22, 52, 50, 55, 54, 84, 82, 87, 86, 116,114,123,122], # a5
[125,128, 29, 27, 33, 32, 61, 59, 65, 64, 93, 91, 97, 96,129,127], # a6
[119,118, 28, 26, 31, 30, 60, 58, 63, 62, 92, 90, 95, 94,124,126], # a7
]
# Optimized (a_pair, b_pair) iteration order for better GPU scheduling
# Interleaves A and B pairs to maximize instruction-level parallelism
FMAC_PAIR_ORDER = [
(0,0),(0,1),(1,1),(1,0), (2,0),(2,1),(3,1),(3,2), (0,2),(0,3),(1,3),(1,2), (2,2),(2,3),(3,3),(3,4),
(0,4),(0,5),(1,5),(1,4), (2,4),(2,5),(3,5),(3,6), (0,6),(0,7),(1,7),(1,6), (2,6),(2,7),(3,7),(3,0),
]
def derive_fmac_pattern(acc_grid, a_tile_regs=None, b_tile_regs=None):
"""Generate 64 dual FMAC ops from accumulator grid with optimized iteration order."""
if a_tile_regs is None: a_tile_regs = V_A_TILE_REGS
if b_tile_regs is None: b_tile_regs = V_B_TILE_REGS
pattern = []
for idx, (a_pair, b_pair) in enumerate(FMAC_PAIR_ORDER):
a_even, a_odd = a_pair * 2, a_pair * 2 + 1
b_even, b_odd = b_pair * 2, b_pair * 2 + 1
a_base, b_base = a_tile_regs[a_pair], b_tile_regs[b_pair]
# Op 1: normal order -> C[a_even, b_even] + C[a_odd, b_odd]
pattern.append((acc_grid[a_even][b_even], acc_grid[a_odd][b_odd],
a_base, b_base, a_base+1, b_base+1))
# Op 2: alternate swapping A vs B to vary register banks
if idx % 2 == 0: # swap B
pattern.append((acc_grid[a_even][b_odd], acc_grid[a_odd][b_even],
a_base, b_base+1, a_base+1, b_base))
else: # swap A
pattern.append((acc_grid[a_odd][b_even], acc_grid[a_even][b_odd],
a_base+1, b_base, a_base, b_base+1))
return pattern
# Derived: 64 dual FMAC operations
FMAC_PATTERN = derive_fmac_pattern(ACC_GRID)
def derive_permute_swaps(acc_grid, out_regs):
"""Derive swap sequence to permute accumulators from FMAC layout to output order.
After FMAC loop: acc_grid[a][b] holds C[a,b]
Output order: for row_half in 0,1; col_group in 0-3; row_in_group in 0-3; b_off in 0-3
-> need C[row_half*4 + row_in_group, col_group*4 + b_off] in descending reg order
"""
def target_ab(i):
row_half, col_group = i // 64, (i // 16) % 4
row_in_group, b_off = (i // 4) % 4, i % 4
return (row_half * 4 + row_in_group, col_group * 4 + b_off)
reg_contents = {acc_grid[a][b]: (a, b) for a in range(8) for b in range(16)}
ab_location = {ab: r for r, ab in reg_contents.items()}
swaps = []
for i in range(128):
target_reg, needed_ab = out_regs[i], target_ab(i)
current_reg = ab_location[needed_ab]
if current_reg != target_reg:
swaps.append((current_reg, target_reg))
ab_at_target = reg_contents.get(target_reg)
reg_contents[target_reg], ab_location[needed_ab] = needed_ab, target_reg
if ab_at_target is not None:
reg_contents[current_reg], ab_location[ab_at_target] = ab_at_target, current_reg
return swaps
# Derived: swap sequence to arrange accumulators for output
OUT_REGS = list(range(129, 1, -1))
PERMUTE_SWAPS = derive_permute_swaps(ACC_GRID, OUT_REGS)
# =============================================================================
# LDS tile staging registers - COMPACT LAYOUT
# =============================================================================
# DATA regs receive contiguous global prefetch, then write to LDS
# TILE regs receive scattered LDS loads (ds_load_b64 pairs), then feed FMACs
# These are SEPARATE - DATA lives during prefetch/store, TILE lives during inner loop
V_LDS_A_ADDR = 153 # single base register for A stores (use +512 offsets)
V_LDS_A_DATA = list(range(155, 163)) # 8 data registers for A prefetch (v155-162)
V_LDS_B_ADDR = 145 # single base register for B stores (use 16-bit offsets)
V_LDS_B_DATA = list(range(163, 171)) # 8 data registers for B prefetch (v163-170)
# Global memory prefetch schedule: (vdst1, vdst2, addr_vreg, saddr_lo1, saddr_lo2)
# First 2 pairs from B prefetch pointers (s[32:39]), next 4 pairs from A prefetch pointers (s[40:55])
PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR, S_PREFETCH_B+8+4*i, S_PREFETCH_B+10+4*i) for i in range(2)] + \
[(V_LDS_B_DATA[2*(i-2)], V_LDS_B_DATA[2*(i-2)+1], V_GLOBAL_A_ADDR, S_PREFETCH_A+4*(i-2), S_PREFETCH_A+2+4*(i-2)) for i in range(2, 6)]
# Initial tile prefetch: (vdst, saddr_lo) - load into A data regs using B prefetch pointers (s[24:31])
INIT_PREFETCH = [(V_LDS_A_DATA[i], S_PREFETCH_B+2*i) for i in range(4)]
# Initial tile loads: (vdst, addr_lo) pairs - use temp regs in accumulator gaps
INIT_TILE_LOADS = [(23,5),(24,9),(25,7),(26,2),(27,11),(28,13),(29,6),(30,8),(31,10),(12,12),(13,14),(3,2),(4,4),(5,8),(6,6),(7,10)]
# A matrix row offset registers (scattered to avoid accumulator conflicts)
ROW_REGS = list(range(137, 145)) # v137-v144 (8 regs)
# =============================================================================
# Kernel class
# =============================================================================
class Kernel:
def __init__(self, arch='gfx1100'):
self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
def emit(self, inst): self.instructions.append(inst); return inst
def label(self, name): self.labels[name] = len(self.instructions)
def branch_to(self, label): self.branch_targets[len(self.instructions) - 1] = label
def add64(self, dst_lo, dst_hi, src_lo, src_hi, off):
"""s[dst_lo:dst_hi] = s[src_lo:src_hi] + off"""
if off: self.emit(s_add_u32(s[dst_lo], s[src_lo], off)); self.emit(s_addc_u32(s[dst_hi], s[src_hi], 0))
elif dst_lo != src_lo: self.emit(s_mov_b64(s[dst_lo:dst_hi], s[src_lo:src_hi]))
def global_load(self, vdst, addr, saddr=None):
"""Global load b32"""
self.emit(global_load_b32(vdst=v[vdst], addr=v[addr:addr+1],
saddr=s[saddr:saddr+2] if saddr else RawImm(124)))
def waitcnt(self, lgkm=None, vm=None):
"""Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
from extra.assembly.amd.asm import waitcnt as encode_waitcnt
if lgkm == 0 and vm is None: self.emit(s_waitcnt(simm16=WAIT_LGKM))
elif vm == 0 and lgkm is None: self.emit(s_waitcnt(simm16=WAIT_VMEM))
elif lgkm == 0 and vm == 0: self.emit(s_waitcnt(simm16=WAIT_ALL))
elif vm is not None and lgkm is None:
self.emit(s_waitcnt(simm16=encode_waitcnt(vmcnt=vm, expcnt=7, lgkmcnt=63)))
elif lgkm is not None and vm is None:
self.emit(s_waitcnt(simm16=encode_waitcnt(vmcnt=63, expcnt=7, lgkmcnt=lgkm)))
else: raise ValueError(f"unsupported waitcnt: lgkm={lgkm}, vm={vm}")
def barrier(self): self.emit(s_barrier())
def to_asm(self):
import re
# Instruction stream with labels
label_at = {pos: name for name, pos in self.labels.items()}
body = []
for i, inst in enumerate(self.instructions):
if i in label_at: body.append(f'.{label_at[i]}:')
asm = inst.disasm()
if i in self.branch_targets:
asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
body.append('\t' + asm)
# HSA kernel descriptor attributes (zeros included for compatibility)
hsa = [
('group_segment_fixed_size', LDS_SIZE), ('private_segment_fixed_size', 0), ('kernarg_size', 36),
('user_sgpr_count', 14), ('user_sgpr_dispatch_ptr', 0), ('user_sgpr_queue_ptr', 0),
('user_sgpr_kernarg_segment_ptr', 1), ('user_sgpr_dispatch_id', 0), ('user_sgpr_private_segment_size', 0),
('wavefront_size32', 1), ('uses_dynamic_stack', 0), ('enable_private_segment', 0),
('system_sgpr_workgroup_id_x', 1), ('system_sgpr_workgroup_id_y', 1), ('system_sgpr_workgroup_id_z', 0),
('system_sgpr_workgroup_info', 0), ('system_vgpr_workitem_id', 0), ('next_free_vgpr', 214),
('next_free_sgpr', 16), ('float_round_mode_32', 0), ('float_round_mode_16_64', 0),
('float_denorm_mode_32', 3), ('float_denorm_mode_16_64', 3), ('dx10_clamp', 1), ('ieee_mode', 1),
('fp16_overflow', 0), ('workgroup_processor_mode', 0), ('memory_ordered', 1), ('forward_progress', 0),
('shared_vgpr_count', 0)]
return '\n'.join([
'\t.text', f'\t.amdgcn_target "amdgcn-amd-amdhsa--{self.arch}"',
'\t.protected\tkernel', '\t.globl\tkernel', '\t.p2align\t8', '\t.type\tkernel,@function', 'kernel:',
*body,
'\t.section\t.rodata,"a",@progbits', '\t.p2align\t6, 0x0', '\t.amdhsa_kernel kernel',
*[f'\t\t.amdhsa_{k} {v}' for k, v in hsa],
'\t.end_amdhsa_kernel', '\t.text', '.Lfunc_end0:', '\t.size\tkernel, .Lfunc_end0-kernel',
'\t.amdgpu_metadata', '---', 'amdhsa.kernels:', ' - .args:',
*[f' - .address_space: global\n .offset: {i*8}\n .size: 8\n .value_kind: global_buffer' for i in range(3)],
f' .group_segment_fixed_size: {LDS_SIZE}', ' .kernarg_segment_align: 8',
' .kernarg_segment_size: 24', ' .max_flat_workgroup_size: 128', ' .name: kernel',
' .private_segment_fixed_size: 0', ' .sgpr_count: 60', ' .symbol: kernel.kd',
' .vgpr_count: 214', ' .wavefront_size: 32', f'amdhsa.target: amdgcn-amd-amdhsa--{self.arch}',
'amdhsa.version:', ' - 1', ' - 2', '...', '\t.end_amdgpu_metadata'])
# =============================================================================
# Kernel builder
# =============================================================================
def build_kernel(arch='gfx1100'):
k = Kernel(arch)
# ===========================================================================
# PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
# ===========================================================================
k.emit(s_load_b128(sdata=s[S_KERNARG_A[0]:S_KERNARG_B[1]], sbase=s[0:1], offset=0x0, soffset=RawImm(124)))
k.emit(s_load_b64(sdata=s[S_KERNARG_OUT[0]:S_KERNARG_OUT[1]], sbase=s[0:1], offset=0x10, soffset=RawImm(124)))
k.emit(s_mov_b32(s[S_DIM_N], MATRIX_DIM))
k.emit(s_mov_b32(s[S_LOOP_CTR], 0)) # used by LDS swizzle, always 0 for valid workgroups
k.emit(s_lshl_b32(s[S_TILE_X], s[S_WORKGROUP_X], 7))
k.emit(s_lshl_b32(s[S_TILE_Y], s[S_WORKGROUP_Y], 7))
# Lane-derived values
k.emit(v_and_b32_e32(v[V_LANE_ID_MOD8], 7, v[0]))
k.emit(v_lshrrev_b32_e32(v[4], 3, v[0]))
k.emit(v_or_b32_e32(v[1], s[S_TILE_X], v[0]))
k.emit(v_or_b32_e32(v[22], s[S_TILE_Y], v[4]))
k.emit(v_lshlrev_b32_e32(v[V_LANE_MOD8_X4], 2, v[V_LANE_ID_MOD8]))
k.emit(v_mov_b32_e32(v[2], 0)) # v[1] always positive, sign extension is 0
k.emit(v_lshlrev_b64(v[5:6], 2, v[1:2]))
k.waitcnt(lgkm=0)
# Copy pointers to working registers
k.emit(s_mov_b64(s[S_OUT_PTR[0]:S_OUT_PTR[1]], s[S_KERNARG_OUT[0]:S_KERNARG_OUT[1]]))
k.emit(s_mov_b64(s[S_A_PTR[0]:S_A_PTR[1]], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))
k.emit(s_mov_b64(s[S_B_PTR[0]:S_B_PTR[1]], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))
# Compute 8 A and B matrix tile base pointers for prefetch
for i in range(8): k.add64(S_PREFETCH_B + i*2, S_PREFETCH_B + i*2 + 1, S_KERNARG_B[0], S_KERNARG_B[1], i * 0x4000) # B: 16KB apart
for i in range(8): k.add64(S_PREFETCH_A + i*2, S_PREFETCH_A + i*2 + 1, S_KERNARG_A[0], S_KERNARG_A[1], i * 0x40000) # A: 256KB apart
# Global prefetch addresses: B = (tile_x + lane_id) * 4, A = ((tile_y << 12) + (lane_id/8)*4K + lane_id%8) * 4
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[0]))
k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_B_ADDR], 2, v[V_GLOBAL_B_ADDR]))
k.emit(s_lshl_b32(s[19], s[S_TILE_Y], 12))
k.emit(v_lshl_add_u32(v[V_GLOBAL_A_ADDR], v[4], 12, v[V_LANE_ID_MOD8])) # (lane_id/8)*4K + lane_id%8
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], s[19], v[V_GLOBAL_A_ADDR]))
k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_A_ADDR], 2, v[V_GLOBAL_A_ADDR]))
# ===========================================================================
# Tile address computation for initial A/B matrix loads
# ===========================================================================
k.emit(s_lshl_b32(s[S_LOOP_BOUND], s[S_DIM_N], 4)) # row stride = 16*N
k.emit(v_mul_lo_u32(v[ROW_REGS[0]], v[22], s[S_DIM_N])) # A matrix row offsets
for i in range(1, 8): k.emit(v_add_nc_u32_e32(v[ROW_REGS[i]], s[S_LOOP_BOUND], v[ROW_REGS[i-1]]))
def addr64(dst, base_s): # 64-bit address: v[dst:dst+1] = s[base_s:base_s+1] + v[dst]*4
k.emit(v_mov_b32_e32(v[dst+1], 0)) # offset always positive, sign ext = 0
k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[dst:dst+1]))
k.emit(v_add_co_u32(v[dst], VCC_LO, s[base_s], v[dst]))
k.emit(v_add_co_ci_u32_e32(v[dst+1], s[base_s+1], v[dst+1]))
def b_addr(dst, mult, tmp=None): # B address for col + mult*N
tmp = tmp if tmp is not None else dst
k.emit(v_mad_u32_u24(v[tmp], s[S_DIM_N], mult, v[1]))
if tmp != dst:
k.emit(v_mov_b32_e32(v[tmp+1], 0)) # offset always positive
k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[tmp:tmp+1]))
k.emit(v_add_co_u32(v[dst], VCC_LO, s[S_B_PTR[0]], v[dst]))
k.emit(v_add_co_ci_u32_e32(v[dst+1], s[S_B_PTR[1]], v[dst+1]))
else: addr64(dst, S_B_PTR[0])
def a_addr(dst, row_reg, tmp): # A address for row_reg + lane_id_mod8
k.emit(v_add_nc_u32_e32(v[tmp], v[row_reg], v[V_LANE_ID_MOD8]))
k.emit(v_mov_b32_e32(v[tmp+1], 0)) # offset always positive
k.emit(v_lshlrev_b64(v[dst:dst+1], 2, v[tmp:tmp+1]))
k.emit(v_add_co_u32(v[dst], VCC_LO, s[S_A_PTR[0]], v[dst]))
k.emit(v_add_co_ci_u32_e32(v[dst+1], s[S_A_PTR[1]], v[dst+1]))
# Batch 1: B addresses (cols 0-5) and loads
k.emit(v_add_co_u32(v[5], VCC_LO, s[S_B_PTR[0]], v[5]))
k.emit(v_add_co_ci_u32_e32(v[6], s[S_B_PTR[1]], v[6]))
for dst, mult in [(9,1), (7,2), (2,3), (11,4), (13,5)]: b_addr(dst, mult)
k.emit(s_clause(simm16=5)) # 6 consecutive global loads
for vdst, addr in INIT_TILE_LOADS[:6]: k.global_load(vdst, addr)
# Batch 2: A addresses (rows 0-4) and loads
for dst, ri in [(6,0), (8,1), (10,2), (12,3), (14,4)]:
k.emit(v_add_nc_u32_e32(v[dst], v[ROW_REGS[ri]], v[V_LANE_ID_MOD8]))
addr64(dst, S_A_PTR[0])
k.emit(s_clause(simm16=4)) # 5 consecutive global loads
for vdst, addr in INIT_TILE_LOADS[6:11]: k.global_load(vdst, addr)
# Batch 3: B cols 6-7, A rows 5-7, and loads
for dst, mult, tmp in [(2,6,15), (4,7,4)]: b_addr(dst, mult, tmp)
for dst, ri, tmp in [(8,5,16), (6,6,18), (10,7,20)]: a_addr(dst, ROW_REGS[ri], tmp)
k.emit(s_clause(simm16=4)) # 5 consecutive global loads
for vdst, addr in INIT_TILE_LOADS[11:]: k.global_load(vdst, addr)
# ===========================================================================
# LDS store address computation (bank-conflict-avoiding swizzle)
# ===========================================================================
# This section computes LDS store addresses with a swizzle pattern to avoid bank conflicts.
# Key outputs:
# v[8]: A-tile initial store base (used only for initial stores with stride64)
# V_LDS_B_ADDR (v145): B-tile store base (used for both initial and main loop)
# V_LANE_DIV8_X4 (v135): (lane_id >> 3) << 2 for epilogue
#
# The swizzle ensures that threads in the same wavefront write to different LDS banks.
# Formula: swizzled_addr = base + (lane_id & 7) * LDS_A_STRIDE + swizzle_offset
# where swizzle_offset depends on (lane_id >> 3) to distribute across banks.
# v[22] = tile_y | (lane_id >> 3) from prologue, used as base for row offsets
# Compute 7 row offsets for B-tile rows 1-7 (row 0 computed separately in v[9])
k.emit(v_add_nc_u32_e32(v[9], s[S_LOOP_CTR], v[22])) # row 0 base (S_LOOP_CTR=0)
for i in range(7): k.emit(v_or_b32_e32(v[10 + i if i < 2 else 12 + i], 16 * (i + 1), v[22])) # rows 1-7
# Extract sign bit of workgroup_x (always 0 for valid workgroups, used for masking)
k.emit(s_bfe_i32(s[S_LOOP_BOUND], s[S_WORKGROUP_X], 0x10018))
k.emit(v_and_b32_e32(v[9], ADDR_MASK, v[9]))
k.emit(s_lshr_b32(s[S_LOOP_BOUND], s[S_LOOP_BOUND], 25))
# Compute masked row offsets for bank conflict avoidance pattern
# Pattern: v[row] = row_val - (row_val & ADDR_MASK) extracts lower bits
k.emit(v_add_nc_u32_e32(v[19], s[S_LOOP_CTR], v[10]))
k.emit(v_add_nc_u32_e32(v[8], s[S_LOOP_BOUND], v[1])) # A-tile base computation
for d, r in zip([20, 21, 32, 33, 34, 35], [11, 14, 15, 16, 17, 18]):
k.emit(v_add_nc_u32_e32(v[d], s[S_LOOP_CTR], v[r]))
k.emit(v_and_b32_e32(v[8], ADDR_MASK, v[8]))
k.emit(v_sub_nc_u32_e32(v[9], v[22], v[9])) # row 0 swizzle offset
for d, s_ in zip([19, 20, 21, 22, 32, 33, 34], [20, 21, 22, 32, 33, 34, 35]):
k.emit(v_and_b32_e32(v[d], ADDR_MASK, v[s_]))
k.emit(v_sub_nc_u32_e32(v[8], v[1], v[8])) # A-tile swizzle
# Apply swizzle offsets and scale to byte offsets
k.emit(v_lshlrev_b32_e32(v[9], 2, v[9])) # row 0 offset * 4
for r, t in zip([10, 11, 14, 15, 16, 17, 18], [19, 20, 21, 22, 32, 33, 34]):
k.emit(v_sub_nc_u32_e32(v[r], v[r], v[t])) # rows 1-7 swizzle
k.emit(v_bfe_u32(v[2], v[0], 3, 2)) # v[2] = (lane_id >> 3) & 3
k.emit(v_lshlrev_b32_e32(v[8], 2, v[8])) # A-tile base * 4
# Compute B-tile base address: LDS_A_STRIDE * (lane_id % 8) + row0_offset
k.emit(v_mad_u32_u24(v[V_LDS_B_ADDR], LDS_A_STRIDE, v[V_LANE_ID_MOD8], v[9]))
# Scale row offsets 1-7 to byte offsets (row 0 already in v[9])
for d, r in zip([9, 10, 11, 14, 15, 16, 17], [10, 11, 14, 15, 16, 17, 18]):
k.emit(v_lshlrev_b32_e32(v[d], 2, v[r]))
k.emit(v_lshlrev_b32_e32(v[V_LANE_DIV8_X4], 2, v[2]))
k.emit(v_add_nc_u32_e32(v[8], 0x80, v[8])) # A-tile initial store base + 128
# Store initial tile data to LDS
k.waitcnt(vm=0)
for i, (d0, d1) in enumerate([(0,1), (2,3), (4,5), (11,12)]):
k.emit(ds_store_2addr_stride64_b32(addr=v[8], data0=v[INIT_TILE_LOADS[d0][0]], data1=v[INIT_TILE_LOADS[d1][0]], offset0=16+i*4, offset1=18+i*4))
# B stores: single base with offsets 0,64,128,192,256,320,384,448
for i, idx in enumerate([6,7,8,9,10,13,14,15]):
offset = i * 64
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[INIT_TILE_LOADS[idx][0]], offset0=offset & 0xFF, offset1=offset >> 8))
k.waitcnt(lgkm=0)
k.barrier()
# ===========================================================================
# INIT: Compute LDS base addresses, then zero accumulators
# ===========================================================================
# v[3] = v[1] & 0x7F (lower 7 bits) since S_LOOP_BOUND=0 for valid workgroups
k.emit(v_lshlrev_b32_e32(v[2], 4, v[2]))
k.emit(v_add_nc_u32_e32(v[3], s[S_LOOP_BOUND], v[1]))
k.emit(v_and_b32_e32(v[3], ADDR_MASK, v[3]))
k.emit(v_sub_nc_u32_e32(v[3], v[1], v[3]))
k.emit(v_lshl_or_b32(v[V_LDS_B_BASE], v[V_LANE_ID_MOD8], 4, LDS_BASE_OFFSET))
k.emit(v_lshl_add_u32(v[V_LDS_A_ADDR], v[3], 2, LDS_BASE_OFFSET))
k.emit(v_lshlrev_b32_e32(v[3], 2, v[0]))
k.emit(v_and_or_b32(v[V_LDS_A_BASE], 0x180, v[3], v[2]))
# Zero all 128 accumulators using VOPD dual moves (64 instructions instead of 128)
for i in range(0, len(OUT_REGS), 2):
k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[OUT_REGS[i]], vdsty=v[OUT_REGS[i+1]], srcx0=0, srcy0=0))
k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
k.emit(s_add_u32(s[S_A_PTR[0]], s[S_A_PTR[0]], 32))
k.emit(s_addc_u32(s[S_A_PTR[1]], s[S_A_PTR[1]], 0))
# S_LOOP_CTR is already 0 from prologue initialization
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_ENTRY')
# ===========================================================================
# MAIN GEMM LOOP
# ===========================================================================
NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
k.label('LOOP_INC')
k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
k.emit(s_cbranch_scc1(simm16=0)); k.branch_to('EPILOGUE')
k.label('LOOP_ENTRY')
k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0)) # s_cselect doesn't modify SCC
k.emit(s_cbranch_scc0(simm16=0)); k.branch_to('SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
# Advance prefetch pointers
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], 0x20000, v[V_GLOBAL_B_ADDR]))
k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], 0x20, v[V_GLOBAL_A_ADDR]))
if not NO_GLOBAL:
for vdst, saddr_lo in INIT_PREFETCH:
k.global_load(vdst, V_GLOBAL_B_ADDR, saddr_lo)
k.label('SKIP_PREFETCH')
# 8 inner loop iterations
for iter in range(8):
# Load A tile (4 pairs) and B tile (8 pairs) from LDS
if not NO_DS:
k.emit(s_clause(simm16=11)) # 12 loads total: 4 A + 8 B
# A tile: 4 ds_load_b64
for i, vdst in enumerate(V_A_TILE_REGS):
a_off = (i & 1) * 8 + (i >> 1) * 64 + iter * LDS_A_STRIDE
k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_A_BASE], offset0=a_off & 0xFF, offset1=a_off >> 8))
# B tile: 8 ds_load_b64
for i, vdst in enumerate(V_B_TILE_REGS):
b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 + iter * LDS_B_STRIDE
k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_B_BASE], offset0=b_off & 0xFF, offset1=b_off >> 8))
k.waitcnt(lgkm=0)
# 64 dual FMACs
k.emit(s_clause(simm16=63))
for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
# Issue global prefetch AFTER FMACs (first 6 iterations only)
if iter < 6 and not NO_GLOBAL:
vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
k.global_load(vdst1, addr, slo1)
k.global_load(vdst2, addr, slo2)
k.emit(s_and_not1_b32(VCC_LO, EXEC_LO, s[S_PREFETCH_FLAG]))
k.waitcnt(vm=0)
k.barrier()
k.emit(s_cbranch_vccnz(simm16=0)); k.branch_to('LOOP_INC')
# Store prefetched data to LDS
# NOTE: Register naming reflects LDS tile organization, not source matrix:
# V_LDS_A_DATA (v155-162) holds data that goes to LDS A-tile region
# V_LDS_B_DATA (v163-170) holds data that goes to LDS B-tile region
# The data sources are swapped: A-tile receives B matrix rows, B-tile receives A matrix columns
for i in range(4): # A tile: 8 values via 4 stride64 stores
k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
for i in range(8): # B tile: 8 values via 8 scalar stores with 64-byte spacing
offset = i * 64
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
k.waitcnt(lgkm=0)
k.barrier()
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_INC')
# ===========================================================================
# EPILOGUE: Permute and store results
# ===========================================================================
k.label('EPILOGUE')
# Rearrange accumulators from FMAC layout to contiguous output order
for a, b in PERMUTE_SWAPS:
k.emit(v_swap_b32_e32(v[a], v[b]))
# Compute output coordinates: v[V_LANE_ID_MOD8] = col, v[V_OUTPUT_ROW] = row
k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32,
vdstx=v[149], vdsty=v[150], srcx0=v[V_LANE_MOD8_X4], vsrcx1=v[0], srcy0=v[V_LANE_DIV8_X4], vsrcy1=v[0]))
k.emit(v_and_b32_e32(v[0], 0x60, v[0]))
k.emit(v_or_b32_e32(v[V_LANE_ID_MOD8], s[S_TILE_X], v[149]))
k.emit(v_add_nc_u32_e32(v[0], s[S_TILE_Y], v[0]))
k.emit(v_or_b32_e32(v[V_OUTPUT_ROW], v[0], v[150]))
# Precompute row offsets: v[144-147] for rows 0-3, v[148-151] for rows 16-19
for base, row_off in [(144, 0), (148, 16)]:
if row_off: k.emit(v_or_b32_e32(v[1], row_off, v[V_OUTPUT_ROW]))
k.emit(v_mul_lo_u32(v[base], v[1] if row_off else v[V_OUTPUT_ROW], s[S_DIM_N]))
for i in range(3): k.emit(v_add_nc_u32_e32(v[base + 1 + i], s[S_DIM_N], v[base + i]))
k.emit(v_mov_b32_e32(v[V_ADDR_HI_ZERO], 0))
k.emit(s_lshl_b32(s[S_PREFETCH_FLAG], s[S_DIM_N], 2)) # row stride in bytes
# Store 128 output values as 32 groups of 4 (128-bit stores)
# Layout: 2 row halves (0-3, 16-19) x 4 col groups x 4 rows = 32 stores of 4 floats
epilogue_reserved = {V_LANE_ID_MOD8, V_OUTPUT_ROW, V_LANE_MOD8_X4, V_LANE_DIV8_X4, V_ADDR_HI_ZERO}
for i, (row_half, col_off, row_in_group) in enumerate([(rh, co, ri)
for rh in range(2) for co in [0, 32, 64, 96] for ri in range(4)]):
row = row_half * 16 + row_in_group
srcs = OUT_REGS[i*4:(i+1)*4]
# Find temp register for scaled values (must not conflict with reserved regs)
tmp = max(srcs) + 5
while any(r in epilogue_reserved for r in range(tmp, tmp + 4)): tmp += 1
# Copy values to temp regs for output (alpha=1.0 hardcoded, so just move)
for j, src in enumerate(srcs):
k.emit(v_mov_b32_e32(v[tmp + j], v[src]))
# Compute output address
if row_in_group == 0: # first row: compute base address for this column group
if col_off == 0: k.emit(v_mov_b32_e32(v[0], v[V_LANE_ID_MOD8]))
else: k.emit(v_add_nc_u32_e32(v[0], col_off, v[V_LANE_ID_MOD8]))
row_base = 144 + row if row < 4 else 148 + row - 16
k.emit(v_add_nc_u32_e32(v[0], v[row_base], v[0]))
k.emit(v_lshlrev_b32_e32(v[0], 2, v[0]))
k.emit(v_add_co_u32(v[0], VCC_LO, s[S_OUT_PTR[0]], v[0]))
k.emit(v_add_co_ci_u32_e32(v[1], s[S_OUT_PTR[1]], v[V_ADDR_HI_ZERO]))
else: # subsequent rows: just add stride
k.emit(v_add_co_u32(v[0], VCC_LO, s[S_PREFETCH_FLAG], v[0]))
k.emit(v_add_co_ci_u32_e32(v[1], v[1], v[V_ADDR_HI_ZERO]))
k.emit(global_store_b128(addr=v[0:1], data=v[tmp:tmp+3], saddr=RawImm(124)))
k.emit(s_sendmsg(simm16=3))
k.emit(s_endpgm())
return k.to_asm()
# =============================================================================
# Test harness
# =============================================================================
N = getenv("N", 4096)
BLOCK_M, BLOCK_N = 128, 128
THREADS = 128
def test_matmul():
dev = Device[Device.DEFAULT]
print(f"Device arch: {dev.arch}")
if getenv("STOCK", 0):
# Load the stock kernel from amd_seb/kernel8_batched_gmem.s
stock_path = Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s"
asm = stock_path.read_text()
print(f"Loaded stock kernel from {stock_path}")
else:
asm = build_kernel(dev.arch)
if getenv("PRINT_ASM", 0): print(asm)
binary = dev.compiler.compile(asm)
print(f"Compiled! Binary size: {len(binary)} bytes")
rng = np.random.default_rng(42)
a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
b = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
c = Tensor.empty(N, N)
Tensor.realize(a, b, c)
grid, local = (N // BLOCK_N, N // BLOCK_M, 1), (THREADS, 1, 1)
print(f"Grid: {grid}, Local: {local}")
_prg = dev.runtime("kernel", binary)
class AsmRunner(Runner):
def __init__(self):
super().__init__(colored("kernel", "cyan"), Device.DEFAULT, Estimates(ops=N*N*N*2, mem=N*N*4*3))
def __call__(self, rawbufs, var_vals, wait=False):
c_buf, a_buf, b_buf = [x.ensure_allocated()._buf for x in rawbufs]
return _prg(a_buf, b_buf, c_buf, global_size=grid, local_size=local, wait=wait)
ei = ExecItem(None, [c.uop.buffer, a.uop.buffer, b.uop.buffer], prg=AsmRunner())
ets = []
with Context(DEBUG=2):
for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
if getenv("VERIFY", 1):
GlobalCounters.reset()
with Context(DEBUG=2): tc = (a @ b).realize()
with Context(DEBUG=0): err = (c - tc).square().mean().item()
print(f"mean squared error {err}")
if err > 1e-06: raise RuntimeError("matmul is wrong!")
def run_sqtt():
"""Run with SQTT profiling and write trace files."""
import subprocess, os
# Run test_matmul in a subprocess with SQTT enabled from the start (no verify)
env = {**os.environ, "AMD": "1", "SQTT": "1", "CNT": "1", "PROFILE": "1", "PYTHONPATH": ".", "VERIFY": "0"}
result = subprocess.run(
["python", "-c", "from extra.gemm.amd_asm_matmul import test_matmul; test_matmul()"],
capture_output=True, text=True, env=env, timeout=120
)
print(result.stdout)
# Run roc.py to extract trace data
result = subprocess.run(
["python", "extra/sqtt/roc.py", "--profile", "/tmp/profile.pkl.tiny", "--kernel", "kernel"],
capture_output=True, text=True, env={**os.environ, "DEBUG": "5"}, timeout=60
)
output = result.stdout + result.stderr
# Write full output to trace file
with open("/tmp/sqtt_trace.txt", "w") as f:
f.write(output)
print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")
if __name__ == "__main__":
if getenv("ASM", 0): print(build_kernel(Device[Device.DEFAULT].arch))
elif getenv("SQTT", 0): run_sqtt()
else: test_matmul()
+4 -4
View File
@@ -140,11 +140,11 @@ def hand_spec_kernel3():
return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
def test_matmul(sink:UOp, N=N):
def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
rng = np.random.default_rng()
a = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
b = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
hc = Tensor.empty(N, N)
a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
hc = Tensor.empty(N, N, dtype=dtype)
Tensor.realize(a, b, hc)
ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
@@ -1,16 +1,9 @@
.text
.section .text.
.global gemm
.p2align 8
.type gemm,@function
gemm:
// ** global buffers
s_load_dwordx2 s[28:29], s[0:1], 0x0 // C
s_load_dwordx4 s[32:35], s[0:1], 0x8 // A, B
s_load_dwordx2 s[34:35], s[0:1], 0x08 // A
s_load_dwordx2 s[32:33], s[0:1], 0x10 // B
// ** others kernel args
s_load_dword s24, s[0:1], 0x18 // N
s_load_dword s54, s[0:1], 0x1C // num work groups
s_waitcnt lgkmcnt(0)
// "info"
s_mov_b32 s51, 1 // gemm_info = 1
@@ -221,11 +214,11 @@ gemm:
s_mul_hi_u32 s87, s86, s40 // 000000003288: 96572856
s_mul_i32 s86, s86, s40 // 00000000328C: 92562856
s_and_b32 s84, s50, 0x8000 // 000000003290: 8654FF32 00008000
s_cbranch_scc1 label_GSUC_A // 000000003298: BF850003
s_cbranch_scc1 skip_offset_A // 000000003298: BF850003
s_mul_hi_u32 s85, 64, s6 // 00000000329C: 965506C0
s_mul_i32 s84, 64, s6 // 0000000032A0: 925406C0
label_GSUC_A:
skip_offset_A:
s_add_u32 s86, s86, s84 // 000000003330: 80565456
s_addc_u32 s87, s87, s85 // 000000003334: 82575557
s_mov_b64 s[60:61], 1 // 000000003338: BEBC0181
@@ -259,11 +252,11 @@ label_GSUC_A:
s_mul_hi_u32 s87, s86, s42 // 0000000033B4: 96572A56
s_mul_i32 s86, s86, s42 // 0000000033B8: 92562A56
s_and_b32 s84, s50, 0x8000 // 0000000033BC: 8654FF32 00008000
s_cbranch_scc1 label_GSUC_B // 0000000033C4: BF850003
s_cbranch_scc1 skip_offset_B // 0000000033C4: BF850003
s_mul_hi_u32 s85, 64, s6 // 0000000033C8: 965506C0
s_mul_i32 s84, 64, s6 // 0000000033CC: 925406C0
label_GSUC_B:
skip_offset_B:
s_add_u32 s86, s86, s84 // 00000000345C: 80565456
s_addc_u32 s87, s87, s85 // 000000003460: 82575557
s_mov_b64 s[62:63], 1 // 000000003464: BEBE0181
@@ -308,8 +301,6 @@ label_GSUC_B:
s_and_b32 s87, s10, 0xe000 // 0000000035A4: 8657FF0A 0000E000
s_and_b32 s10, s10, 0xff // 0000000035AC: 860AFF0A 000000FF
s_mov_b32 s84, s10 // 0000000035B4: BED4000A
label_beginStaggerUIter:
s_lshl_b32 s85, s84, s86 // 0000000035B8: 8E555654
s_cmp_ge_u32 s13, s85 // 0000000035BC: BF09550D
s_sub_u32 s85, s84, 1 // 0000000035CC: 80D58154
@@ -344,7 +335,7 @@ label_beginStaggerUIter:
s_cselect_b32 s58, s62, -1 // 0000000036A4: 853AC13E
s_add_u32 s51, s51, 2 // 0000000036A8: 80338233
s_cmp_eq_u32 s12, 0 // 0000000036AC: BF06800C
s_cbranch_scc1 label_ShadowInitStart // 0000000036B0: BF850092
s_cbranch_scc1 init_output_buffers // 0000000036B0: BF850092
s_mov_b32 m0, s46 // 0000000036B4: BEFC002E
buffer_load_dwordx4 v0, s[52:55], 0 offen lds // 0000000036B8: E05D1000 800D0000
s_add_u32 m0, m0, 0x1040 // 0000000036C0: 807CFF7C 00001040
@@ -431,7 +422,7 @@ label_beginStaggerUIter:
s_cmp_eq_u32 s63, 0 // 0000000038F4: BF06803F
s_cselect_b32 s58, s62, -1 // 0000000038F8: 853AC13E
label_ShadowInitStart:
init_output_buffers:
s_mov_b64 s[16:17], s[28:29] // 0000000038FC: BE90011C
s_mov_b32 s18, 0x80000000 // 000000003900: BE9200FF 80000000
s_mov_b32 s19, 0x20000 // 000000003908: BE9300FF 00020000
@@ -476,12 +467,10 @@ label_ShadowInitStart:
s_lshl_b64 s[84:85], s[84:85], 2 // 0000000039C4: 8ED48254
s_add_u32 s16, s16, s84 // 0000000039C8: 80105410
s_addc_u32 s17, s17, s85 // 0000000039CC: 82115511
label_NoBranch_T8JHFHKM7BO5OHXW:
s_xor_b32 s46, s48, s46 // 0000000039F0: 882E2E30
s_xor_b32 s47, s49, s47 // 0000000039F4: 882F2F31
s_cmp_eq_u32 s12, 1 // 0000000039F8: BF06810C
s_cbranch_scc1 label_skipPGR2 // 0000000039FC: BF850040
s_cbranch_scc1 after_prefetch // 0000000039FC: BF850040
s_mov_b32 m0, s46 // 000000003A00: BEFC002E
buffer_load_dwordx4 v0, s[52:55], 0 offen lds // 000000003A04: E05D1000 800D0000
s_add_u32 m0, m0, 0x1040 // 000000003A0C: 807CFF7C 00001040
@@ -517,7 +506,7 @@ label_NoBranch_T8JHFHKM7BO5OHXW:
s_xor_b32 s46, s48, s46 // 000000003AF8: 882E2E30
s_xor_b32 s47, s49, s47 // 000000003AFC: 882F2F31
label_skipPGR2:
after_prefetch:
s_waitcnt vmcnt(24) // 000000003B00: BF8C4F78
s_barrier // 000000003B04: BF8A0000
ds_read_b128 v[4:7], v2 // 000000003B08: D9FE0000 04000002
@@ -539,14 +528,12 @@ label_skipPGR2:
ds_read_b128 v[92:95], v3 offset:768 // 000000003B80: D9FE0300 5C000003
ds_read_b128 v[96:99], v3 offset:896 // 000000003B88: D9FE0380 60000003
s_waitcnt lgkmcnt(0) // 000000003B90: BF8CC07F
label_openLoopL:
s_cmp_eq_u32 s12, 1 // 000000003B94: BF06810C
s_cbranch_scc1 label_toPGR1 // 000000003B98: BF8502E5
s_cbranch_scc1 final_compute // 000000003B98: BF8502E5
s_cmp_le_u32 s12, 2 // 000000003B9C: BF0B820C
s_cbranch_scc1 label_LoopEndL // 000000003BA0: BF85019E
s_cbranch_scc1 loop_epilogue // 000000003BA0: BF85019E
label_LoopBeginL:
main_loop:
v_mfma_f32_16x16x32_bf16 a[0:3], v[68:71], v[4:7], a[0:3] // 000000003BA4: D3B58000 04020944
ds_read_b128 v[36:39], v2 offset:64 // 000000003BAC: D9FE0040 24000002
v_mfma_f32_16x16x32_bf16 a[4:7], v[68:71], v[8:11], a[4:7] // 000000003BB4: D3B58004 04121144
@@ -770,9 +757,9 @@ label_LoopBeginL:
s_cmp_eq_i32 s12, 2 // 000000004208: BF00820C
s_waitcnt lgkmcnt(0) // 00000000420C: BF8CC07F
v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004210: D3B580FC 07F28180
s_cbranch_scc0 label_LoopBeginL // 000000004218: BF84FE62
s_cbranch_scc0 main_loop // 000000004218: BF84FE62
label_LoopEndL:
loop_epilogue:
v_mfma_f32_16x16x32_bf16 a[0:3], v[68:71], v[4:7], a[0:3] // 00000000421C: D3B58000 04020944
ds_read_b128 v[36:39], v2 offset:64 // 000000004224: D9FE0040 24000002
v_mfma_f32_16x16x32_bf16 a[4:7], v[68:71], v[8:11], a[4:7] // 00000000422C: D3B58004 04121144
@@ -939,7 +926,7 @@ label_LoopEndL:
v_mfma_f32_16x16x32_bf16 a[248:251], v[128:131], v[60:63], a[248:251]// 000000004720: D3B580F8 07E27980
v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004728: D3B580FC 07F28180
label_toPGR1:
final_compute:
s_and_b32 s8, s50, 0x3fff // 000000004730: 8608FF32 00003FFF
s_and_b32 s84, 0xff, s24 // 000000004750: 865418FF 000000FF
s_add_u32 s85, -1, s14 // 000000004758: 80550EC1
@@ -1095,7 +1082,6 @@ label_toPGR1:
v_mfma_f32_16x16x32_bf16 a[248:251], v[128:131], v[60:63], a[248:251]// 000000004BFC: D3B580F8 07E27980
v_mfma_f32_16x16x32_bf16 a[252:255], v[128:131], v[64:67], a[252:255]// 000000004C04: D3B580FC 07F28180
label_toPGR1end_OptNLL:
v_lshrrev_b32_e32 v4, 6, v134 // 000000004C0C: 20090C86
v_lshrrev_b32_e32 v5, 1, v4 // 000000004C10: 200A0881
v_mul_lo_u32 v5, 16, v5 // 000000004C14: D2850005 00020A90
@@ -1114,7 +1100,6 @@ label_toPGR1end_OptNLL:
s_mul_i32 s8, 0x100, s3 // 000000004C64: 920803FF 00000100
v_add_u32_e32 v1, s8, v1 // 000000004C6C: 68020208
label_GW_B0_E0:
v_add_lshl_u32 v11, v3, v0, 1 // 000000004C70: D1FE000B 02060103
v_accvgpr_read_b32 v16, a0 // 000000004C78: D3D84010 18000100
v_accvgpr_read_b32 v17, a4 // 000000004C80: D3D84011 18000104
@@ -1633,86 +1618,4 @@ label_GW_B0_E0:
s_addc_u32 s17, s17, 0 // 000000005B14: 82118011
buffer_store_dwordx4 v[40:43], v11, s[16:19], 0 offen nt // 000000005B18: E07E1000 8004280B
s_nop 0 // 000000005B20: BF800000
end:
s_endpgm // 00000001F5D0: BF810000
.section .rodata,"a",@progbits
.p2align 6, 0x0
.amdhsa_kernel gemm
# ---- basic memory requirements ----
.amdhsa_group_segment_fixed_size 133120
.amdhsa_private_segment_fixed_size 0
.amdhsa_kernarg_size 32
# ---- register usage (RSRC1) ----
.amdhsa_next_free_vgpr 504
.amdhsa_next_free_sgpr 96
# ---- workgroup / workitem IDs (RSRC2) ----
.amdhsa_system_sgpr_workgroup_id_x 1
.amdhsa_system_sgpr_workgroup_id_y 1
.amdhsa_system_sgpr_workgroup_id_z 1
# ---- user SGPR enables (descriptor bits >448) ----
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_user_sgpr_count 2
.amdhsa_user_sgpr_kernarg_preload_length 0
.amdhsa_user_sgpr_kernarg_preload_offset 0
# ---- gfx90a / gfx940 specific (RSRC3) ----
.amdhsa_accum_offset 248
.amdhsa_uses_dynamic_stack 0
.amdhsa_tg_split 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.kernels:
- .args:
- .address_space: global
.name: C
.offset: 0
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .address_space: global
.name: B
.offset: 8
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .address_space: global
.name: A
.offset: 16
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: sz
.offset: 24
.size: 4
.value_kind: by_value
.value_type: u32
- .name: num_wg
.offset: 28
.size: 4
.value_kind: by_value
.value_type: u32
.group_segment_fixed_size: 133120
.kernarg_segment_align: 8
.kernarg_segment_size: 32
.max_flat_workgroup_size: 256
.name: gemm
.private_segment_fixed_size: 0
.sgpr_count: 88
.sgpr_spill_count: 0
.symbol: gemm.kd
.vgpr_count: 248
.vgpr_spill_count: 0
.wavefront_size: 64
amdhsa.version:
- 1
- 0
...
.end_amdgpu_metadata
+78
View File
@@ -0,0 +1,78 @@
.text
.section .text.
.global gemm
.p2align 8
.type gemm,@function
gemm:
INSTRUCTIONS
.section .rodata,"a",@progbits
.p2align 6, 0x0
.amdhsa_kernel gemm
# basic memory requirements
.amdhsa_group_segment_fixed_size 133120
.amdhsa_private_segment_fixed_size 0
.amdhsa_kernarg_size 28
# register usage (RSRC1)
.amdhsa_next_free_vgpr 504
.amdhsa_next_free_sgpr 96
# workgroup / workitem IDs (RSRC2)
.amdhsa_system_sgpr_workgroup_id_x 1
.amdhsa_system_sgpr_workgroup_id_y 1
.amdhsa_system_sgpr_workgroup_id_z 1
# user SGPRs, we only specify the kernel args ptr in s[0:1]
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_user_sgpr_count 2
.amdhsa_user_sgpr_kernarg_preload_length 0
.amdhsa_user_sgpr_kernarg_preload_offset 0
# gfx90a / gfx940 specifics (RSRC3)
.amdhsa_accum_offset 248
.amdhsa_uses_dynamic_stack 0
.amdhsa_tg_split 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.kernels:
- .name: gemm
.symbol: gemm.kd
.args:
- .name: C
.address_space: global
.offset: 0
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: B
.address_space: global
.offset: 8
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: A
.address_space: global
.offset: 16
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: sz
.offset: 24
.size: 4
.value_kind: by_value
.value_type: u32
.group_segment_fixed_size: 133120
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.kernarg_segment_size: 28
.max_flat_workgroup_size: 256
.sgpr_count: 88
.sgpr_spill_count: 0
.vgpr_count: 248
.vgpr_spill_count: 0
.wavefront_size: 64
amdhsa.version:
- 1
- 0
...
.end_amdgpu_metadata
@@ -2,10 +2,8 @@
# VIZ=2 to profile
import pathlib
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.engine.realize import ExecItem, CompiledRunner
from tinygrad.renderer import ProgramSpec
from tinygrad.uop.ops import track_rewrites, UOp
from tinygrad.helpers import TracingKey, getenv
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.helpers import getenv
fp = pathlib.Path(__file__).parent/"gemm.s"
@@ -23,8 +21,8 @@ import torch
torch.manual_seed(0)
A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
Bt = B.t().contiguous() # transpose B for the baseline gemm
C_torch = A@Bt
Bt = B.t().contiguous() # transpose B for the asm gemm
C_torch = A@B
# ** copy buffers to AMD
@@ -33,31 +31,32 @@ C_torch = A@Bt
def from_torch(t:torch.Tensor) -> Tensor:
return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
C_tiny = Tensor.matmul(from_torch(A), from_torch(Bt), dtype=dtypes.float32).cast(dtypes.bfloat16)
C_tiny = from_torch(A) @ from_torch(B)
C_asm = Tensor.empty_like(C_tiny)
C_asm.uop.buffer.allocate()
# ** assembly custom kernel
def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
lidx = UOp.special(THREADS_PER_WG, "lidx0")
gidx = UOp.special(NUM_WG, "gidx0")
src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
sz = UOp.variable("SZ", 256, 8192)
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)), arg=())
C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
# ** run gemms
# baseline tinygrad
sched = C_tiny.schedule()
assert len(sched) == 1
eis:list[ExecItem] = [sched[-1].lower()]
ast = sched[-1].ast
# assembly gemm
@track_rewrites(name=lambda ret: TracingKey(ret.name, (ret.function_name,), ret))
def get_asm_prg() -> ProgramSpec:
src = fp.read_text()
lib = Device[Device.DEFAULT].compiler.compile(src)
return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[NUM_WG, 1, 1], local_size=[THREADS_PER_WG, 1, 1],
globals=[0, 1, 2], vars=[UOp.variable("SZ", 256, 8192), UOp.variable("NUM_WG", 1, 1024)])
eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], fixedvars={"SZ":N, "NUM_WG":NUM_WG},
prg=CompiledRunner(get_asm_prg())))
sched = Tensor.schedule(C_tiny, C_asm)
eis = [si.lower() for si in sched]
with Context(DEBUG=2):
for ei in eis:
et = ei.run(wait=True)
et = ei.run({"SZ":N}, wait=True)
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
# ** correctness
File diff suppressed because it is too large Load Diff
+76
View File
@@ -0,0 +1,76 @@
.text
.section .text.
.global gemm
.p2align 8
.type gemm,@function
gemm:
INSTRUCTIONS
.section .rodata,"a",@progbits
.p2align 6, 0x0
.amdhsa_kernel gemm
# basic memory requirements
.amdhsa_group_segment_fixed_size 30336
.amdhsa_private_segment_fixed_size 0
.amdhsa_kernarg_size 32
# register usage (RSRC1)
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 100
# workgroup / workitem IDs (RSRC2)
.amdhsa_system_sgpr_workgroup_id_x 1
.amdhsa_system_sgpr_workgroup_id_y 1
.amdhsa_system_sgpr_workgroup_id_z 1
# user SGPRs: kernarg ptr in s[0:1]
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_user_sgpr_count 2
# gfx10+ / gfx11 specifics (RSRC1[29..31])
.amdhsa_wavefront_size32 1
.amdhsa_workgroup_processor_mode 1
.amdhsa_memory_ordered 1
.amdhsa_forward_progress 1
# misc for gfx11
.amdhsa_dx10_clamp 1
.amdhsa_ieee_mode 1
.amdhsa_uses_dynamic_stack 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.kernels:
- .args:
- .address_space: generic
.name: C
.offset: 0
.size: 8
.value_kind: global_buffer
.value_type: f16
- .address_space: generic
.name: A
.offset: 8
.size: 8
.value_kind: global_buffer
.value_type: f16
- .address_space: generic
.name: B
.offset: 16
.size: 8
.value_kind: global_buffer
.value_type: f16
.group_segment_fixed_size: 30336
.kernarg_segment_align: 8
.kernarg_segment_size: 32
.max_flat_workgroup_size: 128
.name: gemm
.private_segment_fixed_size: 0
.sgpr_count: 70
.sgpr_spill_count: 0
.symbol: gemm.kd
.vgpr_count: 256
.vgpr_spill_count: 0
.wavefront_size: 32
amdhsa.version:
- 1
- 1
...
.end_amdgpu_metadata
+30
View File
@@ -0,0 +1,30 @@
import math, pathlib
from tinygrad import Device, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from extra.gemm.amd_uop_matmul import test_matmul
N = 4096
TN = 96
THREADS_PER_WG = 128
NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
dname:str = Device.DEFAULT
template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
def asm_kernel() -> UOp:
lidx = UOp.special(THREADS_PER_WG, "lidx0")
gidx = UOp.special(NUM_WG, "gidx0")
a = UOp.placeholder((N*N,), dtypes.half, slot=1)
b = UOp.placeholder((N*N,), dtypes.half, slot=2)
c = UOp.placeholder((N*N,), dtypes.half, slot=0)
src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)), arg=())
if __name__ == "__main__":
test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
+5 -5
View File
@@ -1,12 +1,12 @@
# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
# unpack the complete kernel descriptor of an amdgpu ELF
# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
import struct, pathlib
import struct, pathlib, sys
from tinygrad.runtime.support.elf import elf_loader
def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
with open(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
with open(sys.argv[1], "rb") as f:
lib = f.read()
image, sections, relocs = elf_loader(lib)
@@ -49,7 +49,7 @@ print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
# rsrc 3
# rsrc 3 (gfx950)
accum_offset_raw = bits(pgm_rsrc3, 0, 5)
assert_zero(pgm_rsrc3, 6, 15)
@@ -169,10 +169,10 @@ assert_zero(desc, 458, 459)
uses_dynamic_stack = bits(desc, 459, 460)
print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
# gfx950 only
assert_zero(desc, 460, 463)
kernarg_preload_spec_length = bits(desc, 464, 470)
print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
kernarg_preload_spec_offset = bits(desc, 471, 479)
print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
+47 -30
View File
@@ -1,8 +1,42 @@
import argparse, os, hashlib
from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
import argparse, os, hashlib, functools
from typing import Iterator, Callable
from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch, ceildiv
from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
from tinygrad import Tensor, dtypes, Device, Variable, TinyJit
# rounds up hevc input data to 32 bytes, so more optimal kernels can be generated
HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
@functools.cache
def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
if outbuf is not None: outbuf.assign(x).realize()
return x.realize()
return TinyJit(hevc_decode_frame)
def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
history:list[Tensor]|None=None, preallocated_outputs:list[Tensor]|None=None, warmup=False) -> Iterator[Tensor]:
out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
v_pos = Variable("pos", 0, max_hist + 1)
v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
v_sz = Variable("sz", 1, ceildiv(hevc_tensor.numel(), HEVC_ROUNDUP))
v_i = Variable("i", 0, len(frame_info)-1)
decode_jit = _hevc_jitted_decoder(out_image_size, max_hist, preallocated_outputs is not None)
history = history or [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
assert len(history) == max_hist, f"history length {len(history)} does not match max_hist {max_hist}"
for i, (offset, sz, frame_pos, _, is_hist) in enumerate(frame_info):
history = history[-max_hist:] if max_hist > 0 else []
img = decode_jit(v_pos.bind(frame_pos), hevc_tensor, v_offset.bind(offset), v_sz.bind(ceildiv(sz, HEVC_ROUNDUP)),
opaque, v_i.bind(i), *history, outbuf=preallocated_outputs[i] if preallocated_outputs else None)
res = preallocated_outputs[i] if preallocated_outputs else img.clone().realize()
if is_hist: history.append(res)
yield res
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--input_file", type=str, default="")
@@ -19,49 +53,32 @@ if __name__ == "__main__":
dat_hash = hashlib.md5(dat).hexdigest()
with Timing("prep infos: "):
dat_nv = hevc_tensor.to("NV")
opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]
# move all needed data to gpu
#all_slices = []
with Timing("copy to gpu: "):
opaque_nv = opaque.to("NV").contiguous().realize()
hevc_tensor = hevc_tensor.to("NV")
out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
max_hist = max(history_sz for _, _, _, history_sz, _ in frame_info)
# define variables
v_pos = Variable("pos", 0, max_hist + 1)
v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
v_sz = Variable("sz", 0, hevc_tensor.numel())
v_i = Variable("i", 0, len(frame_info)-1)
# preallocate output/hist buffers
max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
hist = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
out_images = [Tensor.zeros(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(len(frame_info))]
@TinyJit
def decode_jit(pos:Variable, src:Tensor, data:Tensor, *hist:Tensor):
return src.decode_hevc_frame(pos, out_image_size, data, hist).realize()
# warmup decode
_ = list(hevc_decode(hevc_tensor, opaque_nv, frame_info[:3], luma_h, luma_w, history=hist, preallocated_outputs=out_images))
Device.default.synchronize()
# warm up
history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
for i in range(3):
hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(frame_info[0][0])), bound_offset+v_sz.bind(frame_info[0][1])),))
decode_jit(v_pos.bind(0), hevc_frame, opaque_nv[v_i.bind(0)], *history)
out_images = []
# decode all frames using the iterator
with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
for i, (offset, sz, frame_pos, history_sz, is_hist) in enumerate(frame_info):
history = history[-max_hist:] if max_hist > 0 else []
# TODO: this shrink should work as a slice
hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(offset)), bound_offset+v_sz.bind(sz)),))
outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
out_images.append(outimg)
if is_hist: history.append(outimg)
images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
Device.default.synchronize()
# validation
if getenv("VALIDATE", 0):
import pickle
if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
@@ -70,7 +87,7 @@ if __name__ == "__main__":
else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
else: import cv2
for i, img in tqdm(enumerate(out_images)):
for i, img in tqdm(enumerate(images)):
if getenv("VALIDATE", 0):
if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
img = untile_nv12(img, h, w, luma_w, chroma_off).realize()
+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
+1 -1
View File
@@ -129,7 +129,7 @@ class LSTM:
return self.do_step(x_, hc_)
if hc is None:
hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False)
hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False).contiguous().realize()
output = None
for t in range(x.shape[0]):
+28 -2
View File
@@ -1,9 +1,13 @@
import unittest
import numpy as np
from tinygrad.helpers import BEAM, Timing, CI, Context
from tinygrad import Variable, Tensor
from tinygrad.helpers import BEAM, Timing, CI, prod
from tinygrad import Variable, Device, Tensor
from tinygrad.nn import Conv2d
from tinygrad.uop.ops import AxisType
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.codegen.opt.postrange import Scheduler
from tinygrad.codegen.opt.search import get_kernel_actions
def rand(*shape):
return Tensor(np.random.rand(*shape).astype(np.float32))
@@ -75,5 +79,27 @@ class TestBeamSearch(unittest.TestCase):
a = (a + a) * a
a.realize()
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
def test_tc_up(self):
tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
size = max(tc.dims[0], tc.dims[1]) * 8
a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast
s = Scheduler(ast, Device[Device.DEFAULT].renderer)
s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
actions = get_kernel_actions(s, include_0=False, max_up=int(up))
upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
def test_max_up(self):
a = Tensor.rand(16, 16)
ast = a.schedule()[-1].ast
s = Scheduler(ast, Device[Device.DEFAULT].renderer)
for max_up in (2, 4):
actions = get_kernel_actions(s, include_0=False, max_up=max_up)
for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0
if __name__ == '__main__':
unittest.main()
+9 -15
View File
@@ -3,21 +3,18 @@
import numpy as np
import unittest
import subprocess, struct, math, textwrap
from tinygrad import Tensor, dtypes, Device, UOp
from tinygrad.uop.ops import Ops
import subprocess, struct, math, functools
from tinygrad import Tensor, dtypes, Device
from tinygrad.helpers import getenv
from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
from tinygrad.renderer import ProgramSpec
from tinygrad.engine.realize import CompiledRunner
from extra.assembly.amd.autogen.rdna3 import *
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.asm import waitcnt
from test.testextra.test_cfg_viz import template
from test.testextra.test_cfg_viz import asm_kernel
def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
out = Tensor([0]*n_threads, dtype=dtypes.uint32).realize()
src = "\n".join(inst.disasm() for inst in [
insts = [
s_load_b64(s[0:1], s[0:1], NULL),
*asm,
v_lshlrev_b32_e32(v[0], 2, v[0]),
@@ -25,12 +22,9 @@ def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
#global_store_b32(v[0], v[1], s[0:1]),
global_store_b32(addr=v[0], data=vdst, saddr=s[0:1]),
s_endpgm()
])
prg = ProgramSpec("test", template.replace("fn_name", "test").replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
global_size=[1, 1, 1], local_size=[n_threads, 1, 1], globals=[0])
car = CompiledRunner(prg)
if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
car([out.uop.buffer], {}, wait=True)
]
out = Tensor.custom_kernel(out, fxn=functools.partial(asm_kernel, name="test", insts=insts, device=out.device, n_threads=n_threads))[0]
out.realize()
return out.tolist()
def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
+1 -1
View File
@@ -8,7 +8,7 @@ SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get p
`SQTT_ITRACE_SE_MASK=X` to select for which shader engines instruction tracing will be enabled, -1 is all, 0 is none (instruction tracing disabled), >0 is
bitfield/mask for SEs to enable instruction tracing on. Masking shader engines will give smaller file sizes at a cost of less hits and kernels that
don't have any wavefront on first simd of shdaer engine with instruction tracing enabled will not have instruction timings.
don't have any wavefront on first simd of shader engine with instruction tracing enabled will not have instruction timings.
The default is 2 (second shader engine only), only one for file size reasons, second instead of first because dispatch starts from it so there is
greater chance that kernels with small global size will have instruction tracing data.
+6 -1
View File
@@ -113,12 +113,17 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[st
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
exc:Exception|None = None
def worker():
nonlocal exc
try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
except AttributeError as e:
raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install")
exc.__cause__ = e
(t:=threading.Thread(target=worker, daemon=True)).start()
t.join()
if exc is not None:
raise exc
return ROCParseCtx
def print_data(data:dict) -> None:
+103 -29
View File
@@ -45,7 +45,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
@@ -69,9 +68,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
scale_vec = warp.ones(scale_vec)
# load q tile
q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
q_reg = warp.copy(q_reg, q_reg_fl)
q_reg = warp.load(q_reg, q, (), (batch, q_seq, head, 0), axis=1)
q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
for kv_idx in ker.range(N // KV_BLOCK_SIZE):
@@ -85,6 +82,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
att_block = warp.zero(att_block.after(kv_idx))
k_reg_transposed = warp.transpose(k_reg_transposed, k_reg)
att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
# apply attention mask
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_seq, kv_idx), axis=2)
@@ -217,11 +215,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
return ker.finish()
def custom_backward_kv(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
def custom_backward_k(dku:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_k", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dk, dv, do, q, k, v, mask = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
dk, do, q, k, v, mask = GL(dku, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
head_kv = ker.blockIdx_x
@@ -242,7 +240,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dk_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
@@ -256,6 +253,98 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
delta_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
dk_reg = warp.zero(dk_reg)
# load kv tile
k_reg = warp.load(k_reg, k, (), (batch, kv_seq, head_kv, 0), axis=1)
k_reg_t = warp.transpose(k_reg_t, k_reg)
v_reg = warp.load(v_reg, v, (), (batch, kv_seq, head_kv, 0), axis=1)
for q_idx in ker.range(N // Q_BLOCK_SIZE):
for g in ker.range(GROUP_SIZE):
head_q = head_kv * GROUP_SIZE + g
# load q and do
q_smem = warp.load(q_smem, q, (), (batch, q_idx, head_q, 0), axis=1)
do_smem = warp.load(do_smem, do, (), (batch, q_idx, head_q, 0), axis=1)
q_reg = warp.load(q_reg, q_smem)
q_reg_t = warp.transpose(q_reg_t, q_reg)
q_reg_col = warp.load(q_reg_col, q_smem)
do_reg = warp.load(do_reg, do_smem)
do_reg_col = warp.load(do_reg_col, do_smem)
# load l_vec and delta_vec
l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head_q, 0, q_idx), axis=2)
l_vec_reg *= 1.0 / math.log(2)
delta_vec_reg = warp.load(delta_vec_reg, delta_vec, (), (batch, head_q, 0, q_idx), axis=2)
# mma qk^t
att_block = warp.zero(att_block.after(g))
att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
# apply attention mask
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
att_block += mask_reg_transposed
att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
att_block -= l_vec_reg
att_block = att_block.exp2()
dp_block = warp.zero(dp_block.after(g, q_idx))
dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
dp_block -= delta_vec_reg
att_block *= dp_block
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dk_reg = warp.mma_AB(dk_reg, att_block_row, q_reg_col)
dk_reg = ker.endrange(2)
dk_reg *= 1.0 / math.sqrt(D)
dk = warp.store(dk, dk_reg, (batch, kv_seq, head_kv, 0), axis=1)
return ker.finish()
def custom_backward_v(dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_v", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dv, do, q, k, v, mask = GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
head_kv = ker.blockIdx_x
batch = ker.blockIdx_z
kv_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
q_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
att_smem = ker.st((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
q_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
k_reg_t = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
att_block_transposed = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
att_block_row = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
l_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
delta_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
dv_reg = warp.zero(dv_reg)
# load kv tile
@@ -299,27 +388,12 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dv_reg_ = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
dv_reg = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
dv_reg = ker.endrange(2)
dp_block = warp.zero(dp_block.after(g, q_idx, dv_reg_))
dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
dp_block -= delta_vec_reg
att_block *= dp_block
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dk_reg = warp.mma_AB(dk_reg, att_block_row, q_reg_col)
dk_reg = ker.endrange(2)
dv_reg = dv_reg.after(dk_reg)
dk_reg *= 1.0 / math.sqrt(D)
dk = warp.store(dk, dk_reg, (batch, kv_seq, head_kv, 0), axis=1)
dv = warp.store(dv, dv_reg, (batch, kv_seq, head_kv, 0), axis=1)
return ker.finish(2)
return ker.finish()
if is_causal:
if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
@@ -339,11 +413,11 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
grad_v = Tensor.empty_like(v := Tensor(kernel.src[4]))
mask = Tensor(kernel.src[5])
delta_vec = (grad * attn).sum(-1).transpose(1, 2).unsqueeze(-2).detach()
print(l_vec.shape, delta_vec.shape, grad.shape, attn.shape, grad_q.shape, grad_k.shape, grad_v.shape)
delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
grad_q = Tensor.custom_kernel(grad_q, grad, q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
grad_k = Tensor.custom_kernel(grad_k, grad, q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_k)[0]
grad_v = Tensor.custom_kernel(grad_v, grad, q, k, v, mask, l_vec, delta_vec, fxn=custom_backward_v)[0]
return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
+38 -30
View File
@@ -24,13 +24,11 @@ class Group:
# ops that only work on a single warp
clear_rid = 1000
def clear(self, reg:ALL_TILES, value:float=0):
reg = cast(UOp, reg)
assert self.warps == 1
rngs_for_shape = tuple(UOp.range(dim, Group.clear_rid + i) for i, dim in enumerate(reg.shape))
Group.clear_rid += len(reg.shape)
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in reg.shape)
reg_store = reg[*rngs_for_shape].store(value).end(*rngs_for_shape)
@@ -41,14 +39,12 @@ class Group:
def ones(self, reg:ALL_TILES): return self.clear(reg, 1)
def neg_inf(self, reg:ALL_TILES): return self.clear(reg, -math.inf)
copy_rid = 300
def copy(self, dst:ALL_TILES, src:ALL_TILES):
dst, src = cast(UOp, dst), cast(UOp, src)
assert self.warps == 1
assert dst.shape == src.shape
rngs_for_shape = tuple(UOp.range(dim, Group.copy_rid + i) for i, dim in enumerate(dst.shape))
Group.copy_rid += len(dst.shape)
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in dst.shape)
src_load = src[*rngs_for_shape]
if src.dtype.base != dst.dtype.base:
@@ -219,8 +215,7 @@ class Group:
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
for height in self.ker.range(src.shape[-3], track=False):
i = UOp.range(red_reg.size, Group.clear_rid)
Group.clear_rid += 1
i = self.ker.raw_range(red_reg.size)
red_reg = red_reg.after(height, *[tkr._rng for tkr in self.ker.range_stack])
reg_store = red_reg.flatten()[i].store(init_value).end(i)
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
@@ -254,8 +249,7 @@ class Group:
red_reg = self.ker.alloc((1,), src.dtype.base, AddrSpace.REG)
for width in self.ker.range(src.shape[-2], track=False):
i = UOp.range(red_reg.size, Group.clear_rid)
Group.clear_rid += 1
i = self.ker.raw_range(red_reg.size)
red_reg = red_reg.after(width, *[tkr._rng for tkr in self.ker.range_stack])
reg_store = red_reg.flatten()[i].store(init_value).end(i)
red_reg = red_reg.after(reg_store).reshape(red_reg.shape)
@@ -302,9 +296,20 @@ class Group:
row = laneid % rt.base_shape.rows
col = rt.base_shape.stride * (laneid // rt.base_shape.rows) + inner
sheight = height
swidth = width
if len(idxs) == 2:
row_idx = idxs[0] * dst.shape[-3] * rt.base_shape.rows
col_idx = idxs[1] * dst.shape[-2] * rt.base_shape.cols
row += row_idx % st.base_shape.rows
col += col_idx % st.base_shape.cols
sheight += row_idx // st.base_shape.rows
swidth += col_idx // st.base_shape.cols
srow, scol = cast(ST, src).swizzle(row, col)
src_load = src[*idxs[:-2], height, width, srow, scol]
src_load = src[*idxs[:-2], sheight, swidth, srow, scol]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
dst_store = dst[*dst_idxs, height, width, inner].store(src_load)
@@ -318,28 +323,31 @@ class Group:
idxs = tuple(idx * st.cols if i == 3 else idx for i, idx in enumerate(idxs))
src_i = ((idxs[0] * src.shape[-3] + idxs[1]) * src.shape[-2] + idxs[2]) * src.shape[-1] + idxs[3]
for height in self.ker.range(dst.shape[-4], track=False):
for width in self.ker.range(dst.shape[-3], track=False):
elements_per_thread = st.base_shape.elements_per_thread
memcpy_per_row = st.base_shape.cols // elements_per_thread
total_calls = st.base_shape.num_elements // (self.group_threads * elements_per_thread)
elements_per_thread = st.base_shape.elements_per_thread
memcpy_per_row = st.cols // elements_per_thread
total_calls = (dst.shape[-4] * dst.shape[-3] * st.base_shape.num_elements) // (self.group_threads * elements_per_thread)
for outer in self.ker.range(total_calls, track=False):
for inner in self.ker.range(elements_per_thread, axis_type=AxisType.UPCAST, track=False):
load_idx = outer * self.group_threads + self.laneid
row = load_idx // memcpy_per_row
col = (load_idx * elements_per_thread) % st.base_shape.cols + inner
for outer in self.ker.range(total_calls, track=False):
for inner in self.ker.range(elements_per_thread, axis_type=AxisType.UPCAST, track=False):
load_idx = outer * self.group_threads + self.laneid
row = load_idx // memcpy_per_row
col = (load_idx * elements_per_thread) % st.cols + inner
height = row // st.base_shape.rows
width = col // st.base_shape.cols
srow, scol = cast(ST, dst).swizzle(row, col)
row = row % st.base_shape.rows
col = col % st.base_shape.cols
src_i += height * st.base_shape.rows * row_stride + width * st.base_shape.cols
src_i += row * row_stride + col
srow, scol = cast(ST, dst).swizzle(row, col)
src_load = srcf[src_i]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
dst_store = dst_store.end(height, width, outer, inner).barrier()
src_i += height * st.base_shape.rows * row_stride + width * st.base_shape.cols
src_i += row * row_stride + col
src_load = srcf[src_i]
if src.dtype.base != dst.dtype.base:
src_load = src_load.cast(dst.dtype.base)
dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
dst_store = dst_store.end(height, width, outer, inner).barrier()
elif dst_dtype.addrspace == AddrSpace.REG and src_dtype.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
srcf = src.flatten()
row_stride = prod(src.shape[axis+1:])
@@ -471,7 +479,7 @@ class Group:
idxs = tuple(idx * rv.length if i == 3 else idx for i, idx in enumerate(idxs))
dst_i = ((idxs[0] * dst.shape[-3] + idxs[1]) * dst.shape[-2] + idxs[2]) * dst.shape[-1] + idxs[3]
for outer in self.ker.range(src.shape[-2]):
for outer in self.ker.range(src.shape[-2], track=False):
dst_i += outer * reductions + (laneid % reductions)
src_load = src[outer, 0]
+5
View File
@@ -55,6 +55,11 @@ class Kernel(AbstractContextManager):
if track: self.range_stack.append(rng)
return rng
def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
rng = UOp.range(end, self.range_id, axis_type=axis_type)
self.range_id += 1
return rng
def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None):
match addrspace:
case AddrSpace.GLOBAL:
+38
View File
@@ -0,0 +1,38 @@
#!/usr/bin/env python3
import argparse, pathlib
from typing import Iterator
from tinygrad.viz import serve as viz
from tinygrad.uop.ops import RewriteTrace
from tinygrad.helpers import temp, ansistrip, colored
def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg
def print_data(data:dict) -> None:
if isinstance(data.get("value"), Iterator):
for m in data["value"]:
if not m["diff"]: continue
fp = pathlib.Path(m["upat"][0][0])
print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
print(m["upat"][1])
for line in m["diff"]:
color = "red" if line.startswith("-") else "green" if line.startswith("+") else None
print(colored(line, color))
if data.get("src") is not None: print(data["src"])
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--kernel', type=str, default=None, metavar="NAME", help='Select a kernel by name (optional name, default: only list names)')
parser.add_argument('--select', type=str, default=None, metavar="NAME",
help='Select an item within the chosen kernel (optional name, default: only list names)')
args = parser.parse_args()
viz.trace = viz.load_pickle(pathlib.Path(temp("rewrites.pkl", append_user=True)), default=RewriteTrace([], [], {}))
viz.ctxs = viz.get_rewrites(viz.trace)
for k in viz.ctxs:
if not optional_eq(k, args.kernel): continue
print(k["name"])
if args.kernel is None: continue
for s in k["steps"]:
if not optional_eq(s, args.select): continue
print(" "*s["depth"]+s['name']+(f" - {s['match_count']}" if s.get('match_count') is not None else ''))
if args.select is not None: print_data(viz.get_render(s['query']))
+10 -7
View File
@@ -1,6 +1,6 @@
[project]
name = "tinygrad"
version = "0.11.0"
version = "0.12.0"
description = "You like pytorch? You like micrograd? You love tinygrad! <3"
authors = [{ name = "George Hotz" }]
@@ -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,15 +61,15 @@ linting = [
# ]
testing_minimal = [
"numpy",
"torch==2.9.0",
"torch==2.9.1",
"pytest",
"pytest-xdist",
"pytest-timeout",
"pytest-split",
"hypothesis",
"hypothesis>=6.148.9",
"z3-solver",
]
testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai"]
testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "ggml-python"]
testing = [
"tinygrad[testing_unit]",
"pillow",
@@ -87,7 +87,6 @@ testing = [
"networkx",
"nibabel",
"bottle",
"ggml-python",
"capstone",
"pycocotools",
"boto3",
@@ -135,9 +134,14 @@ check_untyped_defs = true
explicit_package_bases = true
warn_unreachable = true
warn_redundant_casts = true
strict_equality = true
# NOTE: had to comment this out to make mypy pass on both CI and OSX
#warn_unused_ignores = true
[[tool.mypy.overrides]]
module = "extra.*"
follow_imports = "skip"
[tool.pytest.ini_options]
norecursedirs = [
"extra",
@@ -188,7 +192,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
-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()
-28
View File
@@ -1,28 +0,0 @@
from tinygrad import Tensor, dtypes, GlobalCounters
from tinygrad.engine.realize import get_program
if __name__ == "__main__":
t = Tensor.empty(81920, 4096, dtype=dtypes.half)
GlobalCounters.reset()
t.softmax(-1, dtype="half").realize()
GlobalCounters.reset()
t.softmax(-1, dtype="half", _single_kernel=True).realize()
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.helpers import get_single_element
GlobalCounters.reset()
si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())
k = Kernel(si.ast)
#k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
k.apply_opt(Opt(OptOps.LOCAL, 1, 32))
#k.apply_opt(Opt(OptOps.LOCAL, 0, 8))
k.apply_opt(Opt(OptOps.UNROLL, 1, 4))
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
#k.apply_opt(Opt(OptOps.GROUP, 1, 256))
#k.apply_opt(Opt(OptOps.GROUP, 0, 32))
#k.apply_opt(Opt(OptOps.GROUP, 1, 32))
#k.apply_opt(Opt(OptOps.GROUP, 0, 32))
from tinygrad.engine.realize import CompiledRunner, ExecItem
run = CompiledRunner(prg:=get_program(k.ast, k.opts, k.applied_opts))
ExecItem(k.ast, list(si.bufs), prg=run).run()
-56
View File
@@ -1,56 +0,0 @@
# ruff: noqa: E501
from tinygrad import dtypes
from tinygrad.helpers import Timing, getenv
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.engine.realize import get_program, CompiledRunner
from tinygrad.uop.ops import UOp, Ops, AxisType
if __name__ == "__main__":
if getenv("TC", 0) == 0:
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1179648), arg=0, src=())
c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
c2 = UOp.range(UOp.const(dtypes.int, 64), 1, AxisType.GLOBAL)
c3 = UOp.range(UOp.const(dtypes.int, 6), 2, AxisType.GLOBAL)
c4 = UOp.range(UOp.const(dtypes.int, 6), 3, AxisType.GLOBAL)
c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(2097152), arg=1, src=())
c6 = UOp.range(UOp.const(dtypes.int, 64), 1004, AxisType.REDUCE)
c7 = UOp.range(UOp.const(dtypes.int, 3), 1005, AxisType.REDUCE)
c8 = UOp.range(UOp.const(dtypes.int, 3), 1006, AxisType.REDUCE)
c9 = c5.index(((((((c1*UOp.const(dtypes.int, 4096))+(c3*UOp.const(dtypes.int, 8)))+c4)+(c6*UOp.const(dtypes.int, 64)))+(c7*UOp.const(dtypes.int, 8)))+c8), UOp.const(dtypes.bool, True)).load()
c10 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(36864), arg=2, src=())
c11 = c10.index(((((c2*UOp.const(dtypes.int, 576))+(c6*UOp.const(dtypes.int, 9)))+(c7*UOp.const(dtypes.int, 3)))+c8), UOp.const(dtypes.bool, True)).load()
c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=3, src=())
c13 = c12.index(c2, UOp.const(dtypes.bool, True)).load()
c14 = ((c9*c11).reduce(c6, c7, c8, arg=Ops.ADD)+c13)
c15 = c0.index(((((c1*UOp.const(dtypes.int, 2304))+(c2*UOp.const(dtypes.int, 36)))+(c3*UOp.const(dtypes.int, 6)))+c4), UOp.const(dtypes.bool, True)).store(c14, c1, c2, c3, c4)
ast = c15.sink()
# this does have tons of locals
opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=0),
Opt(op=OptOps.LOCAL, axis=0, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=2),
Opt(op=OptOps.GROUPTOP, axis=0, arg=16)]
else:
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(10616832), arg=0, src=())
c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
c2 = UOp.range(UOp.const(dtypes.int, 64), 1, AxisType.GLOBAL)
c3 = UOp.range(UOp.const(dtypes.int, 36), 2, AxisType.GLOBAL)
c4 = UOp.range(UOp.const(dtypes.int, 9), 3, AxisType.GLOBAL)
c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(36864), arg=1, src=())
c6 = UOp.range(UOp.const(dtypes.int, 64), 1004, AxisType.REDUCE)
c7 = c5.index((((c2*UOp.const(dtypes.int, 9))+c4)+(c6*UOp.const(dtypes.int, 576))), UOp.const(dtypes.bool, True)).load()
c8 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1179648), arg=2, src=())
c9 = c8.index((((c1*UOp.const(dtypes.int, 2304))+c3)+(c6*UOp.const(dtypes.int, 36))), UOp.const(dtypes.bool, True)).load()
c10 = (c7*c9).reduce(c6, arg=Ops.ADD)
c11 = c0.index(((((c1*UOp.const(dtypes.int, 20736))+(c2*UOp.const(dtypes.int, 324)))+(c3*UOp.const(dtypes.int, 9)))+c4), UOp.const(dtypes.bool, True)).store(c10, c1, c2, c3, c4)
ast = c11.sink()
opts = [Opt(op=OptOps.TC, axis=0, arg=(0, 0, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=4),
Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=0)]
prg = get_program(ast, opts=opts)
print(prg.src)
for i in range(10):
with Timing(f"try {i}: "):
# NOTE: this doesn't even run the kernel
try: CompiledRunner(prg)
except RuntimeError: pass
+5 -29
View File
@@ -64,8 +64,6 @@ backend_test.exclude('test_qlinearmatmul_2D_int8_float32_cpu')
backend_test.exclude('test_qlinearmatmul_3D_int8_float32_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_maxunpool_export_with_output_shape
backend_test.exclude('test_maxunpool_export_with_output_shape_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True
backend_test.exclude('test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_linear_align_corners
backend_test.exclude('test_resize_downsample_scales_linear_align_corners_cpu')
# tested in external_test_onnx_ops.py::TestMainOnnxOps.test_resize_downsample_scales_cubic_align_corners
@@ -142,17 +140,11 @@ backend_test.exclude('test_affine_grid_3d_expanded_cpu')
backend_test.exclude('test_range_int32_type_negative_delta_expanded_cpu')
# unsupported (strange) ops
backend_test.exclude('test_blackmanwindow_*')
backend_test.exclude('test_bernoulli_*')
backend_test.exclude('test_det_*')
backend_test.exclude('test_col2im_*')
backend_test.exclude('test_hammingwindow_*')
backend_test.exclude('test_hannwindow_*')
backend_test.exclude('test_hardmax_*')
backend_test.exclude('test_gridsample_*')
backend_test.exclude('test_dft_*')
backend_test.exclude('test_einsum_batch_diagonal_cpu*') # TODO: equation = '...ii ->...i'
backend_test.exclude('test_einsum_inner_prod_cpu*') # TODO: equation = 'i,i'
backend_test.exclude('test_unique_*')
backend_test.exclude('test_sequence_*')
backend_test.exclude('test_nonmaxsuppression_*')
@@ -170,17 +162,11 @@ backend_test.exclude('test_scan_*')
backend_test.exclude('test_split_to_sequence_*')
backend_test.exclude('test_ai_onnx_ml_tree_ensemble_*') # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/aionnxml/op_tree_ensemble.py#L121
# TODO: not yet implemented
backend_test.exclude('test_tensorscatter_*')
backend_test.exclude('test_l1normalization_*')
backend_test.exclude('test_l2normalization_*')
backend_test.exclude('test_lpnormalization_*')
backend_test.exclude('test_einsum_scalar_cpu')
backend_test.exclude('test_mod_mixed_sign_float16_cpu')
backend_test.exclude('test_qlinearmatmul_2D_uint8_float16_cpu')
backend_test.exclude('test_qlinearmatmul_3D_uint8_float16_cpu')
backend_test.exclude('test_attention_3d_*')
backend_test.exclude('test_attention_4d_*')
backend_test.exclude('test_attention_4d_diff_heads_mask4d_padded_kv_cpu') # needs nonpad_kv_seqlen handling
backend_test.exclude('test_attention_4d_fp16_cpu') # fp16 numerical issues
backend_test.exclude('test_attention_4d_fp16_expanded_cpu') # fp16 numerical issues
backend_test.exclude('test_attention_4d_gqa_with_past_and_present_fp16_cpu') # fp16 numerical issues
backend_test.exclude('test_attention_4d_gqa_with_past_and_present_fp16_expanded_cpu') # fp16 numerical issues
# rest of the failing tests
@@ -197,16 +183,6 @@ backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad d
backend_test.exclude('test_if_opt_cpu') # ValueError: 13 is not a valid AttributeType
backend_test.exclude('test_if_seq_cpu') # NotImplementedError: op='SequenceConstruct' is not supported
# regression from removing StrEnum in Domain
backend_test.exclude('test_adam_cpu')
backend_test.exclude('test_gradient_of_add_and_mul_cpu')
backend_test.exclude('test_gradient_of_add_cpu')
if Device.DEFAULT in ['CL', 'METAL']:
backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu')
backend_test.exclude('test_resize_upsample_sizes_nearest_axes_3_2_cpu')
backend_test.exclude('test_resize_upsample_sizes_nearest_cpu')
if Device.DEFAULT == "METAL" or (OSX and Device.DEFAULT == "CL"):
# numerical inaccuracy
backend_test.exclude('test_mish_cpu')
+15 -12
View File
@@ -477,18 +477,21 @@ class TestContribOnnxOps(TestOnnxOps):
def test_qlinear_global_average_pool(self):
for dtype, zero_point in [(np.uint8, 128), (np.int8, 0)]:
with self.subTest(dtype=dtype, zero_point=zero_point):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
inputs = {
"X": np.random.randint(dtype_min, dtype_max + 1, [1, 3, 32, 32], dtype=dtype),
"x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"x_zero_point": np.array(zero_point, dtype=dtype),
"y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"y_zero_point": np.array(zero_point, dtype=dtype)
}
attributes = {"channels_last": 0}
outputs = ["C"]
self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs)
for channels_last in [0, 1]:
with self.subTest(dtype=dtype, zero_point=zero_point, channels_last=channels_last):
dtype_min, dtype_max = np.iinfo(dtype).min, np.iinfo(dtype).max
# NCHW for channels_last=0, NHWC for channels_last=1
shape = [1, 3, 32, 32] if channels_last == 0 else [1, 32, 32, 3]
inputs = {
"X": np.random.randint(dtype_min, dtype_max + 1, shape, dtype=dtype),
"x_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"x_zero_point": np.array(zero_point, dtype=dtype),
"y_scale": np.array(np.random.uniform(0.01, 0.1), dtype=np.float32),
"y_zero_point": np.array(zero_point, dtype=dtype)
}
attributes = {"channels_last": channels_last}
outputs = ["C"]
self.helper_test_single_op("QLinearGlobalAveragePool", inputs, attributes, outputs)
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -54,7 +54,7 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str,
def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
# the ast.arg is non None if we are inside of search.py
sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None)
input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
input_ast = ast if ast.op is Ops.PROGRAM else ast.replace(arg=replace(sink_arg, name=p.name))
p2 = get_program(input_ast, renderer=renderer)
def to_str(ret:ProgramSpec) -> str:
# PYTHON renderer pickles UOps, first unpickle and decode here
+1 -1
View File
@@ -187,7 +187,7 @@ class PM4Executor(AMDQueue):
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 and rsrc2 to Python emulator for bounds checking and SGPR layout
# 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)
+3 -3
View File
@@ -12,10 +12,10 @@ libc.mmap.restype = ctypes.c_void_p
drivers = [AMDDriver(), NVDriver()]
tracked_fds = {}
orignal_memoryview = builtins.memoryview
original_memoryview = builtins.memoryview
class TrackedMemoryView:
def __init__(self, data, rcb, wcb):
self.mv = orignal_memoryview(data)
self.mv = original_memoryview(data)
self.rcb, self.wcb = rcb, wcb
def __getitem__(self, index):
@@ -41,7 +41,7 @@ def _memoryview(cls, mem):
for d in drivers:
for st,en,rcb,wcb in d.tracked_addresses:
if st <= addr <= en: return TrackedMemoryView(mem, rcb, wcb)
return orignal_memoryview(mem)
return original_memoryview(mem)
builtins.memoryview = type("memoryview", (), {'__new__': _memoryview}) # type: ignore
def _open(path, flags):
+2 -2
View File
@@ -8,9 +8,9 @@ import torch
def get_question_samp(bsz, seq_len, vocab_size, seed):
np.random.seed(seed)
in_ids= np.random.randint(vocab_size, size=(bsz, seq_len))
in_ids = np.random.randint(vocab_size, size=(bsz, seq_len), dtype=np.int32)
mask = np.random.choice([True, False], size=(bsz, seq_len))
seg_ids = np.random.randint(2, size=(bsz, seq_len)) # type_vocab_size
seg_ids = np.random.randint(2, size=(bsz, seq_len), dtype=np.int32) # type_vocab_size
return in_ids, mask, seg_ids
def set_equal_weights(mdl, torch_mdl):
+4
View File
@@ -1,6 +1,10 @@
import unittest
from extra.models import resnet
from tinygrad import dtypes
from tinygrad.device import is_dtype_supported
# pretrained weights contain num_batches_tracked as int64
@unittest.skipUnless(is_dtype_supported(dtypes.int64), "need int64 support")
class TestResnet(unittest.TestCase):
def test_model_load(self):
model = resnet.ResNet18()
+2
View File
@@ -52,6 +52,8 @@ def wer_helper(result: str, reference: str)->float:
@unittest.skipIf(Device.DEFAULT in ["CPU"], "slow")
@unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support")
# TODO: WEBGPU GPU dispatch dimensions limit
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU GPU dispatch dimensions limit")
class TestWhisper(unittest.TestCase):
@classmethod
def setUpClass(cls):
+2 -1
View File
@@ -11,6 +11,7 @@ from tinygrad.helpers import AMX, AMD_LLVM, CPU_LLVM, Context
from test.helpers import slow
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
from tinygrad.codegen.opt.tc import amd_cdna_1616128
# TODO: write a clean version of this
from test.test_linearizer import helper_realized_ast, helper_linearizer_opt
@@ -120,7 +121,7 @@ class TestTensorCores(unittest.TestCase):
# check excessive padding doesn't trigger padded TC in TC_OPT=2
helper_tc_ensure_uops_and_opts_count(tc.dims[0]//4, tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1]//4, tc.dims[2], tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
if not AMX: # AMX tc.dims[2] == 1
if not AMX and tc not in amd_cdna_1616128: # AMX tc.dims[2] == 1
helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1], tc.dims[2]//8, tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
@Context(ALLOW_TF32=1)
+2 -3
View File
@@ -94,10 +94,9 @@ class TestIndexing(unittest.TestCase):
X = dataset[idxs]
assert X.shape == (4,DDIM)
sched = X.schedule()
# TODO: enable these asserts when the scheduler can handle this
#self.assertEqual(len(sched), 1)
self.assertEqual(len(sched), 1)
run_schedule(sched)
#assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
np.testing.assert_allclose(real_index, X.numpy())
def test_index_fused(self, noopt=1):

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