Compare commits

..
Author SHA1 Message Date
George HotzandGitHub 2fe45b0660 Merge branch 'master' into call_inline 2026-02-10 14:12:08 +08:00
sirhcmandGitHub cdb78954cb better cl compiler name (#14660)
cl_compiler instead of compiler because overriding Compiled.compiler seems more confusing
2026-02-10 01:03:46 -05:00
George HotzandGitHub 1d88723aa0 Merge branch 'master' into call_inline 2026-02-10 14:02:43 +08:00
George HotzandGitHub cc9bf8ccbc move more to null/unit tests (#14658)
* move more to null tests

* move test_gc

* no test fusion op
2026-02-10 13:35:17 +08:00
geohot b0dd3af093 inline all these calls 2026-02-10 13:25:29 +08:00
chenyuandGitHub 83f6d28579 two less realize in setitem (#14655) 2026-02-09 23:45:24 -05:00
geohot e89221e9aa add inline flag for call 2026-02-10 12:19:51 +08:00
wozeparrotandGitHub 69574542ab fix: use correct fa implementation in eval (#14651) 2026-02-09 18:20:44 -08:00
chenyuandGitHub 0dedf4063c minor test_setitem cleanup (#14654) 2026-02-09 20:40:29 -05:00
sirhcmandGitHub b36b62eb59 don't push docker cache for PRs (#14652) 2026-02-09 19:55:55 -05:00
sirhcmandGitHub e6562a5061 remove CompilerPair (#14638) 2026-02-09 19:51:18 -05:00
sirhcmandGitHub 396e1320fb bump cache version for z3 (#14650) 2026-02-09 19:32:07 -05:00
chenyuandGitHub 9e3f24db9f assign realize fix (#14649)
fix the need for explicit assign. track pending assigns for each buffer, and run those before the main realize in order
2026-02-09 17:46:46 -05:00
chenyuandGitHub 0913c068ea clean up setitem disk path (#14648) 2026-02-09 15:58:04 -05:00
chenyuandGitHub 205a1212b7 delegate non Tensor src setitem to assign (#14647)
cannot do this for DISK in the unified path
2026-02-09 13:53:20 -05:00
chenyuandGitHub e9f40f49d4 explicitly check advanced setitem (#14644)
advanced setitem DISK would failed in rangeify with bad error, now it's checked directly in setitem. eventully DISK can use regular setitem path
2026-02-09 13:36:46 -05:00
chenyuandGitHub 20a132b1c4 relax atol for test_uop_scan_matmul (#14646)
flaky, also log max diff
2026-02-09 13:25:19 -05:00
qazalandGitHub 50d3f6cea5 EVAL_BS=0 in llama profile (#14643) 2026-02-10 00:49:43 +09:00
chenyuandGitHub 8a2c23d3dc raise RuntimeError for setitem dtype mismatch (#14642) 2026-02-09 10:37:08 -05:00
qazalandGitHub 80b0119cef llama: add new asm gemm shape (#14611)
* llama: add new asm gemm shape

* work

* cleanup

* half dtype

* more comment
2026-02-10 00:34:29 +09:00
chenyuandGitHub a49e038c0c dont manually broadcast in setitem (#14641)
handled by assign
2026-02-09 09:34:09 -05:00
chenyuandGitHub 2c3e3559eb remove a contiguous in basic setitem (#14640)
handled in rangeify
2026-02-09 09:19:46 -05:00
chenyuandGitHub 6c0c8e2ac3 setitem push a realize to basic setitem (#14637)
advanced setitem does not need it
2026-02-09 08:54:07 -05:00
nimlgenandGitHub e087c58ae0 print tables in llama/profile.sh (#14639) 2026-02-09 12:32:54 +03:00
sirhcmandGitHub 27f7ea478b new style DSP renderer (#14636)
* new style DSP renderer

* cleanup
2026-02-09 00:39:03 -05:00
sirhcmandGitHub efac5b9ef6 new style NV/CUDA renderers, try 2 (#14634)
* new style NV/CUDA renderers, try 2

* fix diskcache
2026-02-08 22:58:48 -05:00
sirhcmandGitHub 0ebb508b85 new style metal compiler (#14632) 2026-02-08 21:58:25 -05:00
sirhcmandGitHub 9eef9f38ad new style python renderer (#14631) 2026-02-08 21:45:07 -05:00
sirhcmandGitHub 5f2f2cc956 Revert "new style NV/CUDA renderers (#14627)" (#14633)
This reverts commit 0e505951b0.
2026-02-08 21:16:03 -05:00
sirhcmandGitHub 4ad787ece2 new style CPULLVMRenderer (#14629) 2026-02-08 21:05:01 -05:00
sirhcmandGitHub 0e505951b0 new style NV/CUDA renderers (#14627)
* new style NV/CUDA renderers

* fix pickle

* oops

* fix CUDA_CC=NVCC

* mockgpu uses PTXCompiler

* oops

* ruff

* dont discard stderr

* ugh
2026-02-08 21:04:51 -05:00
Filip BrzekandGitHub 1667669c46 fix: python3 -m tinygrad.device reporting on AMD/CPU (#14622)
* test: device module expects PASS in -m tinygrad.device for CPU

* fix: use device._compiler_name instead of unwrap_class_type(compiler).__name__ in enumerate_devices_str
2026-02-08 20:22:35 +03:00
nimlgenandGitHub 01a4ee4d66 do not hive_reset when amdgpu (#14624) 2026-02-08 19:14:13 +03:00
nimlgenandGitHub a615b9d781 am: f8_mode for gfx94x only (#14620) 2026-02-08 17:38:48 +03:00
chenyuandGitHub c28f7d0167 remove realize in Tensor.svd (#14623) 2026-02-08 09:36:31 -05:00
qazalandGitHub 087dab4c3b gemm/asm: split out cdna tests from CI (#14619)
* gemm/asm: split out cdna tests from CI

* reorder

* work
2026-02-08 21:33:42 +09:00
George HotzandGitHub 183d38b128 remove CUSTOM_KERNEL / directly construct it (#14604)
* remove CUSTOM_KERNEL / directly construct it

* clean that up

* simpler multi

* custom kernel spec

* remove Kernel

* fix multi

* use sharded shape

* explicit regression test
2026-02-08 18:43:33 +08:00
nimlgenandGitHub e29a88ca09 hive_reset respects lock (#14618) 2026-02-08 10:47:25 +03:00
qazalandGitHub b10802eb53 use existing VIZ ContextVar instead of getenv (#14610) 2026-02-08 15:37:55 +09:00
chenyuandGitHub 510b65489e style change rangeify assign [pr] (#14616)
consistent naming, also a standalone fucntion to replace complicated lambda
2026-02-07 15:47:32 -05:00
chenyuandGitHub b7afd4471c use arg instead of 3rd op for ASSIGN [pr] (#14613) 2026-02-07 14:17:10 -05:00
nimlgenandGitHub 88c3022223 amd: kfd iface early exit (#14612)
* amd: kfd iface early exit

* l

* revert
2026-02-07 18:57:10 +03:00
nimlgenandGitHub ce7bfc6ce8 nv: use nv_flags for all fields (#14607) 2026-02-07 15:01:38 +03:00
qazalandGitHub c2544e2252 viz: remove outdated comment (#14608) 2026-02-07 20:05:24 +09:00
nimlgenandGitHub 6838b35cff mockgpu: hevc (#14606)
* mockgpu: hevc

* eng
2026-02-07 12:27:55 +03:00
chenyuandGitHub 884592f6c8 pin z3-solver version (#14605)
found exact input that crashes z3 4.15.4
2026-02-06 22:49:31 -05:00
George HotzandGitHub 7a2a3b5c71 Remove Ops.KERNEL, it's all Ops.CALL now (#14603) 2026-02-07 10:21:54 +08:00
George HotzandGitHub ca6604eae2 kernel is call (#14577)
* call is kernel

* closer

* fix bugs

* dedup

* pm_gate_kernel_sink

* better

* Revert "better"

This reverts commit b4c799b810.

* Reapply "better"

This reverts commit e53f094ce7.

* cleanups

* work

* remove junk

* subtle fix

* index

* viz cleanups

* disable assert for now
2026-02-07 10:10:14 +08:00
wozeparrotandGitHub d87ae1c84c feat: tinyfs load test in benchmark (#14602) 2026-02-06 18:00:00 -08:00
ttomsaandGitHub 462b455562 cleanup linearize (#14523) 2026-02-07 08:54:02 +08:00
ttomsaandGitHub d5652e4da2 new dtype aliases (#14596) 2026-02-07 08:53:35 +08:00
sirhcmandGitHub ad9e2f0de7 decompose bf16 (#14601) 2026-02-06 19:24:09 -05:00
sirhcmandGitHub 7bb45e7df0 decompose fp8 to bigger floats [skip_process_replay] (#14554)
* decompose fp8 also

* it works

* cleanup

* no shift required

* default to float

* cleanup

* fixes

* fp8e5m2

* don't rely on behavior comparing nans

* cleanup
2026-02-06 19:05:40 -05:00
chenyuandGitHub 81f6cdb4ab delete realize_assign [pr] (#14575)
use realize and realize_srcs like COPY and STORE. src[0] always has BUFFER for base
2026-02-06 17:12:33 -05:00
chenyuandGitHub 7d193a6e26 fix wgsl bitcast (#14600)
was wrong for signed int
2026-02-06 16:57:36 -05:00
chenyuandGitHub b9fe8b7591 fix opt in process replay [pr] (#14599) 2026-02-06 16:49:56 -05:00
chenyuandGitHub 197ebcbbbc log seed with flush=True in fuzz_symbolic (#14597)
* log seed with flush=True in fuzz_symbolic

i think z3 can crash. added reading seed from argv to see if we repro later

* fuzz_symbolic_symbolic_div
2026-02-06 15:03:57 -05:00
nimlgenandGitHub fbb67a3f95 am_smi: fix after regen (#14594) 2026-02-06 20:57:41 +03:00
qazalandGitHub a80fb4e641 viz: better ordering of device engines in profiler (#14590) 2026-02-06 23:08:09 +09:00
qazalandGitHub b7e3fbe07e llama: add VIZ=-1 to dev_run (#14583)
* llama: add VIZ=-1 to dev_run

* readme

* cleaner

* add profile.sh script

* better grouping of options

* add other row

* readme edits

* work
2026-02-06 22:59:22 +09:00
nimlgenandGitHub fbeb978170 diff devices for sdma (#14589)
* start

* x

* fix

* sdma

* c

* clean

* x

* hm

* cleaer
2026-02-06 16:39:12 +03:00
George HotzandGitHub 7cb996e153 bottom up earliest rewrites (#14587)
* better

* bottom up earliest rewrites

* fix
2026-02-06 18:13:07 +08:00
George HotzandGitHub 03af2404e2 small changes and test fixes from kernel is call (#14586) 2026-02-06 17:08:33 +08:00
George HotzandGitHub 3c26ce29b2 make disk tensor tests process safe (#14584) 2026-02-06 15:39:55 +08:00
qazalandGitHub cf73d7e2a7 hotfix: disable slower asm gemm shape from llama seqlen 8192 (#14582) 2026-02-06 15:05:19 +09:00
qazalandGitHub be77873974 llama: contig backward for wk / wv matmul backward (#14581) 2026-02-06 14:54:00 +09:00
chenyuandGitHub 15d3344d9e use int inputs in test_assign (#14580)
int is less flaky
2026-02-06 00:07:31 -05:00
qazalandGitHub 50a166a5fa viz: cleanup amdgpu target mapping (#14579)
* viz: cleanup amdgpu target mapping

* linter

* unwraps
2026-02-06 13:51:51 +09:00
chenyuandGitHub b09dc646f5 revert some late_buffer_view change (#14578)
revert #14478 which breaks tinyfs
2026-02-05 22:51:40 -05:00
chenyuandGitHub d41836f135 remove KERNEL special case in realize_assign [pr] (#14573) 2026-02-05 21:55:44 -05:00
George HotzandGitHub 6cbcf98627 KernelInfo is required on get_program (#14571)
* rangeify always adds KernelInfo

* fix tests

* skip flaky test
2026-02-06 10:49:27 +08:00
George HotzandGitHub 28c56a783c add CallInfo and viz call toggle (#14570) 2026-02-06 09:30:58 +08:00
wozeparrotandGitHub f73468d516 fa: block skipping for fa kv bwd (#14569) 2026-02-05 16:13:53 -08:00
chenyuandGitHub b7ef775677 more cleanup in create_schedule [pr] (#14566)
fixed wrong comments and simplified queue building
2026-02-05 16:12:17 -05:00
Garret CastroandGitHub cee7ef7ab2 disable threads (#14555) 2026-02-05 16:11:32 -05:00
chenyuandGitHub 79b7799dba clean up linearize schedule [pr] (#14565)
* clean up linearize schedule [pr]

don't mix ScheduleItem and UOp in schedule queue

* ok
2026-02-05 15:24:09 -05:00
chenyuandGitHub 41a179f542 fix test_xlm_roberta_large (#14564)
onnxruntime does not allow symlink that's outside model dir. update snapshot_download to use local_dir instead of cache_dir. some ad hoc migration step to copy the existing model too
2026-02-05 14:56:06 -05:00
sirhcmandGitHub aa9dc50577 dtype decomps don't require bitshifts (#14542)
* dtype decomps don't require bitshifts

* simplify shr/shl

* ruff
2026-02-05 14:42:30 -05:00
sirhcmandGitHub b47397ab17 list ml_dtypes as dependency for DSP (#14562)
* pin onnxruntime to 1.23.2 for DSP

* list ml_dtypes instead

This reverts commit 84bb2cc0fc.
2026-02-05 14:27:50 -05:00
chenyuandGitHub 2b47a9a1b5 skip test_xlm_roberta_large (#14563)
symlink model not allowed in latest onnxruntime
2026-02-05 14:00:24 -05:00
chenyuandGitHub 42c18da88a add Ops asserts in toposort sched_sink [pr] (#14561)
more explicit
2026-02-05 12:40:02 -05:00
nimlgenandGitHub 483bba4f05 nv: use prof_exec_counter (#14559) 2026-02-05 19:00:14 +03:00
qazalandGitHub 190042358f llama: faster bf16 matmul / rope backward (#14558) 2026-02-05 23:57:25 +09:00
b398335f62 assembly/amd: fix saturation in python remu (#14557)
* PYTHONREMU: failing test for V_SUB_NC_U32_E64 clamp

* fix saturation in PYTHON_REMU

* simpler

* more tests, less lines

---------

Co-authored-by: Christopher Milan <[email protected]>
2026-02-05 18:35:57 +08:00
wozeparrotandGitHub c1ea6687e5 fa: simpler is faster (#14548) 2026-02-05 01:13:17 -08:00
43e7eda4e7 grad_b uses custom gemm (#14550)
* grad_b uses custom gemm

* fix multi backward, acc is in float32

* test_gemm_batched

* square gemm

---------

Co-authored-by: qazal <[email protected]>
Co-authored-by: qazal <[email protected]>
2026-02-05 15:22:27 +09:00
qazalandGitHub f9cfb64cd9 test asm_gemm in CI (#14551)
* test asm_gemm in CI

* default float16

* use a smaller shape for multi

* smaller size

* smaller for CI

* smaller for ci

* need half
2026-02-05 13:32:22 +09:00
chenyuandGitHub c0ca7f9c51 use more UOp.sum and UOp.prod [pr] (#14549) 2026-02-04 22:05:20 -05:00
chenyuandGitHub e8dace41b6 clean up UOp.vars [pr] (#14547) 2026-02-04 20:52:25 -05:00
sirhcmandGitHub 232848d086 PYTHONREMU: VOP3P integer operations with constants don't cast to fp16 (#14546)
* PYTHONREMU: VOP3P integer operations with constants don't cast to fp16

* put that back

* cleaner

* do that once
2026-02-04 20:10:59 -05:00
wozeparrotandGitHub 2966619834 feat: llama uses enable_gqa during training (#14545) 2026-02-04 16:22:31 -08:00
chenyuandGitHub 664f1bf76d minor ops/jit cleanups [pr] (#14543) 2026-02-04 17:21:34 -05:00
chenyuandGitHub 03d0fa9c3f merge as_buf into buf_uop [pr] (#14541) 2026-02-04 16:32:23 -05:00
chenyuandGitHub 43ef24a8af remove buf_target [pr] (#14540)
not really needed
2026-02-04 15:03:47 -05:00
chenyuandGitHub 8b7343b950 clean up is_realized [pr] (#14538)
base cannot be Ops.MULTI since MULTI is a view now
2026-02-04 14:24:10 -05:00
sirhcmandGitHub 5338ce6b74 test S_PACK in extra/assembly/amd/test/hw (#14537)
* S_PACK_LL_B32_B16 in test/hw

* add rest of S_PACK instructions
2026-02-04 14:17:16 -05:00
chenyuandGitHub 9052db678f remove allow_shape_mismatch in Tensor.replace (#14536)
move all logic to torch_backend and not hacking Tensor method
2026-02-04 12:38:18 -05:00
nimlgenandGitHub ec2b6bbda8 hcq: update signal logic (#14531) 2026-02-04 19:32:56 +03:00
nimlgenandGitHub 62786d488a am: mi3xx perf (#14529) 2026-02-04 19:32:43 +03:00
chenyuandGitHub d57d24c7d4 Buffer.as_buffer -> Buffer.as_memoryview [pr] (#14535)
it casts to memoryview. also inline the as_typed_buffer checks to Tensor._data
2026-02-04 11:31:11 -05:00
chenyuandGitHub 024f57ecf5 jit input_buffers cleanup [pr] (#14532) 2026-02-04 10:14:38 -05:00
chenyuandGitHub 67f91e897b UOp.is_contiguous -> UOp.has_buffer_identity [pr] (#14530)
one more confusing buffer related method, but it's definitely not is_contiguous
2026-02-04 09:21:26 -05:00
George HotzandGitHub fb9df1e031 pretty print binary (#14520) 2026-02-04 18:04:35 +08:00
sirhcmandGitHub 8c3c026d86 decomp float16 to float32 (#14417)
* decomp float16 to float32

* denormals arent zero

* add test

* denormals are zero

* fix

* oops

* bitcast works

* fix LOADs

* test_dtype passing

* cleanup

* mypy

* debug print

* only emulate if EMULATED

* very ugly, but passes spec

* add test_dtype_alu tests

* Revert "very ugly, but passes spec"

This reverts commit fdc3999b654d630678bf208927ab2f55e026b4ca.

* bottom up decompositions

* that should have symbolic

* simplify a bit

* SPEC really works

* run with DEBUG

* debug=4

* rm debug
2026-02-04 01:37:47 -05:00
sirhcmandGitHub ecbce5269e PYTHONREMU properly supports S_PACK_LL_B32_B16 (#14527)
* PYTHONREMU properly supports S_PACK_LL_B32_B16

* default
2026-02-03 23:45:33 -05:00
wozeparrotandGitHub 720c9597a9 feat: llama uses is_causal on sdpa during training (#14528) 2026-02-03 20:24:30 -08:00
chenyuandGitHub 9c2fc118ef relax setitem target check (#14526)
old check was too conservative
2026-02-03 22:32:49 -05:00
qazalandGitHub d1bfbe9ce3 isolate slow llama gemm (#14525) 2026-02-04 12:20:10 +09:00
nimlgenandGitHub 2f55005ad9 qcom: sync cpu cache when from_blob (#14518)
* um

* fx

* d

* x

* x

* x

* x

* f

* ren
2026-02-03 21:51:03 +03:00
chenyuandGitHub ee9d6a1f36 remove DEFINE_VAR in to_define_global [pr] (#14522)
not needed
2026-02-03 10:12:33 -05:00
Nino RisteskiandGitHub af4c74bb41 delete extra cast (#14517) 2026-02-03 08:29:04 -05:00
chenyuandGitHub 9d1e9e643e removed a duplicated remove_bufferize rule [pr] (#14519) 2026-02-03 08:28:07 -05:00
George HotzandGitHub d59e6e7a37 move more tests to test/null, split some existing ones (#14512)
* move more tests to test/null, split some existing ones

* null work

* null work

* move more

* fixes

* move PIL

* PIL in CLIP

* don't move that
2026-02-03 20:20:20 +08:00
qazalandGitHub a98c53769a ASM_GEMM=1 runs the UOp gemm on non cdna (#14516)
* ASM_GEMM=1 runs the UOp gemm on non cdna

tests run on mac in 3 seconds

* min diff
2026-02-03 20:42:02 +09:00
qazalandGitHub 5c1d21349e viz: profiler command line tool (#14515) 2026-02-03 19:51:25 +09:00
George HotzandGitHub dd2de4f838 rename all DEFINE_GLOBAL to PARAM (#14511) 2026-02-03 15:09:38 +08:00
George HotzandGitHub dc77b3318b move files that pass with NULL=1 to test/null (#14508)
* move files that pass with NULL=1 to test/null

* fix windows

* cpu 0

* bugfix + durations
2026-02-03 13:52:36 +08:00
George HotzandGitHub 888819ee09 call autodiff gradient (#14510) 2026-02-03 13:51:02 +08:00
wozeparrotandGitHub bbcd3d67a3 fa: faster (#14453) 2026-02-02 21:34:17 -08:00
sirhcmandGitHub e579613b90 IR3 has aux (#14509) 2026-02-02 23:46:41 -05:00
George HotzandGitHub 85c7b23160 add pytest -nauto to benchmark for mac (#14458)
* add pytest -nauto to benchmark

* 3 minute timeout

* 3 min

* setup env

* comment

* fresh db

* in the pyenv
2026-02-03 12:26:09 +08:00
sirhcmandGitHub a5d7eb37db IR3 works on versions earlier than 3.14 (#14507) 2026-02-02 23:10:19 -05:00
George HotzandGitHub 33c886cafa disable copyout on NULL backend by default (#14506)
* disable copyout on NULL backend

* gate it

* allow copyout on some tests
2026-02-03 11:57:47 +08:00
chenyuandGitHub 3c5845e8a5 remove cut_store_range (#14505)
special scheduling for CPU
2026-02-02 21:58:36 -05:00
chenyuandGitHub 4f2e7aed24 fix multiple REDUCE on same RANGE (#14504)
each RANGE maps to one END, but reduce_to_acc is local and would not know this
2026-02-02 20:42:09 -05:00
chenyuandGitHub 93c41a78fa clean up NOOP [pr] (#14503)
should not be used as a COPY, started with removing from ALWAYS_RUN_OPS
2026-02-02 19:46:45 -05:00
chenyuandGitHub 66d2b02f11 delete files that depends on extra.optimization.helpers (#14499) 2026-02-02 13:33:33 -05:00
George HotzandGitHub ec0398fceb test amd gpu crashes (#14459)
* test amd gpu crashes

* cleanup

* less sketch tests
2026-02-02 18:57:47 +03:00
nimlgenandGitHub 6e4238c016 amd: recovery (#14461)
* rec

* ?

* rv

* cleaner

* post merge

* not used

* um

* clnr

* x

* x

* d

* move
2026-02-02 18:57:35 +03:00
chenyuandGitHub 61ca19ff24 after with empty src is self [pr] (#14496) 2026-02-02 10:19:05 -05:00
George HotzandGitHub 6e958dbfd4 assembly/amd: add RDNA4 support to emulator (#14341)
* start new rdna4

* work

* plus works

* more pass

* rdna4

* assembly/amd: fix RDNA4 emulator for float16 and VOP3 clamp

* stale

* rev

* rr

* rdna4 emu tests

* cleanup

* cleanup

* simp

* works

* better factorizaion

* hacks

* fix mockgpu

* guard both

* cleaner

* gate

* bug fix and a few tests

* all test_tiny
2026-02-02 21:35:59 +08:00
chenyuandGitHub a908f447d5 remove disk special case in mstack_early_shrink [pr] (#14494) 2026-02-02 08:34:45 -05:00
qazalandGitHub 965940dd00 sqtt: update examples after event field change (#14493)
* regen sqtt examples

* cdna

* rdna4

* packet counts for rdna3

* sqttmap work
2026-02-02 21:39:48 +09:00
George HotzandGitHub 965149a46d assembly/amd: add ds perm instructions (#14486)
* assembly/amd: add ds perm instructions

* NO SKIP

* fix preexisting RDNA3 issues

* pcode

* assert

* asserts

* unify

* simp

* good fix
2026-02-02 16:02:00 +08:00
qazalandGitHub 1746d1f997 remove SPEC=0 context in custom_kernel tests, pyrender always skips it (#14489) 2026-02-02 16:32:01 +09:00
George HotzandGitHub d4007f36e0 remove DEFINE_GLOBAL (it is PARAM now) (#14488) 2026-02-02 14:56:37 +08:00
qazalandGitHub 6c487656f9 viz: kernel metadata from rodata entry (#14487) 2026-02-02 15:41:42 +09:00
Robbe DerksandGitHub d75a1b0d5a usbgpu: use BOT interface for patch.py (#13644)
* BOT usage

* cleanup

* fix lint

* fix ruff

* fix -7?
2026-02-02 11:54:46 +08:00
sirhcmandGitHub 2931b52875 skip autogen if MTLCompiler is loaded (#14466) 2026-02-01 22:12:27 -05:00
George HotzandGitHub 9a32d6e090 add depth limit for SPEC=2 (#14485)
* make SPEC=2 work for everything

* that's a horrible fix

* add depth limit
2026-02-02 10:43:28 +08:00
George HotzandGitHub 368a692e1a make SPEC=2 work for everything (#14476)
* make SPEC=2 work for everything

* that's a horrible fix
2026-02-02 10:30:56 +08:00
chenyuandGitHub ea1f1d2b9d test_assign_to_bitcast_view (#14483)
currently disk allows assign same size dtype into a bitcasted view
2026-02-01 16:46:04 -05:00
chenyuandGitHub 6deeccc192 fix RING with single dest (#14482) 2026-02-01 12:14:46 -05:00
chenyuandGitHub 3ff390159b don't implicitly change dtype in assign (#14481)
broadcast shape is fine, but implicitly cast dtype is hard to find
2026-02-01 11:48:54 -05:00
2111762a48 failed test case for RING output device (#14191)
* Add enable/disable scheduler cache ContextVar

* add allreduce ring and naive to() tests

* clearer test comparing native vs ring allreduce

* split tests, add helper

* removing trailing whitespace

---------

Co-authored-by: chenyu <[email protected]>
2026-02-01 11:48:43 -05:00
chenyuandGitHub 02afae04f4 atol in test_call_gemm (#14480)
flaky
2026-02-01 11:24:58 -05:00
chenyuandGitHub 5705398a1f assign cleanup [pr] (#14479)
share more code path between disk and non-disk. also raise RuntimeError instead of Assert for mismatches
2026-02-01 09:10:22 -05:00
chenyuandGitHub da500dbe06 simplify late_buffer_view [pr] (#14478)
check the only allowed Ops in the chain, and offset cannot be negative
2026-01-31 22:38:40 -05:00
chenyuandGitHub b4f96301e0 remove unused rules [pr] (#14477) 2026-01-31 21:29:30 -05:00
qazalandGitHub 54e78dbec8 viz: remove hardcoded strings in cfg tests (#14462) 2026-02-01 09:30:43 +09:00
chenyuandGitHub 5d38db9da6 generic bitcast assign (#14474)
a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))
2026-01-31 17:29:20 -05:00
chenyuandGitHub b38fc43b07 assert assign dtype mismatch for disk [pr] (#14473)
the disk hack is generally wrong, now force bitcast on the source before assign
2026-01-31 17:08:54 -05:00
chenyuandGitHub ced886f26c failed test case for assign into bitcast (#14469)
* failed test case for assign into bitcast

DISK assign has custom hack for this. need to fix before we can unify assign

* test_assign_bitcast_different_size
2026-01-31 14:26:47 -05:00
chenyuandGitHub 81eee5b30a unused spec [pr] (#14468)
no BUFFER_VIEW in tensor, and no CONTIGUOUS in KERNEL
2026-01-31 13:53:16 -05:00
nimlgenandGitHub f873c7b6c5 amd: fetch_name is file_name (#14465) 2026-01-31 20:11:07 +03:00
chenyuandGitHub c765641215 remove unused allow_any_len [pr] (#14464)
STORE has 2 src, RESHAPE has 2 src, BUFFER has 2 src
added some tests for the untested allow_any_len
2026-01-31 11:05:42 -05:00
chenyuandGitHub b4f5a51ebb move tests to unit (#14463)
test_uop_graph does not need device, test_memory_planner can use NULL
2026-01-31 10:49:31 -05:00
qazalandGitHub 616e9c1483 CDNA assembly gemm in tensor.py with flag (#14310)
* work

* work

* the assembly

* remove the old one

* remove ws bufs, assert splitk

* notes cleanup

* work

* gemm args

* gemm in mixins would be nice

* add gemm gradient

* print counters

* the realize is for DEBUG=2 aesthetics

* dedup

* rewrite to python dsl, no list copies

* leave that

* add B, M, N, K to gemm name

* it's M0 not NULL

* fp16 support

* test cleanup + more gemms

* work from viz

* more work

* gemm batch_size

* xccg path work

* tiny comments on the label naming

* s_waitcnt
2026-01-31 22:34:14 +09:00
chenyuandGitHub 55f806b713 tighter late_buffer_view match [pr] (#14456)
src must be len 2 at that point
2026-01-31 07:28:26 -05:00
qazalandGitHub d69bc5aa1a make DEV=NULL EMULATE=AMD amd_asm_matmul run (#14460) 2026-01-31 20:45:24 +09:00
qazalandGitHub 4976544bf9 multi ram usage tests on the NULL device (#14457) 2026-01-31 14:14:53 +09:00
chenyuandGitHub 99b44121bc failed test case for non-consecutive disk read (#14455)
silently fail now
2026-01-30 23:44:04 -05:00
George HotzandGitHub b705c9143c assembly/amd: test more instructions (#14365)
* assembly/amd: test more instructions

* more

* passing

* revert

* no const fold

* remove junk

* cleaner
2026-01-31 12:40:22 +08:00
George HotzandGitHub c9a3ddb341 benchmark llama walltime script (#14454)
* benchmark llama walltime script

* adj layers
2026-01-31 10:21:54 +08:00
George HotzandGitHub f5346d6a1a fix USE_ATOMICS for non float dtypes and make it the default (#14444)
* embedded multistep test

* complex test

* with jit

* fix dtypes and reenable USE_ATOMICS

* that test didn't catch anything
2026-01-31 09:44:16 +08:00
sirhcmandGitHub e575dd8275 prevent UB in long decomp and more emulated tests (#14447) 2026-01-30 19:38:41 -05:00
chenyuandGitHub 3204f94454 correct var_vals schedule filter (#14451)
complete_create_schedule_with_vars returns var_vals that's used in schedule
2026-01-30 17:10:07 -05:00
chenyuandGitHub cfcd1debb5 test schedule with multiple AFTER (#14449) 2026-01-30 15:59:00 -05:00
nimlgenandGitHub 486d53d646 device: call free for external_ptr (#14448)
* device: call free for external_ptr

* lin
2026-01-30 23:53:17 +03:00
nimlgenandGitHub e0978498dc amd: read_ptr/write_ptr/doorbells are not lists (#14445) 2026-01-30 23:11:57 +03:00
sirhcmandGitHub 1803ee939d EMULATED_DTYPES=long works with CPU_LLVM (#14446) 2026-01-30 13:54:43 -05:00
chenyuandGitHub 03613e83ad update TestTensorMetadata (#14443)
run with SCACHE=0 some more TODOs
2026-01-30 12:39:01 -05:00
geohot cbb1eed57b hotfix: partial revert of 9eb449f88, caused llama NaN 2026-01-30 17:19:27 +00:00
chenyuandGitHub 26f5c00265 move TestTensorMetadata to unit (#14442) 2026-01-30 12:14:21 -05:00
chenyuandGitHub c05a0b85ae flip unique const src order [pr] (#14441)
* flip unique const src order [pr]

matches buffer, simplifies replace_input_buffer

* combine rules
2026-01-30 11:44:18 -05:00
geohot ee2c78709d mlperf/llama: disable USE_ATOMICS for now 2026-01-31 00:42:08 +08:00
chenyuandGitHub beecac4d85 expand ranges -> unroll outer ranges [pr] (#14440) 2026-01-30 11:26:05 -05:00
chenyuandGitHub 9eb449f882 clean up toposort sched_sink [pr] (#14439) 2026-01-30 10:18:28 -05:00
George HotzandGitHub 838cd078bc use atomics for embedding backward (#14400)
* embedding is slow

* failing

* float is fine

* null

* it fails

* simplify embedding with broadcasting

* ATOMIC_ADD incoming

* min change

* simpler test

* better test

* fix test

* real test

* simpler

* cleanups

* types and names

* _zero_kernel

* grad multi

* hack

* none

* multi unshard

* more for call

* don't tag in call

* good

* call_multi

* call_multi wow claude is useless

* embedding backward mutli test

* test passes

* fix as_param

* shape_to_shape_arg

* add clip

* before cast

* fix spec=2, use atomics
2026-01-30 18:10:59 +08:00
nimlgenandGitHub 1998e0bb28 nv: add prof props to dev (#14437) 2026-01-30 12:51:43 +03:00
George HotzandGitHub 7a9dee4e50 add call/param UOps (#14433)
* add call/param UOps

* resolve call

* skip that for now

* grad on call

* fix tests
2026-01-30 14:51:45 +08:00
qazalandGitHub 66d6a68016 viz: sqtt work from cdna gemm (#14434)
* it's the tag

* initialize rows based on the disasm

* test_cfg with Ops.BINARY

* pyremu wants s_code_end?

* test_diamond

* diff cleanup
2026-01-30 14:00:56 +09:00
sirhcmandGitHub 88caf57ef4 ci: unify python versions (#14430) 2026-01-29 21:42:03 -05:00
chenyuandGitHub 86a204d22a allow Tensor setitem input to be list/tuple (#14432)
matches assign, and generally matches numpy
2026-01-29 21:26:58 -05:00
chenyuandGitHub 4a80319093 clean up split_store final logic [pr] (#14429)
explicitly check the structure
2026-01-29 18:40:07 -05:00
sirhcmandGitHub e47f12f671 ci: replace testing_minimal with testing_unit (#14427) 2026-01-29 18:02:43 -05:00
wozeparrotandGitHub c2fb8b208f fa: 32 block size (#14416) 2026-01-29 13:59:13 -08:00
chenyuandGitHub a979fafae5 cleanup around disk buffer [pr] (#14428)
style change, prep for refactor
2026-01-29 16:18:44 -05:00
nimlgenandGitHub dc977a03b0 nv_pma: bw decoder (#14424)
* nv_pma: bw decoder

* decoder fix

* better
2026-01-30 00:12:39 +03:00
chenyuandGitHub ddc041854b failed test case for disk setitem (#14426)
strided setitem is wrong
2026-01-29 14:54:19 -05:00
chenyuandGitHub 31706bf6bc add few more types [pr] (#14425) 2026-01-29 14:04:09 -05:00
nimlgenandGitHub 2d5c24879f nv: pma for 5090 (#14420)
* nv: pma for 5090

* hm

* 4090
2026-01-29 20:06:01 +03:00
nimlgenandGitHub c8dc6332d2 memory: read_fields is not universal (#14348) 2026-01-29 20:00:00 +03:00
chenyuandGitHub dbe8f034a7 pass z3.Context in validate ctx [pr] (#14423)
does not need to pass the whole solver
2026-01-29 11:11:47 -05:00
chenyuandGitHub 033ce1b885 types for validate.py (#14422) 2026-01-29 10:56:50 -05:00
nimlgenandGitHub 230d08ec70 test for am recovery and faults handling (#14421)
* test for am recovery and faults handling

* linter
2026-01-29 17:11:24 +03:00
George HotzandGitHub 793afbd473 simplify nn.Embedding, support AFTER in CUSTOM_KERNEL (#14419) 2026-01-29 17:22:13 +08:00
sirhcmandGitHub 0c855d6149 ci: remove unused pydeps (#14418) 2026-01-29 01:51:26 -05:00
wozeparrotandGitHub 4845e42135 llama3 gradacc fixes (#14414) 2026-01-28 19:12:39 -08:00
chenyuandGitHub 37cde4a01a add one line mypy report (#14415) 2026-01-28 20:39:32 -05:00
chenyuandGitHub 15aed51544 return types for all math.py function (#14413)
calling int() on sint -> int, i think it's better support since some UOp can be safely cast to int
2026-01-28 20:10:11 -05:00
nimlgenandGitHub aec1ae0de1 llama: set manual_seed (#14409) 2026-01-28 14:40:00 -08:00
chenyuandGitHub 0870ed28b1 add Self type to MathMixin (#14411)
these don't cause error
2026-01-28 16:59:38 -05:00
chenyuandGitHub 079f33c208 fix type in Tensor.mean and Tensor.var (#14410)
use Tensor.from_uop to wrap UOp from symbolic shape, kernels are the same
2026-01-28 15:24:02 -05:00
chenyuandGitHub 2b5e99ccc1 minor type cleanups [pr] (#14408)
mypy --warn-redundant-casts has false negative
2026-01-28 14:11:50 -05:00
chenyuandGitHub 726415dbc8 import sint directly in movement.py TYPE_CHECKING (#14406)
avoid creating string TypeAlias, fixed warning in `TYPED=1 python test/test_tiny.py`
2026-01-28 12:47:26 -05:00
nimlgenandGitHub acb2fc36ba nv_pma: add decoder (#14404)
* nv_pma: add decoder

* cl
2026-01-28 20:44:02 +03:00
chenyuandGitHub 7b9bc1d8cf _MockMemoryviewMeta for mockgpu (#14405)
fixed `PYTHONPATH=. TYPED=1 DEV=AMD MOCKGPU=1 python test/test_tiny.py`. basically make `isinstance(TrackedMemoryView_instance, memoryview)` true
2026-01-28 11:59:00 -05:00
chenyuandGitHub 93793a645b use cl.cl_mem instead of internal ctypes._CData (#14403)
fixed `CHECK_OOB=0 DEV=CL TYPED=1 python test/test_tiny.py`
2026-01-28 10:56:41 -05:00
chenyuandGitHub a9b44070a8 fix webgpu runtime types (#14402)
`CHECK_OOB=0 DEV=WEBGPU TYPED=1 python test/test_tiny.py` passed, also skip tests that failed locally
2026-01-28 10:37:25 -05:00
George HotzandGitHub 0c6b3f50aa add marker to llama training (#14401) 2026-01-28 22:44:28 +08:00
Jakob SachsandGitHub 2b7c00d3d2 fix sd-example dtype for CLIP embeddings (#14397) 2026-01-28 09:07:19 -05:00
qazalandGitHub a5a9ce3fdf viz: disasm cleanups from null emulate (#14399)
* it's AMDHIPRenderer

* don't need that indent

* less assignment stuff

* that arg order did not make sense

* pmc
2026-01-28 22:03:30 +09:00
nimlgenandGitHub 544928766d hcq_smi: kill mac pids (#14398) 2026-01-28 15:00:28 +03:00
George HotzandGitHub 202b74b369 assembly/amd: continue refactors (#14386)
* simpler

* merge

* flat

* no ctx

* use the correct apis

* dup code

* write clean code

* remove bad helpers

* bits junk remove

* junk remove

* smem test

* fix tests

* correct fix + tests

* Fmt matters it seems

* wmma refactor

* a lil more

* kimi cleanups

* line
2026-01-28 17:33:03 +08:00
qazalandGitHub 5bffa17f82 llama train: better NULL=1 EMULATE=AMD_CDNA4 dev experience (#14395)
* beam opens devices

* switch to hip renderer

* amd: true?

* llvm true is for test_autogen
2026-01-28 17:31:22 +09:00
qazalandGitHub 0294014108 fix bufferize cost function for multi, improve VIZ=-1 cli (#14394)
* improve cli

* remove_bufferize change
2026-01-28 15:53:18 +09:00
qazalandGitHub c158acea29 failing multi ram usage test from llama gemm (#14392) 2026-01-28 14:32:32 +09:00
sirhcmandGitHub 067e27857e nested composite actions don't work (#14393) 2026-01-28 00:13:30 -05:00
sirhcmandGitHub 9dddf3d478 don't save caches for PRs, try 2 (#14391) 2026-01-27 23:30:17 -05:00
sirhcmandGitHub 68fe5d8b36 Revert "don't save caches for PRs (#14389)" (#14390) 2026-01-27 23:22:26 -05:00
sirhcmandGitHub 4ab228b498 don't save caches for PRs (#14389) 2026-01-27 23:21:31 -05:00
sirhcmandGitHub 5e36482314 decompose long to ints where unsupported, try 2 (#14383) 2026-01-27 23:20:43 -05:00
wozeparrotandGitHub e496547720 llama3 gradacc (#14291) 2026-01-27 19:48:10 -08:00
George HotzandGitHub 88bc5ee212 assembly/amd: rename to better names (#14384)
* assembly/amd: rename to better names

* might help fuzzing segfault

* emu2 -> emu
2026-01-28 10:00:54 +08:00
George HotzandGitHub 065b95cfb0 Revert "add retry to fetch (#14370)" (#14385)
This reverts commit dc4d7f2d55.
2026-01-28 09:35:37 +08:00
Eitan TurokandGitHub dc4d7f2d55 add retry to fetch (#14370) 2026-01-27 14:04:25 -08:00
chenyuandGitHub 8d1f3c8885 fix copysign for inf input (#14381)
* fix copysign for inf input

* llvm olt
2026-01-27 16:45:48 -05:00
sirhcmandGitHub 289a3e415e also skip test_nonoverlapping_shrink_assignment (#14382) 2026-01-27 16:26:26 -05:00
sirhcmandGitHub f34efc1ad1 DISABLE_FAST_IDIV actually works as a ContextVar (#14378) 2026-01-27 16:12:42 -05:00
chenyuandGitHub 8c899e4aaf fix copysign for -0 (#14380)
test both x and 1/x < 0 work too. and found another big with the * 0 hack
2026-01-27 15:44:58 -05:00
chenyuandGitHub 62884585a7 failed test case for copysign -0.0 (#14379)
* failed test case for copysign -0.0

* skip those
2026-01-27 14:37:17 -05:00
nimlgenandGitHub ec1b28bc2c am: exit early in case of failures (#14376)
* am: exit early in case of failures

* sorry, pre-linter

* reset when error state
2026-01-27 22:10:02 +03:00
chenyuandGitHub cd22ee9ed0 add InvalidType to ConstType [pr] (#14373)
* add InvalidType to ConstType [pr]

TYPED=1 python test/test_tiny.py passes.
added PyConst = float|int|bool for some Tensor level input types

* hcq
2026-01-27 14:09:34 -05:00
sirhcmandGitHub 5b42a1357b SCACHE=0 works with DEBUG (#14377) 2026-01-27 13:12:43 -05:00
chenyuandGitHub db010a31be IGNORE_OOB -> CHECK_OOB [pr] (#14374)
flip the meaning
2026-01-27 12:20:59 -05:00
chenyuandGitHub c22667b0c4 also skip test_overlapping_shrink_assignment_reverse (#14375)
crashing
2026-01-27 12:20:39 -05:00
nimlgenandGitHub e52d58b041 autogen: update amd (#14372) 2026-01-27 19:53:14 +03:00
nimlgenandGitHub cbf94a0a95 nv: exit early in case of failures (#14363)
* nv: exit early in case of failures

* f

* cleaner
2026-01-27 19:16:22 +03:00
nimlgenandGitHub ec691cb299 am: print sq intrs (#14366)
* am: print sq intrs

* cleaner
2026-01-27 18:28:13 +03:00
qazalandGitHub a5f3d46423 hcq: do not assume kernel names are unique (#14371)
* hcq: do not assume kernel names are unique

* colored kernel name
2026-01-27 23:03:15 +09:00
George HotzandGitHub e5df7e640b fix branches in amd_asm_matmul (#14369) 2026-01-27 20:48:42 +08:00
geohot 0ced258726 HOTFIX: skip crashing assign test 2026-01-27 20:35:17 +08:00
George HotzandGitHub 131ae604de force_transcendental on sqrt (#14368) 2026-01-27 20:24:41 +08:00
14574c68fa Add ContextVar to disable the scheduler cache (#14257)
* add scheduler cache ContextVar

* test scheduler cache context var

---------

Co-authored-by: George Hotz <[email protected]>
2026-01-27 19:55:29 +08:00
George HotzandGitHub bfc88bcfb8 assembly/amd: emu refactors + enable PYTHON_REMU by default (#14361)
* assembly/amd: start refactors

* cleanups

* those are global

* methods on ctx

* const cleanup

* range helper

* types and imports

* cleanups

* cleanups

* remove stale name

* fix emu2 types

* more typing

* more mypy

* cleanups

* fxns

* scc cleanup

* cleanups

* cleanups

* simpler parse_pcode

* laneid

* no defaults for pcode

* pcode is not optional

* cleanups

* functions cleanup

* splat

* expr_parser functions

* single tok

* invert global loops

* try_eat

* minor

* run parser on all

* no silent 0

* tests
2026-01-27 17:42:24 +08:00
sirhcmandGitHub 2e72625652 Revert "decompose dtypes.long to ints where unsupported (#14261)" (#14362) 2026-01-27 02:04:59 -05:00
qazalandGitHub f866b2a513 mfma loop in asm dsl (#14349)
* mfma loop in asm dsl

* work
2026-01-27 11:11:37 +09:00
sirhcmandGitHub 0793319929 decompose dtypes.long to ints where unsupported (#14261)
* add works

* use carry not overflow

* bitwise ops

* use tag instead of vec

* cleaner

* mul somewhat works

* mul actually works

* SUB and NEG work

* SHL/SHR

* ulong support

* this should work?

* oops

* fix indexing

* all ALU mostly works

* refactor

* test_dtype passing

* signed division works

* format

* clean

* some tests

* ruff
2026-01-26 18:34:13 -05:00
wozeparrotandGitHub a987a4abc3 feat: llama8b dev_beam.sh (#14358) 2026-01-26 14:51:23 -08:00
sirhcmandGitHub c9c533fc78 libclang path is homebrew on macos (#14357)
* libclang path is homebrew macos

* typo

* ugh

* typo

* regen

* no LIBCLANG_PATH
2026-01-26 17:32:09 -05:00
chenyuandGitHub d641e63189 improve min/max for AND (#14356) 2026-01-26 15:44:18 -05:00
chenyuandGitHub f16372487a fix assign hazard on shrink (#14355)
* fix assign hazard on shrink

possible to have race if both assign src and dest are shrink

* test_nonoverlapping_shrink_assignment
2026-01-26 14:46:30 -05:00
chenyuandGitHub 145df879c1 find_permutes -> fix_assign_hazard [pr] (#14354)
some noop tweaks and comment updates
2026-01-26 14:05:19 -05:00
nimlgenandGitHub e152f1b0f5 llama: use ALL2ALL (#14353) 2026-01-26 22:01:53 +03:00
nimlgenandGitHub 3f25eb3026 am: ih (#14346)
* am: ih

* um

* fix

* line

* no trap and fix ring

* keep

* fix
2026-01-26 20:11:04 +03:00
chenyuandGitHub 823bc17fb5 failed test case for shrink overlap assigns (#14350)
* failed test case for shrink overlap assigns

current logic can create a race resulted in wrong output

* skip for now
2026-01-26 11:58:45 -05:00
George HotzandGitHub 204f51e739 assembly/amd: bug fixes for PYTHON_REMU (#14347)
* default PYTHON_REMU to 1

* mockgpu

* less size

* normal compile path

* uniqie

* more

* fix clamp

* Change PYTHON_REMU default to 0 in _try_dlopen_remu
2026-01-27 00:48:22 +08:00
chenyuandGitHub 231305603d remove REAL_DEV [pr] (#14337)
it's just Device.DEFAULT now
2026-01-26 10:08:16 -05:00
Martin SzewieczekandGitHub 9cbe99348a func meshgrid: change param index to type str (#14331) 2026-01-26 10:07:56 -05:00
George HotzandGitHub 3b43d26f10 assembly/amd: emu speed (#14344)
* assembly/amd: emu speed

* fix spec

* go

* don't do this

* simpler

* no stupid consts

* hack

* simpler

* no index

* no where

* faster linearizer

* fix spec

* no index dtype
2026-01-26 22:21:34 +08:00
George HotzandGitHub 774a454bb5 assembly/amd: fix scratch SVE (#14340)
* assembly/amd: default python REMU

* mem_used

* no lane

* sve

* remove that

* needs s_code_end in tests
2026-01-26 21:03:51 +08:00
qazalandGitHub 2d91fe6310 use amdgpu dsl in mmapeak (#14342)
* use amdgpu dsl in mmapeak

* don't rely on llvm for vgpr counting

* llvm roundtrip assert

* rm it, add ci

* vgpr_count

* move emulated test to amd, it needs comgr

* env

* arch

* inst._fields -> inst.operands

* vgpr offset
2026-01-26 22:03:43 +09:00
qazalandGitHub b2e2ace85b viz: remove ci check, it's VIZ=-1/-2 (#14343) 2026-01-26 20:36:23 +09:00
George HotzandGitHub be23776ba7 assembly/amd: replace pcode with ucode (#14002)
* a bunch of todos for my boy claude

* uops have types

* lil cleanups

* simpler ucode

* isNAN

* calls

* move more

* cleanup pcode_parse

* cvt functions

* fix parser bugs

* no void

* minmax

* more pcode parse

* pretty print

* transform

* comments

* move to transform

* assign/declare

* simpler norm

* single PM

* just Uops

* simpler

* more typed

* all rewrite

* less verbose

* work

* spec

* transform

* work

* simpler spec

* less spec

* bitcast

* simpler

* simp ucode

* work

* more in pcode_transform

* remove junk

* more functions

* bug

* no void assign

* load/store

* wave

* fixes

* move denorm

* move more functions

* tests

* cat is shape None

* uop syntax

* move a few more

* program_spec

* cat stuff

* assign fix clear

* unused

* nans

* fp bits

* works with simplify

* remove junk

* special

* meh

* more

* more

* update test pcode parse

* improve parser

* parse some for loops

* merge master

* dead files

* tests pass

* emu2

* better emu2

* test_plus works

* uselessly write more instructions

* use pcode

* something

* something

* bench_emu

* progress

* ds works

* work

* work

* more passing

* run compare

* bench_emu

* more pcode

* a few more

* bugfixes

* bugfix

* test fixes

* tests pass without USE_HW

* all hw tests pass

* add more hw tests

* new hw tests

* bit

* less handcode

* parse more

* consolidate pcode

* fixes

* rsrc

* lane pcode

* cleanups

* simpler

* emu bugs

* one cmp test fails

* fix decode and upd name

* fix name and test harness

* _ftz_f32

* fix denorm

* fix VOPD and use load

* fix carry bug

* no load where / just invalid

* clean

* simpler

* merge sops

* refactoring

* simplifications

* bugfixes

* new tests

* f16 sin fix

* assertion and hw tests

* cvt functions

* one more failure

* bugfixes

* bugfix + regression

* more tests

* fmac

* no manual unrolling

* ordering

* LLVM backend is a lot faster

* compile inst

* more bugs

* f16

* bugfix

* fix regression

* one clang call

* 1M inst

* scratch works

* do scratch correctly

* cleanup

* regression

* cmp

* fmamk fixes

* merge

* fix vcmpx

* unify memory

* remove unused code

* ignore oob for test

* cleanups

* fix mbs

* unify cmp

* test

* minor cleanups

* bump timeout

* fix tests

* revert the CMPLE stuff

* remove opt

* less diff

* simpler

* revert

* support multiple backends

* memset is a lot faster

* split out in bench emu

* improve timing

* timing

* cache that

* cache that

* simpler and faster

* tokenize

* binop table

* simpler

* move to parser

* tok for lambda

* refactor

* expr_parser

* delete emu2_pcode

* import cleanup

* lil

* if parse

* work

* simpler

* no v

* trig preop is faster

* durations for tests

* fix cmp bug

* sdst

* remove scartch_size hack

* null behavior

* _MXCSRContext

* bugfixes

* DEBUG >= 3

* test smem crashes my gpu

* debug

* test

* test smem

* profiler

* full inst

* bugfix

* rtag(1)

* pc is 64-bit and word

* pc is real code now

* dynamic

* more dynamic

* fix oob access

* fix crash, more dyn

* all dyn

* really all dyn

* correct null mask

* lit + format

* 21s on the tests

* 13s on the tests

* canonical name

* simm16

* more dyn

* 14s

* proper saddr dedup

* dyn

* debug 5

* better 5

* revert dynamic stuff

* that can be dyn

* negative offsets

* dyn wmma

* f16 wmma support / ops / dtype / dtype_alu

* symbolic changes not needed

* ConstFloat

* more uop.const

* __eq__

* uop tests

* fix f16

* bf16 tensor cores

* whitespace

* remove cast roundtrip

* Revert "remove cast roundtrip"

This reverts commit c5bb0381c3.

* just the fix

* remove dead paths

* llvm runs
2026-01-26 18:04:29 +08:00
George HotzandGitHub 984cdc4840 add wrapper class for the -0.0 != 0.0 issue (#14339)
* add wrapper class for the -0.0 != 0.0 issue

* fixes

* spec fix

* missed one
2026-01-26 16:52:37 +08:00
qazalandGitHub 92bfe92138 assembly/amd: fix cdna mfma xml (#14329)
* handwritten failing test

* new amdxml

* more mfma from fixes

* ci

* move arch of test integration

* alt

* amdxml human cleanup

* _TestIntegration rename to IntegrationTestBase

* it's the same problem as _LIT

* better comment

* better variable name
2026-01-26 17:51:26 +09:00
6c109f4d75 LLVM: CPU threading support (#14320)
* make generic llvmrenderer class for cpu and amd

* move `tensor_cores` back to parent

* remove empty line

* restore extra matcher position

* add threading

* dont need to add core_id here

* dont move code for workitem

* cleanup

---------

Co-authored-by: TheVanadium <[email protected]>
2026-01-26 13:12:39 +08:00
George HotzandGitHub cc49e47ea2 tinygrad changes from ucode (#14336)
* tinygrad changes from ucode

* dtype
2026-01-26 11:30:18 +08:00
8477368d07 generic LLVMRenderer class for CPU and AMD (#14321)
* make generic llvmrenderer class for cpu and amd

* move `tensor_cores` back to parent

* remove empty line

* restore extra matcher position

* cleanup

---------

Co-authored-by: TheVanadium <[email protected]>
2026-01-26 09:11:49 +08:00
geohot 11ce1e847d llama train: null device support 2026-01-26 08:53:05 +08:00
chenyuandGitHub e3601788fa update torch backend function (#14333)
those have tensor.py implementation
2026-01-25 16:39:34 -05:00
nimlgenandGitHub 9865f51e39 cupti: ref collector (#14330)
* cupti: ref collector

* ll
2026-01-25 20:35:21 +03:00
nimlgenandGitHub 21ab23ae18 nv: add pma for ada (#14328)
* nv: add pma for ada

* um

* fix

* shorter

* mock
2026-01-25 17:33:37 +03:00
George HotzandGitHub 49db266b96 ReprEnum for repr roundtrips (#14327)
* ReprEnum for repr roundtrips

* dsl

* bugfixes

* vdsty fixes

* cleaner

* fix

* fix cdna fields

* tests all pass
2026-01-25 18:58:31 +08:00
qazalandGitHub bf2d9d138f viz: simplify amdgpu cfg (#14326)
* viz: replace llvm disasm with our disasm

* it starts with more code

* then it becomes less

* simpler, cdna disassembles with decimal simm16

* s_branch is upper case, add test

* simm16s and others
2026-01-25 15:21:45 +09:00
qazalandGitHub 647e527a7e viz: replace llvm disasm with our disasm (#14325) 2026-01-25 13:56:56 +09:00
nimlgenandGitHub 4280a8eef2 am: update fw (#14323) 2026-01-25 01:08:47 +03:00
chenyuandGitHub 7e41da1ae8 fix generate_dataset.sh (#14324)
added `set -e` so wrong pathes would fail the script, then fixed the path
2026-01-24 16:47:10 -05:00
chenyuandGitHub 311bfd91d6 clean up where_on_load [pr] (#14322)
no repeated split_uop and general cleanup
2026-01-24 14:43:43 -05:00
nimlgenandGitHub 8b282ba6d2 memory: reserved vram (#14318) 2026-01-24 19:39:24 +03:00
chenyuandGitHub 00e9ba0b82 update type for split_uop and where_on_load [pr] (#14319)
also variable names in where_on_load, before logic update
2026-01-24 11:17:41 -05:00
chenyuandGitHub cb69b7b2b2 comment out fold_where_closure (#14316) 2026-01-24 10:15:42 -05:00
wozeparrotandGitHub d74587f16d fa multi fix 2 (#14314) 2026-01-23 23:35:02 -08:00
chenyuandGitHub d9f0ad1d87 update return type for Tensor.tolist (#14313)
since sequence is incorrect since it can be list of list, use Any to avoid recursive type
2026-01-23 23:21:49 -05:00
qazalandGitHub 807bc40931 assembly/amd: dsl and disasm cleanup (#14311)
* rdna4 inst helper

* remove dsl aliases
2026-01-24 11:36:12 +09:00
sirhcmandGitHub e782d44918 WEBGPU/NIR truncates ints (#14307)
* WEBGPU truncates ints

* nir has this bug too
2026-01-23 19:28:06 -05:00
nimlgenandGitHub 26220a472e no core_id (#14265)
* no core_id

* kwargs

* est

* linters

* ugh

* revert this

* deps

* glb

* should work?

* nn

* line

* fx

* ym

* z

* d

* um?

* revert

* this one?

* first half

* um p2

* all?

* um

* cleaner

* um
2026-01-23 21:30:12 +03:00
chenyuandGitHub e65bc7a7c5 where closure folding (#14304) 2026-01-23 10:55:13 -05:00
chenyuandGitHub d5a3b02a9c clean up xpow (#14295)
mostly for `ret * (base < 0).where(adj, ret.const_like(1))` -> `(base < 0).where(neg_base, ret)`, since it's good for NAN neg_base but not generic
2026-01-23 10:19:47 -05:00
b913c910c5 assembly/amd: rdna4 passing test_roundtrip (#14300)
* test_roundtrip on different archs

* failing tests

* take RDNA4 xml changes from the emu branch

* work

* min diff to disasm flat

* test_add passes, rdna4 first

* correct vgpr field for the multi dword store stuff

* amdllvm

* recompile in roundtrip, get sources from emulator

* amdllvm, 2

* clean clean

* note, don't rely on that os.environ

---------

Co-authored-by: George Hotz <[email protected]>
2026-01-23 21:33:53 +09:00
qazalandGitHub f3b0e42863 remove extra sqtt pickles in gfx1200 (#14302) 2026-01-23 20:13:48 +09:00
George HotzandGitHub d116312b1a get cdna sqtt working (#14301)
* get cdna sqtt working

* cnd aprser

* wavestart/waveend

* names

* cdna

* test that
2026-01-23 18:46:15 +08:00
George HotzandGitHub a5c4fa39d1 RDNA4 support in SQTT (#14299)
* table test

* cleanups

* dead file

* delta short

* tests

* delta test

* work

* l4 tests pass

* l0

* cnda

* print

* reverT

* wave failure

* wave failure

* test

* encs

* no l0 crap

* L4

* rdna4 sqtt

* notes

* linter
2026-01-23 16:16:45 +08:00
wozeparrotandGitHub 963c59ebdb fix: pull fixes from gradacc branch (#14296) 2026-01-22 23:07:54 -08:00
sirhcmandGitHub 68668b8f28 fix WEBGPU NEG (#14298)
* fix WEBGPU NEG

* add test

* parenthesize
2026-01-23 01:44:52 -05:00
qazalandGitHub 3b8a7bb8c9 use existing roc.py infra for sqtt tests (#14297)
* add pc, per kernel tracing

* work

* remove those imports

* min diff
2026-01-23 14:07:11 +09:00
chenyuandGitHub 5f32f7a06b fix winograd padding order (#14294) 2026-01-22 23:00:14 -05:00
George HotzandGitHub 52b989c6c8 don't place consts early + fixes from anthropic challenge (#14286)
* don't place consts early

* add anthropic challenge

* with ref

* do we still have to devectorize bools?

* tests pass

* just WHERE

* fine, revert that

* fine, revert

* only index

* z3 validator doesn't support vectorized

* Revert "z3 validator doesn't support vectorized"

This reverts commit 1b7930ecb3.

* z3 not for vec

* no spec

* VLIWRenderer

* loop unrolling

* better comments

* cleanups

* skip cast

* renderer

* cleanups

* prints

* no hack

* hacks

* bump to 11

* reg warning

* lil clean

* cleaner renderer
2026-01-23 10:48:39 +09:00
chenyuandGitHub 0903782bc0 remove few dead or unneeded codes [pr] (#14275) 2026-01-22 20:05:43 -05:00
chenyuandGitHub 3eb5cd7d32 stronger test_rand_is_lazy (#14293) 2026-01-22 18:58:53 -05:00
chenyuandGitHub c15b6e6709 update test_randn_finite skipped device (#14292) 2026-01-22 18:26:02 -05:00
chenyuandGitHub 073c6a81b5 raise if Tensor._buffer is called during jit (#14114)
* raise if Tensor._buffer is called during jit

* cleaner
2026-01-22 17:30:18 -05:00
nimlgenandGitHub 8cd22df2dd amd: alive wgps (#14149)
* amd: disabled wgps

* l

* wgp

* uoops

* mockgpu

* drm

* ad this

* fi

* reg
2026-01-23 00:08:45 +03:00
chenyuandGitHub a738c4bb22 test symbolic view broken with jit (#14290) 2026-01-22 13:44:47 -05:00
chenyuandGitHub f22fa6a5be test rand is lazy (#14289) 2026-01-22 13:07:55 -05:00
chenyuandGitHub 1726b884f2 update test_jit_v_nojit_random_regen (#14288)
current behavior is that jit and non-jit consume random seed differently, still the random values are different
2026-01-22 12:21:47 -05:00
chenyuandGitHub fbed36fa15 jit graph handle input==output aliasing (#14287)
a position that wasn't an input during capture should never become an input during execution, but graph cannot tell this by jit_cache and input_buffers only
2026-01-22 11:37:41 -05:00
chenyuandGitHub 8bb61c2490 stronger test_graph_input_output_aliasing (#14282)
* stronger test_graph_input_output_aliasing

* comfirmed failure
2026-01-22 09:59:34 -05:00
qazalandGitHub d7afa02085 clean up the extra/sqtt directory (#14284)
* remove legacy test_timing stuff

* remove legacy test_pmc, update active_sqtt_parse
2026-01-22 19:10:59 +09:00
qazalandGitHub dff5f361b0 support rendering assembly kernels on the NULL backend (#14283)
* assembly custom kernels in DEV=NULL, use renderer arch

* update mmapeak

* llvm
2026-01-22 15:49:07 +09:00
qazalandGitHub dfefeddeed add tflops to cdna gemm custom kernel (#14281) 2026-01-22 12:48:28 +09:00
qazalandGitHub 18f408a35a custom assembly kernel with variable tests (#14280)
* custom assembly kernel with variable tests

* different threads

* sink

* zeros like / flatten
2026-01-22 11:34:17 +09:00
chenyuandGitHub 4de107b764 jit graph bug when input is output (#14278)
* jit graph bug when input is output

wrong result in llm

* not just metal
2026-01-21 18:49:52 -05:00
wozeparrotandGitHub 76a9242a66 fa: merge kv bwd into one kernel (#14277) 2026-01-21 15:24:41 -08:00
chenyuandGitHub 6279ae4a94 remove llm generate always reset start_pos (#14276)
* remove llm generate always reset start_pos

by itself seems like a bug, also added a test to repro forward_jit.reset() issue

* issue is jit graph, so revert that test
2026-01-21 16:54:30 -05:00
nimlgenandGitHub da1fedc3c8 working ioctls (#14272) 2026-01-21 20:29:04 +03:00
chenyuandGitHub 574d171fa6 fix onnx Pad constant_value=None (#14271)
also removed a dead branch in _resolve_pool_pads
2026-01-21 11:51:34 -05:00
chenyuandGitHub a18d34be1e simpler split_store outer range check [pr] (#14273)
also fixed comment
2026-01-21 11:51:14 -05:00
chenyuandGitHub e64111ad08 update all_same [pr] (#14270)
add type annotation and unit test
2026-01-21 11:26:15 -05:00
chenyuandGitHub 9ad3c865ac fix bug in logsumexp keepdim=True (#14268) 2026-01-21 09:49:55 -05:00
George HotzandGitHub 41d00a046d add device to local, fix PCONTIG=2 (#14266)
* add device to local, fix PCONTIG=2

* regression test

* remove the device when we render

* viz slowness

* no long
2026-01-21 22:12:18 +09:00
362 changed files with 49282 additions and 11693 deletions
+1 -1
View File
@@ -11,5 +11,5 @@ runs:
git fetch origin $CURRENT_SHA
export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
export CURRENT_HEAD=$(git rev-parse HEAD)
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && CHECK_OOB=0 PYTHONPATH=. python3 process_replay.py
git checkout $CURRENT_HEAD # restore to branch
+35 -12
View File
@@ -56,7 +56,15 @@ runs:
# **** Caching packages ****
- name: Cache Python packages (PR)
if: github.event_name == 'pull_request'
id: restore-venv-pr
uses: actions/cache/restore@v4
with:
path: ${{ github.workspace }}/.venv
key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
- name: Cache Python packages
if: github.event_name != 'pull_request'
id: restore-venv
uses: actions/cache@v4
with:
@@ -65,23 +73,23 @@ runs:
# **** Caching downloads ****
- name: Cache downloads (Linux)
if: inputs.key != '' && runner.os == 'Linux'
uses: actions/cache@v4
- name: Cache downloads (PR)
if: inputs.key != '' && github.event_name == 'pull_request'
uses: actions/cache/restore@v4
with:
path: ~/.cache/tinygrad/downloads/
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
- name: Cache downloads (macOS)
if: inputs.key != '' && runner.os == 'macOS'
- name: Cache downloads
if: inputs.key != '' && github.event_name != 'pull_request'
uses: actions/cache@v4
with:
path: ~/Library/Caches/tinygrad/downloads/
path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
# **** Python deps ****
- name: Install dependencies in venv (with extra)
if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
shell: bash
run: |
python -m venv .venv
@@ -92,7 +100,7 @@ runs:
fi
python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
- name: Install dependencies in venv (without extra)
if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
shell: bash
run: |
python -m venv .venv
@@ -182,8 +190,14 @@ runs:
echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
- name: Cache apt (PR)
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
uses: actions/cache/restore@v4
with:
path: /var/cache/apt/archives/
key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
- name: Cache apt
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
uses: actions/cache@v4
with:
path: /var/cache/apt/archives/
@@ -239,8 +253,17 @@ runs:
ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
- name: Cache gpuocelot (PR)
if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
id: cache-build-pr
uses: actions/cache/restore@v4
env:
cache-name: cache-gpuocelot-build-1
with:
path: ${{ github.workspace }}/gpuocelot/ocelot
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
- name: Cache gpuocelot
if: inputs.ocelot == 'true'
if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
id: cache-build
uses: actions/cache@v4
env:
@@ -249,7 +272,7 @@ runs:
path: ${{ github.workspace }}/gpuocelot/ocelot
key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
- name: Clone/compile gpuocelot
if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
shell: bash
run: |
git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
+3 -3
View File
@@ -40,13 +40,13 @@ jobs:
mesa: 'true'
pydeps: 'pyyaml mako'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
- name: Regenerate autogen files
run: |
find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
python3 -c "from tinygrad.runtime.autogen import opencl"
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
python3 -c "from tinygrad.runtime.autogen import llvm"
@@ -84,7 +84,7 @@ jobs:
- name: Regenerate autogen files
run: |
rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
- name: Check for differences
run: |
if ! git diff --quiet; then
+59 -11
View File
@@ -16,6 +16,48 @@ on:
workflow_dispatch:
jobs:
# the goal of this test is to replicate a normal person on a laptop running the test
# no process replay, no benchmarks, no CI, just a normal laptop person
# the 3 minute timeout should not be raised
testmacpytest:
name: Mac pytest
runs-on: [self-hosted, macOS]
timeout-minutes: 3
defaults:
run:
shell: bash -e -o pipefail {0}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v4
# brew install uv
- name: setup python environment
run: |
rm -rf /tmp/tinygrad_pytest_ci
uv venv /tmp/tinygrad_pytest_ci
source /tmp/tinygrad_pytest_ci/bin/activate
uv pip install .[testing]
- name: setup staging db
run: |
echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
rm -f /tmp/pytest-db-ci*
# TODO: remove this step once all old caches are migrated
- name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
run: |
cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
for old_dir in models--*; do
[ -d "$old_dir" ] || continue
repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
[ -n "$snapshot" ] || continue
mkdir -p "$repo_id"
cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
done
- name: Run pytest -nauto
run: |
source /tmp/tinygrad_pytest_ci/bin/activate
pytest -nauto --durations=20
testmacbenchmark:
name: Mac Benchmark
env:
@@ -145,6 +187,10 @@ jobs:
run: |
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: Kill stale pids
run: |
PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
- name: UsbGPU boot time
run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU tiny tests
@@ -332,9 +378,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: ./extra/hcq/hcq_smi.py amd rmmod
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
#- name: Insert amdgpu
# run: sudo modprobe amdgpu
- name: Symlink models and datasets
@@ -444,9 +490,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: ./extra/hcq/hcq_smi.py amd rmmod
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -479,6 +525,8 @@ jobs:
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Test full tinyfs load
run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
- name: Run process replay tests
run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -496,9 +544,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: ./extra/hcq/hcq_smi.py amd rmmod
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -561,7 +609,7 @@ jobs:
- name: openpilot compile3 0.10.1 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.10.1 dmonitoring
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
@@ -587,9 +635,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove amd modules
run: ./extra/hcq/hcq_smi.py amd rmmod
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
- name: Kill stale pids
run: ./extra/hcq/hcq_smi.py amd kill_pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
@@ -651,9 +699,9 @@ jobs:
- name: Setcap to python
run: ./extra/amdpci/setup_python_cap.sh
- name: Remove nv modules
run: ./extra/hcq/hcq_smi.py nv rmmod
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
- name: Kill stale pids
run: ./extra/hcq/hcq_smi.py nv kill_pids
run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
- name: Symlink models and datasets
run: |
mkdir -p weights
+80 -62
View File
@@ -1,11 +1,11 @@
name: Unit Tests
env:
# increment this when downloads substantially change to avoid the internet
CACHE_VERSION: '15'
CACHE_VERSION: '16'
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
IGNORE_OOB: 0
CHECK_OOB: 1
on:
push:
@@ -26,19 +26,19 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: llvm-speed
deps: testing_minimal
deps: testing_unit
llvm: 'true'
- name: Speed Test
run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
- name: Speed Test (BEAM=2)
run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
docs:
name: Docs
runs-on: ubuntu-22.04
timeout-minutes: 10
env:
IGNORE_OOB: 1
CHECK_OOB: 0
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -98,7 +98,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_minimal
deps: testing_unit
pydeps: "pillow torchvision expecttest"
llvm: 'true'
- name: Install ninja
@@ -134,7 +134,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_minimal
deps: testing_unit
llvm: 'true'
- name: Install ninja
run: |
@@ -156,7 +156,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: be-minimal
deps: testing_minimal
deps: testing_unit
- name: Test dtype with Python emulator
run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
- name: Test ops with Python emulator
@@ -207,11 +207,11 @@ jobs:
run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
- name: Test device flop counts
run: |
DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
linter:
name: Linters
@@ -239,9 +239,10 @@ jobs:
- name: Run mypy with lineprecision report
run: |
python -m mypy --lineprecision-report .
grep -v autogen lineprecision.txt | awk 'NR>2 {lines+=$2; precise+=$3; imprecise+=$4; any+=$5; empty+=$6} END {t=lines-empty; printf "TOTAL: %d lines, %d precise (%.1f%%), %d imprecise (%.1f%%), %d any (%.1f%%)\n", t, precise, 100*precise/t, imprecise, 100*imprecise/t, any, 100*any/t}'
cat lineprecision.txt
- name: Run TYPED=1
run: TYPED=1 python -c "import tinygrad"
run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py
unittest:
name: Unit Tests
@@ -255,26 +256,29 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: unittest-13
pydeps: "pillow numpy ftfy regex pre-commit"
pydeps: "pillow ftfy regex pre-commit"
deps: testing_unit
llvm: 'true'
amd: 'true'
- name: Run pre-commit test hooks
run: SKIP=ruff,mypy pre-commit run --all-files
- name: Check Device.DEFAULT
run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
- name: Run unit tests
run: |
CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
CPU=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Run NULL backend tests
run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
- name: Run targetted tests on NULL backend
run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
# TODO: too slow
# - name: Run SDXL on NULL backend
# run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
- name: Run Clip tests for SD MLPerf on NULL backend
run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
- name: Run AMD emulated BERT training on NULL backend
run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
# TODO: support fake weights
#- name: Run LLaMA 7B on 4 fake devices
# run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -312,7 +316,7 @@ jobs:
deps: testing_unit
python-version: '3.14'
- name: Test SPEC=2
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --ignore test/unit/test_autogen.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore=test/null --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
fuzzing:
name: Fuzzing
@@ -346,7 +350,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: gpu-image
deps: testing_minimal
deps: testing_unit
opencl: 'true'
- name: Test CL IMAGE=2 ops
run: |
@@ -367,7 +371,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: gen-dataset
deps: testing_minimal
deps: testing
opencl: 'true'
- name: Generate Dataset
run: CL=1 extra/optimization/generate_dataset.sh
@@ -422,7 +426,7 @@ jobs:
with:
key: onnxoptc
deps: testing
python-version: '3.11'
python-version: '3.12'
llvm: 'true'
- name: Test ONNX (CPU)
run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -450,7 +454,7 @@ jobs:
key: onnxoptl
deps: testing
pydeps: "tensorflow==2.19"
python-version: '3.11'
python-version: '3.12'
opencl: 'true'
- name: Test ONNX (CL)
run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -463,11 +467,11 @@ jobs:
- name: Test MLPerf stuff
run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
- name: NULL=1 beautiful_mnist_multigpu
run: NULL=1 python examples/beautiful_mnist_multigpu.py
run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
- name: Test Bert training
run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Test llama 3 training
run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -476,7 +480,7 @@ jobs:
runs-on: ubuntu-24.04
timeout-minutes: 15
env:
IGNORE_OOB: 1
CHECK_OOB: 0
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -524,7 +528,7 @@ jobs:
with:
key: metal
deps: testing
python-version: '3.11'
python-version: '3.12'
- name: Test models (Metal)
run: METAL=1 python -m pytest -n=auto test/models --durations=20
- name: Test LLaMA compile speed
@@ -543,7 +547,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: devectorize-minimal
deps: testing_minimal
deps: testing_unit
pydeps: "pillow"
llvm: "true"
- name: Test LLVM=1 DEVECTORIZE=0
@@ -564,8 +568,8 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: dsp-minimal
deps: testing_minimal
pydeps: "onnx==1.18.0 onnxruntime pillow"
deps: testing_unit
pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
llvm: "true"
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
@@ -577,7 +581,7 @@ jobs:
load: true
tags: qemu-hexagon:latest
cache-from: type=gha
cache-to: type=gha,mode=min
cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
- name: Set MOCKDSP env
run: printf "MOCKDSP=1" >> $GITHUB_ENV
- name: Run test_tiny on DSP
@@ -598,8 +602,8 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: webgpu-minimal
deps: testing_minimal
python-version: '3.11'
deps: testing_unit
python-version: '3.12'
webgpu: 'true'
- name: Check Device.DEFAULT (WEBGPU) and print some source
run: |
@@ -607,7 +611,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 --durations=20
WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -632,7 +636,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: ${{ matrix.backend }}-minimal
deps: testing_minimal
deps: testing_unit
amd: 'true'
llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
- name: Check Device.DEFAULT and print some source
@@ -650,16 +654,23 @@ jobs:
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
- name: Run TestOps.test_add with SQTT
run: |
VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
VIZ=-2 DEBUG=5 python3 test/test_ops.py TestOps.test_add
extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
- name: Run AMD emulated mmapeak on NULL backend
env:
AMD: 0
run: PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdasm:
name: AMD ASM IDE
runs-on: ubuntu-24.04
timeout-minutes: 10
timeout-minutes: 20
env:
AMD: 1
PYTHON_REMU: 1
MOCKGPU: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -667,12 +678,12 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: rdna3-emu
deps: testing_minimal
deps: testing_unit
amd: 'true'
python-version: '3.13'
python-version: '3.14'
- name: Verify AMD autogen is up to date
run: |
python -m extra.assembly.amd.amdxml
python -m extra.assembly.amd.generate
git diff --exit-code extra/assembly/amd/autogen/
- name: Install LLVM 21
run: |
@@ -685,16 +696,18 @@ jobs:
- 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
run: AMD_LLVM=0 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 dtype tests
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
run: AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
run: AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
# 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"
run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
- name: Run RDNA4 emulator tests
run: MOCKGPU_ARCH=rdna4 python -m pytest test/test_tiny.py -v --durations 20
testnvidia:
strategy:
@@ -715,7 +728,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: ${{ matrix.backend }}-minimal
deps: testing_minimal
deps: testing_unit
cuda: 'true'
ocelot: 'true'
- name: Set env
@@ -726,7 +739,9 @@ jobs:
DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run pytest (cuda)
# skip multitensor because it's slow
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_multitensor.py --durations=20
- name: Run TestOps.test_add with PMA
run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -746,7 +761,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: ${{ matrix.backend }}-minimal
deps: testing_minimal
deps: testing_unit
opencl: ${{ matrix.backend == 'opencl' && 'true' }}
llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
mesa: ${{ matrix.backend == 'lvp' && 'true' }}
@@ -757,7 +772,7 @@ jobs:
python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
- name: Run pytest (${{ matrix.backend }})
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
- name: Run TRANSCENDENTAL math
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
- name: Run process replay tests
@@ -777,15 +792,15 @@ jobs:
with:
key: metal
deps: testing
python-version: '3.11'
python-version: '3.12'
amd: 'true'
cuda: 'true'
ocelot: 'true'
llvm: 'true'
- name: Run unit tests
env:
LIBCLANG_PATH: '/opt/homebrew/opt/llvm@20/lib/libclang.dylib'
run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Run NULL backend tests
run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
- name: Run ONNX
run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
- name: Test tensor core ops (fake)
@@ -877,8 +892,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: macos-${{ matrix.backend }}-minimal
deps: testing_minimal
pydeps: "capstone"
deps: testing_unit
llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
mesa: ${{ matrix.backend == 'lvp' && 'true' }}
- name: Set env
@@ -888,7 +902,7 @@ jobs:
python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
- name: Run pytest (${{ matrix.backend }})
run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --durations=20
- name: Run process replay tests
uses: ./.github/actions/process-replay
- name: Run macOS-specific unit test
@@ -921,7 +935,11 @@ jobs:
- name: Run unit tests
if: matrix.backend=='llvm'
# test_newton_schulz hits RecursionError
run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
- name: Run NULL backend tests
if: matrix.backend=='llvm'
shell: bash
run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
- name: Run pytest (${{ matrix.backend }})
shell: bash
run: |
@@ -945,12 +963,12 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: compile-${{ matrix.backend }}
deps: testing_minimal
deps: testing_unit
mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
python-version: '3.14'
python-version: '3.12'
- name: Set env
shell: bash
run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
- name: Run test_ops
shell: bash
run: |
+1 -1
View File
@@ -28,7 +28,7 @@ repos:
pass_filenames: false
- id: tests
name: comprehensive test suite
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/unit/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
language: system
always_run: true
pass_filenames: false
+196
View File
@@ -0,0 +1,196 @@
from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
from tinygrad.uop.ops import Ops, PatternMatcher, UPat
from tinygrad.uop.symbolic import symbolic
from tinygrad.codegen import Renderer
from tinygrad.codegen.opt import Opt, OptOps
# ************************* implementation of the problem ************************
def myhash(a: Tensor) -> Tensor:
a = (a + 0x7ED55D16) + (a << 12)
a = (a ^ 0xC761C23C) ^ (a >> 19)
a = (a + 0x165667B1) + (a << 5)
a = (a + 0xD3A2646C) ^ (a << 9)
a = (a + 0xFD7046C5) + (a << 3)
a = (a ^ 0xB55A4F09) ^ (a >> 16)
return a
def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
n = values.shape[0]
if n == 1: return values[0].expand(relative_idx.shape)
mid = n // 2
left = select_with_where_tree(values[:mid], relative_idx)
right = select_with_where_tree(values[mid:], relative_idx - mid)
go_left = relative_idx < mid
return go_left.where(left, right)
def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
# All walkers start at idx=0
idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
for r in range(rounds):
level = r % (height + 1)
level_start = (1 << level) - 1
level_size = 1 << level
if level == 0:
# At root (level 0), all walkers are at idx=0
# No gather needed, just broadcast the root value
node_val = forest[0].expand(val.shape)
idx = idx * 0 # Reset to 0
elif level <= where_tree_threshold:
# Small level: use where-tree
level_values = forest[level_start : level_start + level_size]
relative_idx = (idx - level_start)
node_val = select_with_where_tree(level_values, relative_idx)
else:
# Large level: use gather
node_val = forest.gather(0, idx)
val = myhash(val ^ node_val)
idx = (idx << 1) + (1 + (val & 1))
# No wrap check needed! At round 10 (level becomes 0), we reset idx above.
return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
# ************************* renderer for VLIW machine *************************
def loop_unrolling(sink:UOp):
rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
if len(rng) == 0: return None
print(f"unrolling loop with size {rng[0].vmax+1}")
unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
return UOp.sink(*unrolled_sinks, arg=sink.arg)
global_addrs = []
vliw_prepare = PatternMatcher([
# loop unrolling (should be a part of tinygrad)
(UPat(Ops.SINK, name="sink"), loop_unrolling),
# cast is fake
(UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
# rewrites to hardcode the addresses in memory
(UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
# INDEX is just plus
(UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
])+symbolic
class VLIWRenderer(Renderer):
has_local = False # TODO: this should be the default / cleaned up
# this says this backend supports MULACC + more. decompositions uses this
code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
# this matcher runs while still in graph form
pre_matcher = vliw_prepare
def render(self, uops:list[UOp]):
# TODO: this is a minimal renderer. for low cycle count, make it good
# to get speed, you need to add VLIW packing
# to get under 1536 regs, you need to add a register allocator
# we left the fun parts to you
print(f"rendering with {len(uops)} uops")
reg, inst = 0, []
r: dict[UOp, int] = {}
for u in uops:
assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
# dumb register allocator
if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
r[u] = reg
reg += u.dtype.count
# render UOps to instructions
match u.op:
case Ops.SINK:
inst.append({"flow": [("halt",)]})
case Ops.CONST:
inst.append({"load": [("const", r[u], u.arg)]})
case Ops.GEP:
# a GEP is just an alias to a special register in the vector
r[u] = r[u.src[0]] + u.arg[0]
case Ops.VECTORIZE:
if all(s == u.src[0] for s in u.src):
# if all sources are the same, we can broadcast
inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
else:
# this is a copy into a contiguous chunk of registers
inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
case Ops.LOAD:
op = "vload" if u.dtype.count > 1 else "load"
inst.append({"load": [(op, r[u], r[u.src[0]])]})
case Ops.STORE:
op = "vstore" if u.src[1].dtype.count > 1 else "store"
inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
case Ops.MULACC:
assert u.dtype.count == 8
inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
case Ops.WHERE:
assert u.dtype.count == 8
inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
case _ if u.op in self.code_for_op:
cat = "valu" if u.dtype.count > 1 else "alu"
inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
case _:
raise NotImplementedError(f"unhandled op {u.op}")
return repr(inst)
# ************************* test and render *************************
import sys, types
PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
sys.modules["problem"] = problem = types.ModuleType("problem")
exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
if __name__ == "__main__":
batch_size = getenv("BS", 256)
height = 10
rounds = getenv("ROUNDS", 16)
# build problem
tree = problem.Tree.generate(height)
inp = problem.Input.generate(tree, batch_size, rounds)
mem = problem.build_mem_image(tree, inp)
global_addrs.extend([mem[6], mem[6], mem[4]]) # output, input, forest
# *** verify the kernel in tinygrad compared to reference ***
forest_t = Tensor(tree.values, dtype=dtypes.uint32)
val_t = Tensor(inp.values, dtype=dtypes.uint32)
if getenv("VERIFY", 1):
# verify on normal tinygrad device
with Context(PCONTIG=2):
out = tree_traversal(forest_t, val_t, height, rounds)
val_out = out.tolist()
problem.reference_kernel(tree, inp)
assert val_out == inp.values
print("verification passed")
# *** render to device ***
from tinygrad.codegen import get_program
with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
out = tree_traversal(forest_t, val_t, height, rounds)
sink = out.schedule()[-1].ast
prg = get_program(sink, VLIWRenderer())
# *** run on Machine and compare ***
# NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
src = eval(prg.src)
max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else " <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
machine.run()
print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else " <-- EVEN CLAUDE GOT 1363"))
# compare to reference
ref_mem = mem.copy()
for _ in problem.reference_kernel2(ref_mem, {}): pass
assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
print("compare passed!")
+8 -8
View File
@@ -72,7 +72,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
#storage_tensor._copyin(img_tensor.numpy())
# faster
X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
# ideal
#X[idx].assign(img.tobytes()) # NOTE: this is slow!
@@ -264,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
x = random_brightness_augmentation(x)
x = gaussian_noise(x)
X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()
queue_out.put(idx)
queue_out.put(None)
@@ -379,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
clipped_match_idxs = np.clip(match_idxs, 0, None)
clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()
imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
queue_out.put(idx)
queue_out.put(None)
+96 -40
View File
@@ -3,7 +3,7 @@ from pathlib import Path
import multiprocessing
from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
@@ -1292,13 +1292,17 @@ def train_llama3():
BASEDIR = config["BASEDIR"] = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
BS = config["BS"] = getenv("BS", 16)
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
assert grad_acc == 1, f"{grad_acc=} is not supported"
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
SEED = config["SEED"] = getenv("SEED", 5760)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
SMALL = config["SMALL"] = getenv("SMALL", 0)
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
EVAL_SAMPLES = config["EVAL_SAMPLES"] = getenv("EVAL_SAMPLES", 5760 if not SMALL else 1024)
MAX_STEPS = config["MAX_STEPS"] = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS))
WARMUP_STEPS = config["WARMUP_STEPS"] = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
LR = config["LR"] = getenv("LR", 8e-5 * GBS / 1152)
END_LR = config["END_LR"] = getenv("END_LR", 8e-7)
EVAL_FREQ = config["EVAL_FREQ"] = getenv("EVAL_FREQ", 46080)
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
EVAL_TARGET = config["EVAL_TARGET"] = getenv("EVAL_TARGET", 5.6)
@@ -1312,10 +1316,12 @@ def train_llama3():
opt_adamw_weight_decay = 0.1
opt_gradient_clip_norm = 1.0
opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
opt_end_learning_rate = getenv("END_LR", 8e-7)
opt_learning_rate_warmup_steps = WARMUP_STEPS
opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
opt_base_learning_rate = LR
opt_end_learning_rate = END_LR
Tensor.manual_seed(SEED) # seed for weight initialization
# ** init wandb **
WANDB = getenv("WANDB")
@@ -1365,6 +1371,12 @@ def train_llama3():
optim = AdamW(get_parameters(model), lr=0.0,
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
# init grads
for p in optim.params:
p.grad = p.zeros_like().contiguous().realize()
grads = [p.grad for p in optim.params]
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
if resume_ckpt := getenv("RESUME_CKPT"):
@@ -1377,9 +1389,7 @@ def train_llama3():
load_state_dict(scheduler, safe_load(fn), realize=False)
@TinyJit
@Tensor.train()
def train_step(model, tokens:Tensor):
optim.zero_grad()
def minibatch(tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.shard(device, 0)
@@ -1389,27 +1399,40 @@ def train_llama3():
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
loss.backward()
assert all(p.grad is g for p,g in zip(optim.params, grads))
Tensor.realize(loss, *grads)
return loss
@TinyJit
def optim_step():
for p in optim.params:
p.grad.assign(p.grad / grad_acc)
# L2 norm grad clip
# https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
# https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
if not getenv("DISABLE_GRAD_CLIP_NORM"):
total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
for p in optim.params:
total_norm += p.grad.float().square().sum()
total_norm = total_norm.sqrt().contiguous()
for p in optim.params:
p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
for g in grads:
total_norm += g.float().square().sum()
total_norm = total_norm.sqrt().contiguous().realize()
for g in grads:
g.assign((g * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype)).realize()
optim.step()
scheduler.step()
for g in grads:
g.assign(g.zeros_like().contiguous()).realize()
lr = optim.lr
loss.realize(lr)
return loss, lr
Tensor.realize(lr, *grads)
return lr
@TinyJit
@Tensor.train(False)
def eval_step(model, tokens:Tensor):
def eval_step(tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.shard(device, 0)
@@ -1436,42 +1459,73 @@ def train_llama3():
eval_dataset = None
else:
from examples.mlperf.dataloader import get_llama3_dataset
eval_dataset = get_llama3_dataset(1024 if SMALL else 5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
def get_eval_iter():
if eval_dataset is None:
return fake_data(EVAL_BS, 5760)
return fake_data(EVAL_BS, EVAL_SAMPLES)
from examples.mlperf.dataloader import iterate_llama3_dataset
return iterate_llama3_dataset(eval_dataset, EVAL_BS)
iter = get_train_iter()
num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
train_iter = get_train_iter()
i, sequences_seen = resume_ckpt, 0
step_times = []
for tokens in tqdm(iter, total=SAMPLES//GBS):
while i < MAX_STEPS:
GlobalCounters.reset()
if getenv("TRAIN", 1):
t = time.perf_counter()
loss, lr = train_step(model, tokens)
profile_marker(f"train @ {i}")
st = time.perf_counter()
stopped = False
for _ in range(grad_acc):
ist = time.perf_counter()
try: tokens = next(train_iter)
except StopIteration:
stopped = True
break
dt = time.perf_counter()
loss = minibatch(tokens)
if stopped: break
gt = time.perf_counter()
lr = optim_step()
ot = time.perf_counter()
loss = loss.float().item()
lr = lr.item()
i += 1
sequences_seen += tokens.shape[0]
et = time.perf_counter()
step_time = et - st
gbs_time = gt - st
optim_time = ot - gt
data_time = dt - ist
dev_time = step_time - data_time * grad_acc
if BENCHMARK: step_times.append(step_time)
sec = time.perf_counter()-t
if BENCHMARK: step_times.append(sec)
i += 1
sequences_seen += GBS
mem_gb = GlobalCounters.mem_used / 1e9
gflops = GlobalCounters.global_ops / 1e9 / sec
gflops = GlobalCounters.global_ops / 1e9 / dev_time
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
tqdm.write(
f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
if (fname:=getenv("LOSS_FILE", "")):
with open(fname, "a") as f:
f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
if WANDB:
wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
wandb.log({
"lr": lr, "train/loss": loss,
"train/step_time": step_time,
"train/gbs_time": gbs_time,
"train/optim_time": optim_time,
"train/dev_time": dev_time,
"train/data_time": data_time,
"train/mem": mem_gb,
"train/GFLOPS": gflops,
"train/MFU": mfu,
"train/sequences_seen": sequences_seen
})
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
tqdm.write("saving checkpoint")
@@ -1491,17 +1545,19 @@ def train_llama3():
f"epoch global_mem: {GlobalCounters.global_mem:_}")
if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
if EVAL_BS == 0: return
tqdm.write(f"evaluating after {sequences_seen} sequences")
profile_marker(f"eval @ {i}")
# run eval
eval_losses = []
eval_iter = get_eval_iter()
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
for j,tokens in tqdm(enumerate(eval_iter), total=5760//EVAL_BS):
eval_losses += eval_step(model, tokens).tolist()
if BENCHMARK and (j+1) == min(BENCHMARK, 5760//EVAL_BS):
for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
eval_losses += eval_step(tokens).tolist()
if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
return
log_perplexity = Tensor(eval_losses).mean().float().item()
@@ -1589,7 +1645,7 @@ def train_stable_diffusion():
loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
Tensor.realize(loss, out_lr)
return loss, out_lr
# checkpointing takes ~9 minutes without this, and ~1 minute with this
@TinyJit
def ckpt_to_cpu():
@@ -1628,7 +1684,7 @@ def train_stable_diffusion():
if i == 3:
for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
total_train_time = time.perf_counter() - train_start_time
if WANDB:
wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
export IGNORE_OOB=1
export CHECK_OOB=0
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
@@ -5,7 +5,7 @@ export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000
export BEAM=0 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export IGNORE_OOB=1
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -0,0 +1,36 @@
#!/usr/bin/env bash
export PYTHONPATH="."
export DEV=${DEV:-AMD}
export EMULATE="AMD_CDNA4"
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEBUG=${DEBUG:-2}
export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
export ALL2ALL=${ALL2ALL:-1}
export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
export GBS=$((BS * GRADIENT_ACC_STEPS))
export MODEL="llama3"
export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}
export SEED=${SEED:-5760}
export JITBEAM=${JITBEAM:-3}
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2
python3 examples/mlperf/model_train.py
@@ -1,26 +1,34 @@
#!/usr/bin/env bash
export PYTHONPATH="." AMD=1
export IGNORE_OOB=1
export PYTHONPATH="."
export DEV=${DEV:-AMD}
export EMULATE="AMD_CDNA4"
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEBUG=${DEBUG:-0}
export FLASH_ATTENTION=1
export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
export ALL2ALL=${ALL2ALL:-1}
export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=8 BS=8 EVAL_BS=8
export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
export GBS=$((BS * GRADIENT_ACC_STEPS))
export MODEL="llama3"
export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8b"}
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="1e-3" END_LR="1e-4" WARMUP_STEPS=1024 MAX_STEPS=1200000
export SAMPLES=$((MAX_STEPS * BS))
export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}
export SEED=5760
export SEED=${SEED:-5760}
export JITBEAM=3
export JITBEAM=${JITBEAM:-3}
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
python3 examples/mlperf/model_train.py
@@ -0,0 +1,6 @@
#!/bin/bash
export BENCHMARK=5
export EVAL_BS=0
export VIZ=${VIZ:--1}
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -0,0 +1,10 @@
#!/bin/bash
export BENCHMARK=5
export EVAL_BS=0
export FAKEDATA=1
export NULL_ALLOW_COPYOUT=1
export HIP_VISIBLE_DEVICES=""
export DEV=NULL
export JITBEAM=0
export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+1 -1
View File
@@ -6,7 +6,6 @@ import argparse, time
from collections import namedtuple
from typing import Dict, Any
from PIL import Image
import numpy as np
from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
@@ -336,6 +335,7 @@ if __name__ == "__main__":
print(x.shape)
profile_marker("save image")
from PIL import Image
im = Image.fromarray(x.numpy())
print(f"saving {args.out}")
im.save(args.out)
@@ -48,7 +48,7 @@ def prepare_browser_chunks(model):
weight_metadata = metadata.get(name, default)
weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
metadata[name] = weight_metadata
data = bytes(state_dict[name].uop.base.realized.as_buffer())
data = bytes(state_dict[name].uop.base.realized.as_memoryview())
data = data if not offsets else data[offsets[0]:offsets[1]]
writer.write(data)
cursor += size
+1 -1
View File
@@ -93,7 +93,7 @@ if __name__ == "__main__":
forward: Any = None
sub_steps = [
Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
Step(name = "textModel", input = [Tensor.randint(1, 77, low=0, high=49408, dtype=dtypes.int32)], forward = model.cond_stage_model.transformer.text_model),
Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)
+4 -3
View File
@@ -92,7 +92,7 @@ class SMICtx:
self.prev_terminal_width = 0
self.prev_terminal_height = 0
remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
for k,v in self.lspci.items():
@@ -153,7 +153,8 @@ class SMICtx:
tables = {}
for dev in self.devs:
match dev.ip_ver[am.MP1_HWIP]:
case (13,0,6)|(13,0,12): table_t = dev.smu.smu_mod.MetricsTableX_t
case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_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
@@ -279,7 +280,7 @@ class SMICtx:
device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
+ [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
+ [f"MEM Usage {draw_bar((mem_used / mem_total) / 100, activity_line_width, opt_text=mem_fmt)}"] \
+ [f"MEM Usage {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
+7 -1
View File
@@ -1,12 +1,18 @@
#!/usr/bin/env python3
import os
from tinygrad.helpers import Context
from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.support.am.amdev import AMDev
if __name__ == "__main__":
gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
for gpu in gpus:
drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
amdevs = []
with Context(DEBUG=2):
for pcidev in pcidevs:
+27 -2
View File
@@ -19,8 +19,9 @@ amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
amdev.AMDev = AMDFake
from tinygrad.runtime.ops_amd import PCIIface
def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = False):
def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
register_map = register_names or {}
register_objs = register_objects or {}
def replace_register(match):
reg = match.group(1)
@@ -37,6 +38,28 @@ def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = Fal
# remove timing prefix
processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)
# decode register values into field dicts
def decode_value(match):
reg_name = match.group(1)
xcc_part = match.group(2) # "xcc=0 " or ""
val_str = match.group(3)
val = int(val_str, 16)
reg_obj = register_objs.get(reg_name)
if reg_obj is not None and reg_obj.fields:
fields = reg_obj.decode(val)
# show raw for unaccounted bits
accounted = 0
for name, (start, end) in reg_obj.fields.items():
accounted |= (((1 << (end - start + 1)) - 1) << start)
unaccounted = val & ~accounted
parts = {k: v for k, v in fields.items() if v != 0}
if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
return match.group(0)
processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
# keep only xcc=0 lines (but keep lines with no xcc at all)
if only_xcc0:
kept = []
@@ -50,16 +73,18 @@ def main():
only_xcc0 = bool(getenv("ONLY_XCC0", 0))
reg_names = {}
reg_objs = {}
dev = PCIIface(None, 0)
for x, y in dev.dev_impl.__dict__.items():
if isinstance(y, AMRegister):
for xcc, addr in y.addr.items():
reg_names[addr] = f"{x}, xcc={xcc}"
reg_objs[x] = y
with open(sys.argv[1], 'r') as f:
log_content = f.read()
processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)
with open(sys.argv[2], 'w') as f:
f.write(processed_log)
@@ -29,15 +29,16 @@ def _matches(data: bytes, cls: type[Inst]) -> bool:
# Import instruction classes for each architecture
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH)
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT,
VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_LIT as R4_VOP1_LIT,
VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT,
SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT, SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT,
SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
SMEM as R4_SMEM, DS as R4_DS, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP1_SDWA as C_VOP1_SDWA, VOP1_DPP16 as C_VOP1_DPP16,
VOP2 as C_VOP2, VOP2_LIT as C_VOP2_LIT, VOP2_SDWA as C_VOP2_SDWA, VOP2_DPP16 as C_VOP2_DPP16,
VOPC as C_VOPC, VOPC_SDWA_SDST as C_VOPC_SDWA_SDST,
VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3PX2 as C_VOP3PX2,
VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3P_MFMA as C_VOP3P_MFMA, VOP3PX2 as C_VOP3PX2,
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPK_LIT as C_SOPK_LIT, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
FLAT as C_FLAT, GLOBAL as C_GLOBAL, SCRATCH as C_SCRATCH, MUBUF as C_MUBUF)
@@ -47,8 +48,9 @@ _FORMATS = {
"rdna3": [VOPD, VOP3P, VINTERP, VOP3SD, VOP3_SDST, VOP3, DS, GLOBAL, SCRATCH, FLAT, SMEM,
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, VOPC, VOP1_SDST, VOP1, VOP1_LIT, VOP2, VOP2_LIT],
"rdna4": [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3SD, R4_VOP3_SDST, R4_VOP3, R4_DS, R4_GLOBAL, R4_SCRATCH, R4_FLAT, R4_SMEM,
R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
"cdna": [C_VOP3PX2, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
R4_SOP1, R4_SOP1_LIT, R4_SOPC, R4_SOPC_LIT, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_VOP1_LIT,
R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
"cdna": [C_VOP3PX2, C_VOP3P_MFMA, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_SOPK_LIT, C_VOPC_SDWA_SDST, C_VOPC,
C_VOP1_DPP16, C_VOP1_SDWA, C_VOP1, C_VOP2_DPP16, C_VOP2_SDWA, C_SOP2, C_VOP2, C_VOP2_LIT],
}
+21 -22
View File
@@ -1,8 +1,7 @@
# autogenerated from AMD ISA XML - do not edit
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(Enum):
class DSOp(ReprEnum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -133,7 +132,7 @@ class DSOp(Enum):
DS_READ_B96 = 254
DS_READ_B128 = 255
class FLATOp(Enum):
class FLATOp(ReprEnum):
FLAT_LOAD_UBYTE = 16
FLAT_LOAD_SBYTE = 17
FLAT_LOAD_USHORT = 18
@@ -189,7 +188,7 @@ class FLATOp(Enum):
FLAT_ATOMIC_INC_X2 = 107
FLAT_ATOMIC_DEC_X2 = 108
class GLOBALOp(Enum):
class GLOBALOp(ReprEnum):
GLOBAL_LOAD_UBYTE = 16
GLOBAL_LOAD_SBYTE = 17
GLOBAL_LOAD_USHORT = 18
@@ -252,7 +251,7 @@ class GLOBALOp(Enum):
GLOBAL_LOAD_LDS_DWORDX4 = 125
GLOBAL_LOAD_LDS_DWORDX3 = 126
class HWREG(Enum):
class HWREG(ReprEnum):
HW_REG_MODE = 1
HW_REG_STATUS = 2
HW_REG_TRAPSTS = 3
@@ -278,7 +277,7 @@ class HWREG(Enum):
HW_REG_SQ_PERF_SNAPSHOT_PC_LO = 23
HW_REG_SQ_PERF_SNAPSHOT_PC_HI = 24
class MTBUFOp(Enum):
class MTBUFOp(ReprEnum):
TBUFFER_LOAD_FORMAT_X = 0
TBUFFER_LOAD_FORMAT_XY = 1
TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -296,7 +295,7 @@ class MTBUFOp(Enum):
TBUFFER_STORE_FORMAT_D16_XYZ = 14
TBUFFER_STORE_FORMAT_D16_XYZW = 15
class MUBUFOp(Enum):
class MUBUFOp(ReprEnum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -372,7 +371,7 @@ class MUBUFOp(Enum):
BUFFER_ATOMIC_INC_X2 = 107
BUFFER_ATOMIC_DEC_X2 = 108
class SCRATCHOp(Enum):
class SCRATCHOp(ReprEnum):
SCRATCH_LOAD_UBYTE = 16
SCRATCH_LOAD_SBYTE = 17
SCRATCH_LOAD_USHORT = 18
@@ -401,7 +400,7 @@ class SCRATCHOp(Enum):
SCRATCH_LOAD_LDS_SSHORT = 41
SCRATCH_LOAD_LDS_DWORD = 42
class SMEMOp(Enum):
class SMEMOp(ReprEnum):
S_LOAD_DWORD = 0
S_LOAD_DWORDX2 = 1
S_LOAD_DWORDX4 = 2
@@ -487,7 +486,7 @@ class SMEMOp(Enum):
S_ATOMIC_INC_X2 = 171
S_ATOMIC_DEC_X2 = 172
class SOP1Op(Enum):
class SOP1Op(ReprEnum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -543,7 +542,7 @@ class SOP1Op(Enum):
S_ANDN2_WREXEC_B64 = 54
S_BITREPLICATE_B64_B32 = 55
class SOP2Op(Enum):
class SOP2Op(ReprEnum):
S_ADD_U32 = 0
S_SUB_U32 = 1
S_ADD_I32 = 2
@@ -598,7 +597,7 @@ class SOP2Op(Enum):
S_PACK_LH_B32_B16 = 51
S_PACK_HH_B32_B16 = 52
class SOPCOp(Enum):
class SOPCOp(ReprEnum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -620,7 +619,7 @@ class SOPCOp(Enum):
S_CMP_EQ_U64 = 18
S_CMP_LG_U64 = 19
class SOPKOp(Enum):
class SOPKOp(ReprEnum):
S_MOVK_I32 = 0
S_CMOVK_I32 = 1
S_CMPK_EQ_I32 = 2
@@ -643,7 +642,7 @@ class SOPKOp(Enum):
S_SETREG_IMM32_B32 = 20
S_CALL_B64 = 21
class SOPPOp(Enum):
class SOPPOp(ReprEnum):
S_NOP = 0
S_ENDPGM = 1
S_BRANCH = 2
@@ -677,7 +676,7 @@ class SOPPOp(Enum):
S_ENDPGM_ORDERED_PS_DONE = 30
S_SET_VALU_COEXEC_MODE = 31
class VOP1Op(Enum):
class VOP1Op(ReprEnum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -855,7 +854,7 @@ class VOP1Op(Enum):
V_PERMLANE32_SWAP_B32 = V_PERMLANE32_SWAP_B32_E32
V_CVT_F32_BF16 = V_CVT_F32_BF16_E32
class VOP2Op(Enum):
class VOP2Op(ReprEnum):
V_CNDMASK_B32_E32 = 0
V_ADD_F32_E32 = 1
V_SUB_F32_E32 = 2
@@ -981,7 +980,7 @@ class VOP2Op(Enum):
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
V_XNOR_B32 = V_XNOR_B32_E32
class VOP3Op(Enum):
class VOP3Op(ReprEnum):
V_CMP_CLASS_F32_E64 = 16
V_CMPX_CLASS_F32_E64 = 17
V_CMP_CLASS_F64_E64 = 18
@@ -1879,7 +1878,7 @@ class VOP3Op(Enum):
V_ADD_LSHL_U32 = V_ADD_LSHL_U32_E64
V_ADD3_U32 = V_ADD3_U32_E64
class VOP3POp(Enum):
class VOP3POp(ReprEnum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1988,11 +1987,11 @@ class VOP3POp(Enum):
V_SMFMAC_F32_32X32X32_FP8_BF8 = 126
V_SMFMAC_F32_32X32X32_FP8_FP8 = 127
class VOP3PX2Op(Enum):
class VOP3PX2Op(ReprEnum):
V_MFMA_SCALE_F32_16X16X128_F8F6F4 = 45
V_MFMA_SCALE_F32_32X32X64_F8F6F4 = 46
class VOP3SDOp(Enum):
class VOP3SDOp(ReprEnum):
V_ADD_CO_U32 = 281
V_SUB_CO_U32 = 282
V_SUBREV_CO_U32 = 283
@@ -2004,7 +2003,7 @@ class VOP3SDOp(Enum):
V_MAD_U64_U32 = 488
V_MAD_I64_I32 = 489
class VOPCOp(Enum):
class VOPCOp(ReprEnum):
V_CMP_CLASS_F32_E32 = 16
V_CMPX_CLASS_F32_E32 = 17
V_CMP_CLASS_F64_E32 = 18
+71 -69
View File
@@ -38,7 +38,7 @@ class GLOBAL(Inst):
addr = VGPRField(39, 32)
data = VGPRField(47, 40)
saddr = SGPRField(54, 48, default=NULL)
offset = BitField(11, 0)
offset = BitField(12, 0)
seg = FixedBitField(15, 14, 2)
acc = BitField(55, 55)
sve = BitField(13, 13)
@@ -53,7 +53,7 @@ class SCRATCH(Inst):
addr = VGPRField(39, 32)
data = VGPRField(47, 40)
saddr = SGPRField(54, 48, default=NULL)
offset = BitField(11, 0)
offset = BitField(12, 0)
seg = FixedBitField(15, 14, 1)
acc = BitField(55, 55)
sve = BitField(13, 13)
@@ -164,7 +164,7 @@ class VOP3(Inst):
class VOP3P(Inst):
encoding = FixedBitField(31, 23, 0b110100111)
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE})
vdst = VGPRField(7, 0)
src0 = SrcField(40, 32)
src1 = SrcField(49, 41)
@@ -174,6 +174,7 @@ class VOP3P(Inst):
clmp = BitField(15, 15)
opsel = BitField(13, 11)
opsel_hi = BitField(60, 59, default=3)
opsel_hi2 = BitField(14, 14, default=1)
class VOP3PX2(Inst):
encoding = FixedBitField(95, 87, 0b110100111)
@@ -309,6 +310,7 @@ class VOP2_SDWA_SDST(VOP2):
s1 = BitField(63, 63)
class VOP3P_MFMA(VOP3P):
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
cbsz = BitField(10, 8)
abid = BitField(14, 11)
acc_cd = BitField(15, 15)
@@ -1647,80 +1649,80 @@ v_dot4_u32_u8 = functools.partial(VOP3P, VOP3POp.V_DOT4_U32_U8)
v_dot8_i32_i4 = functools.partial(VOP3P, VOP3POp.V_DOT8_I32_I4)
v_dot8_u32_u4 = functools.partial(VOP3P, VOP3POp.V_DOT8_U32_U4)
v_mfma_ld_scale_b32 = functools.partial(VOP3P, VOP3POp.V_MFMA_LD_SCALE_B32)
v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
v_pk_fma_f32 = functools.partial(VOP3P, VOP3POp.V_PK_FMA_F32)
v_pk_mul_f32 = functools.partial(VOP3P, VOP3POp.V_PK_MUL_F32)
v_pk_add_f32 = functools.partial(VOP3P, VOP3POp.V_PK_ADD_F32)
v_pk_mov_b32 = functools.partial(VOP3P, VOP3POp.V_PK_MOV_B32)
v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF16)
v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X64_I8)
v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF16)
v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X32_I8)
v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF16)
v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X64_I8)
v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF16)
v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X32_I8)
v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
v_mfma_f32_16x16x8_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X8_XF32)
v_mfma_f32_32x32x4_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_XF32)
v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X2_F32)
v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_F32)
v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_F16)
v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_F16)
v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_F16)
v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_F16)
v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X16_I8)
v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X32_I8)
v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X2_F32)
v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_F32)
v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_F16)
v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_F16)
v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_F16)
v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_F16)
v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X16_I8)
v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X32_I8)
v_accvgpr_read = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_READ)
v_accvgpr_write = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_WRITE)
v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_BF16)
v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_BF16)
v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_16X16X4_F64)
v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_BF16)
v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_BF16)
v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_16X16X4_F64)
v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
v_mfma_scale_f32_16x16x128_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_16X16X128_F8F6F4)
v_mfma_scale_f32_32x32x64_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_32X32X64_F8F6F4)
v_add_co_u32 = functools.partial(VOP3SD, VOP3SDOp.V_ADD_CO_U32)
+4
View File
@@ -1,6 +1,10 @@
# autogenerated from AMD ISA XML - do not edit
from enum import Enum, auto
class ReprEnum(Enum):
"""Enum with clean repr that roundtrips with eval()."""
def __repr__(self): return f"{type(self).__name__}.{self.name}"
class Fmt(Enum):
FMT_ANY = auto()
FMT_BUF = auto()
+26 -27
View File
@@ -1,8 +1,7 @@
# autogenerated from AMD ISA XML - do not edit
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(Enum):
class DSOp(ReprEnum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -130,10 +129,10 @@ class DSOp(Enum):
DS_LOAD_B96 = 254
DS_LOAD_B128 = 255
class EXPOp(Enum):
class EXPOp(ReprEnum):
EXP = 0
class FLATOp(Enum):
class FLATOp(ReprEnum):
FLAT_LOAD_U8 = 16
FLAT_LOAD_I8 = 17
FLAT_LOAD_U16 = 18
@@ -191,7 +190,7 @@ class FLATOp(Enum):
FLAT_ATOMIC_MAX_F32 = 82
FLAT_ATOMIC_ADD_F32 = 86
class GLOBALOp(Enum):
class GLOBALOp(ReprEnum):
GLOBAL_LOAD_U8 = 16
GLOBAL_LOAD_I8 = 17
GLOBAL_LOAD_U16 = 18
@@ -254,7 +253,7 @@ class GLOBALOp(Enum):
GLOBAL_ATOMIC_MAX_F32 = 82
GLOBAL_ATOMIC_ADD_F32 = 86
class HWREG(Enum):
class HWREG(ReprEnum):
HW_REG_MODE = 1
HW_REG_STATUS = 2
HW_REG_TRAPSTS = 3
@@ -280,11 +279,11 @@ class HWREG(Enum):
HW_REG_IB_STS2 = 28
HW_REG_SHADER_CYCLES = 29
class LDSDIROp(Enum):
class LDSDIROp(ReprEnum):
LDS_PARAM_LOAD = 0
LDS_DIRECT_LOAD = 1
class MIMGOp(Enum):
class MIMGOp(ReprEnum):
IMAGE_LOAD = 0
IMAGE_LOAD_MIP = 1
IMAGE_LOAD_PCK = 2
@@ -370,7 +369,7 @@ class MIMGOp(Enum):
IMAGE_GATHER4_C_B_CL = 101
IMAGE_GATHER4H = 144
class MSG(Enum):
class MSG(ReprEnum):
MSG_RTN_GET_DOORBELL = 128
MSG_RTN_GET_DDID = 129
MSG_RTN_GET_TMA = 130
@@ -380,7 +379,7 @@ class MSG(Enum):
MSG_RTN_GET_TBA_TO_PC = 134
MSG_RTN_ILLEGAL_MSG = 255
class MTBUFOp(Enum):
class MTBUFOp(ReprEnum):
TBUFFER_LOAD_FORMAT_X = 0
TBUFFER_LOAD_FORMAT_XY = 1
TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -398,7 +397,7 @@ class MTBUFOp(Enum):
TBUFFER_STORE_D16_FORMAT_XYZ = 14
TBUFFER_STORE_D16_FORMAT_XYZW = 15
class MUBUFOp(Enum):
class MUBUFOp(ReprEnum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -479,7 +478,7 @@ class MUBUFOp(Enum):
BUFFER_ATOMIC_MAX_F32 = 82
BUFFER_ATOMIC_ADD_F32 = 86
class SCRATCHOp(Enum):
class SCRATCHOp(ReprEnum):
SCRATCH_LOAD_U8 = 16
SCRATCH_LOAD_I8 = 17
SCRATCH_LOAD_U16 = 18
@@ -508,7 +507,7 @@ class SCRATCHOp(Enum):
SCRATCH_LOAD_LDS_I16 = 48
SCRATCH_LOAD_LDS_B32 = 49
class SMEMOp(Enum):
class SMEMOp(ReprEnum):
S_LOAD_B32 = 0
S_LOAD_B64 = 1
S_LOAD_B128 = 2
@@ -524,7 +523,7 @@ class SMEMOp(Enum):
S_ATC_PROBE = 34
S_ATC_PROBE_BUFFER = 35
class SOP1Op(Enum):
class SOP1Op(ReprEnum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -606,7 +605,7 @@ class SOP1Op(Enum):
S_TRUNC_F16 = 109
S_RNDNE_F16 = 110
class SOP2Op(Enum):
class SOP2Op(ReprEnum):
S_ADD_U32 = 0
S_SUB_U32 = 1
S_ADD_I32 = 2
@@ -675,7 +674,7 @@ class SOP2Op(Enum):
S_MUL_F16 = 77
S_FMAC_F16 = 78
class SOPCOp(Enum):
class SOPCOp(ReprEnum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -723,7 +722,7 @@ class SOPCOp(Enum):
S_CMP_NEQ_F16 = 93
S_CMP_NLT_F16 = 94
class SOPKOp(Enum):
class SOPKOp(ReprEnum):
S_MOVK_I32 = 0
S_VERSION = 1
S_CMOVK_I32 = 2
@@ -752,7 +751,7 @@ class SOPKOp(Enum):
S_WAITCNT_EXPCNT = 26
S_WAITCNT_LGKMCNT = 27
class SOPPOp(Enum):
class SOPPOp(ReprEnum):
S_NOP = 0
S_SETKILL = 1
S_SETHALT = 2
@@ -793,7 +792,7 @@ class SOPPOp(Enum):
S_ICACHE_INV = 60
S_BARRIER = 61
class VINTERPOp(Enum):
class VINTERPOp(ReprEnum):
V_INTERP_P10_F32 = 0
V_INTERP_P2_F32 = 1
V_INTERP_P10_F16_F32 = 2
@@ -801,7 +800,7 @@ class VINTERPOp(Enum):
V_INTERP_P10_RTZ_F16_F32 = 4
V_INTERP_P2_RTZ_F16_F32 = 5
class VOP1Op(Enum):
class VOP1Op(ReprEnum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -975,7 +974,7 @@ class VOP1Op(Enum):
V_CVT_I32_I16 = V_CVT_I32_I16_E32
V_CVT_U32_U16 = V_CVT_U32_U16_E32
class VOP2Op(Enum):
class VOP2Op(ReprEnum):
V_CNDMASK_B32_E32 = 1
V_DOT2ACC_F32_F16_E32 = 2
V_ADD_F32_E32 = 3
@@ -1069,7 +1068,7 @@ class VOP2Op(Enum):
V_LDEXP_F16 = V_LDEXP_F16_E32
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
class VOP3Op(Enum):
class VOP3Op(ReprEnum):
V_CMP_F_F16_E64 = 0
V_CMP_LT_F16_E64 = 1
V_CMP_EQ_F16_E64 = 2
@@ -1809,7 +1808,7 @@ class VOP3Op(Enum):
V_CVT_I32_I16 = V_CVT_I32_I16_E64
V_CVT_U32_U16 = V_CVT_U32_U16_E64
class VOP3POp(Enum):
class VOP3POp(ReprEnum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1845,7 +1844,7 @@ class VOP3POp(Enum):
V_WMMA_I32_16X16X16_IU8 = 68
V_WMMA_I32_16X16X16_IU4 = 69
class VOP3SDOp(Enum):
class VOP3SDOp(ReprEnum):
V_ADD_CO_CI_U32 = 288
V_SUB_CO_CI_U32 = 289
V_SUBREV_CO_CI_U32 = 290
@@ -1857,7 +1856,7 @@ class VOP3SDOp(Enum):
V_SUB_CO_U32 = 769
V_SUBREV_CO_U32 = 770
class VOPCOp(Enum):
class VOPCOp(ReprEnum):
V_CMP_F_F16_E32 = 0
V_CMP_LT_F16_E32 = 1
V_CMP_EQ_F16_E32 = 2
@@ -2239,7 +2238,7 @@ class VOPCOp(Enum):
V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
class VOPDOp(Enum):
class VOPDOp(ReprEnum):
V_DUAL_FMAC_F32 = 0
V_DUAL_FMAAK_F32 = 1
V_DUAL_FMAMK_F32 = 2
+26 -27
View File
@@ -1,8 +1,7 @@
# autogenerated from AMD ISA XML - do not edit
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(Enum):
class DSOp(ReprEnum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -127,7 +126,7 @@ class DSOp(Enum):
DS_LOAD_B96 = 254
DS_LOAD_B128 = 255
class HWREG(Enum):
class HWREG(ReprEnum):
HW_REG_WAVE_MODE = 1
HW_REG_WAVE_STATUS = 2
HW_REG_WAVE_STATE_PRIV = 4
@@ -148,7 +147,7 @@ class HWREG(Enum):
HW_REG_SHADER_CYCLES_LO = 29
HW_REG_SHADER_CYCLES_HI = 30
class MSG(Enum):
class MSG(ReprEnum):
MSG_RTN_GET_DOORBELL = 128
MSG_RTN_GET_DDID = 129
MSG_RTN_GET_TMA = 130
@@ -159,7 +158,7 @@ class MSG(Enum):
MSG_RTN_GET_SE_HW_ID = 135
MSG_RTN_ILLEGAL_MSG = 255
class SMEMOp(Enum):
class SMEMOp(ReprEnum):
S_LOAD_B32 = 0
S_LOAD_B64 = 1
S_LOAD_B128 = 2
@@ -189,7 +188,7 @@ class SMEMOp(Enum):
S_BUFFER_PREFETCH_DATA = 39
S_PREFETCH_DATA_PC_REL = 40
class SOP1Op(Enum):
class SOP1Op(ReprEnum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -278,7 +277,7 @@ class SOP1Op(Enum):
S_TRUNC_F16 = 109
S_RNDNE_F16 = 110
class SOP2Op(Enum):
class SOP2Op(ReprEnum):
S_ADD_CO_U32 = 0
S_SUB_CO_U32 = 1
S_ADD_CO_I32 = 2
@@ -354,7 +353,7 @@ class SOP2Op(Enum):
S_SUB_NC_U64 = 84
S_MUL_U64 = 85
class SOPCOp(Enum):
class SOPCOp(ReprEnum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -402,7 +401,7 @@ class SOPCOp(Enum):
S_CMP_NEQ_F16 = 93
S_CMP_NLT_F16 = 94
class SOPKOp(Enum):
class SOPKOp(ReprEnum):
S_MOVK_I32 = 0
S_VERSION = 1
S_CMOVK_I32 = 2
@@ -413,7 +412,7 @@ class SOPKOp(Enum):
S_SETREG_IMM32_B32 = 19
S_CALL_B64 = 20
class SOPPOp(Enum):
class SOPPOp(ReprEnum):
S_NOP = 0
S_SETKILL = 1
S_SETHALT = 2
@@ -458,7 +457,7 @@ class SOPPOp(Enum):
S_WAIT_LOADCNT_DSCNT = 72
S_WAIT_STORECNT_DSCNT = 73
class VBUFFEROp(Enum):
class VBUFFEROp(ReprEnum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -549,14 +548,14 @@ class VBUFFEROp(Enum):
TBUFFER_STORE_D16_FORMAT_XYZ = 142
TBUFFER_STORE_D16_FORMAT_XYZW = 143
class VDSDIROp(Enum):
class VDSDIROp(ReprEnum):
DS_PARAM_LOAD = 0
DS_DIRECT_LOAD = 1
class VEXPORTOp(Enum):
class VEXPORTOp(ReprEnum):
EXPORT = 0
class VFLATOp(Enum):
class VFLATOp(ReprEnum):
FLAT_LOAD_U8 = 16
FLAT_LOAD_I8 = 17
FLAT_LOAD_U16 = 18
@@ -615,7 +614,7 @@ class VFLATOp(Enum):
FLAT_ATOMIC_PK_ADD_F16 = 89
FLAT_ATOMIC_PK_ADD_BF16 = 90
class VGLOBALOp(Enum):
class VGLOBALOp(ReprEnum):
GLOBAL_LOAD_U8 = 16
GLOBAL_LOAD_I8 = 17
GLOBAL_LOAD_U16 = 18
@@ -682,7 +681,7 @@ class VGLOBALOp(Enum):
GLOBAL_ATOMIC_PK_ADD_BF16 = 90
GLOBAL_ATOMIC_ORDERED_ADD_B64 = 115
class VIMAGEOp(Enum):
class VIMAGEOp(ReprEnum):
IMAGE_LOAD = 0
IMAGE_LOAD_MIP = 1
IMAGE_LOAD_PCK = 2
@@ -717,7 +716,7 @@ class VIMAGEOp(Enum):
IMAGE_ATOMIC_PK_ADD_F16 = 134
IMAGE_ATOMIC_PK_ADD_BF16 = 135
class VINTERPOp(Enum):
class VINTERPOp(ReprEnum):
V_INTERP_P10_F32 = 0
V_INTERP_P2_F32 = 1
V_INTERP_P10_F16_F32 = 2
@@ -725,7 +724,7 @@ class VINTERPOp(Enum):
V_INTERP_P10_RTZ_F16_F32 = 4
V_INTERP_P2_RTZ_F16_F32 = 5
class VOP1Op(Enum):
class VOP1Op(ReprEnum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -907,7 +906,7 @@ class VOP1Op(Enum):
V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E32
V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E32
class VOP2Op(Enum):
class VOP2Op(ReprEnum):
V_CNDMASK_B32_E32 = 1
V_ADD_F64_E32 = 2
V_ADD_F32_E32 = 3
@@ -1007,7 +1006,7 @@ class VOP2Op(Enum):
V_LDEXP_F16 = V_LDEXP_F16_E32
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
class VOP3Op(Enum):
class VOP3Op(ReprEnum):
V_CMP_LT_F16_E64 = 1
V_CMP_EQ_F16_E64 = 2
V_CMP_LE_F16_E64 = 3
@@ -1732,7 +1731,7 @@ class VOP3Op(Enum):
V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E64
V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E64
class VOP3POp(Enum):
class VOP3POp(ReprEnum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1790,7 +1789,7 @@ class VOP3POp(Enum):
V_SWMMAC_F32_16X16X32_BF8_FP8 = 89
V_SWMMAC_F32_16X16X32_BF8_BF8 = 90
class VOP3SDOp(Enum):
class VOP3SDOp(ReprEnum):
V_ADD_CO_CI_U32 = 288
V_SUB_CO_CI_U32 = 289
V_SUBREV_CO_CI_U32 = 290
@@ -1802,7 +1801,7 @@ class VOP3SDOp(Enum):
V_SUB_CO_U32 = 769
V_SUBREV_CO_U32 = 770
class VOPCOp(Enum):
class VOPCOp(ReprEnum):
V_CMP_LT_F16_E32 = 1
V_CMP_EQ_F16_E32 = 2
V_CMP_LE_F16_E32 = 3
@@ -2128,7 +2127,7 @@ class VOPCOp(Enum):
V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
class VOPDOp(Enum):
class VOPDOp(ReprEnum):
V_DUAL_FMAC_F32 = 0
V_DUAL_FMAAK_F32 = 1
V_DUAL_FMAMK_F32 = 2
@@ -2147,7 +2146,7 @@ class VOPDOp(Enum):
V_DUAL_LSHLREV_B32 = 17
V_DUAL_AND_B32 = 18
class VSAMPLEOp(Enum):
class VSAMPLEOp(ReprEnum):
IMAGE_MSAA_LOAD = 24
IMAGE_SAMPLE = 27
IMAGE_SAMPLE_D = 28
@@ -2207,7 +2206,7 @@ class VSAMPLEOp(Enum):
IMAGE_GATHER4_C_B_CL = 101
IMAGE_GATHER4H = 144
class VSCRATCHOp(Enum):
class VSCRATCHOp(ReprEnum):
SCRATCH_LOAD_U8 = 16
SCRATCH_LOAD_I8 = 17
SCRATCH_LOAD_U16 = 18
+16 -16
View File
@@ -101,11 +101,11 @@ class VFLAT(Inst):
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = BitField(62, 55)
vsrc = VGPRField(62, 55)
ioffset = BitField(95, 72)
class VGLOBAL(Inst):
encoding = FixedBitField(31, 24, 0b11101100)
encoding = FixedBitField(31, 24, 0b11101110)
op = EnumBitField(21, 14, VGLOBALOp, {VGLOBALOp.GLOBAL_LOAD_U8, VGLOBALOp.GLOBAL_LOAD_I8, VGLOBALOp.GLOBAL_LOAD_U16, VGLOBALOp.GLOBAL_LOAD_I16, VGLOBALOp.GLOBAL_LOAD_B32, VGLOBALOp.GLOBAL_LOAD_B64, VGLOBALOp.GLOBAL_LOAD_B96, VGLOBALOp.GLOBAL_LOAD_B128, VGLOBALOp.GLOBAL_STORE_B8, VGLOBALOp.GLOBAL_STORE_B16, VGLOBALOp.GLOBAL_STORE_B32, VGLOBALOp.GLOBAL_STORE_B64, VGLOBALOp.GLOBAL_STORE_B96, VGLOBALOp.GLOBAL_STORE_B128, VGLOBALOp.GLOBAL_LOAD_D16_U8, VGLOBALOp.GLOBAL_LOAD_D16_I8, VGLOBALOp.GLOBAL_LOAD_D16_B16, VGLOBALOp.GLOBAL_LOAD_D16_HI_U8, VGLOBALOp.GLOBAL_LOAD_D16_HI_I8, VGLOBALOp.GLOBAL_LOAD_D16_HI_B16, VGLOBALOp.GLOBAL_STORE_D16_HI_B8, VGLOBALOp.GLOBAL_STORE_D16_HI_B16, VGLOBALOp.GLOBAL_LOAD_ADDTID_B32, VGLOBALOp.GLOBAL_STORE_ADDTID_B32, VGLOBALOp.GLOBAL_INV, VGLOBALOp.GLOBAL_WB, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_ADD_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_CLAMP_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_I32, VGLOBALOp.GLOBAL_ATOMIC_MIN_U32, VGLOBALOp.GLOBAL_ATOMIC_MAX_I32, VGLOBALOp.GLOBAL_ATOMIC_MAX_U32, VGLOBALOp.GLOBAL_ATOMIC_AND_B32, VGLOBALOp.GLOBAL_ATOMIC_OR_B32, VGLOBALOp.GLOBAL_ATOMIC_XOR_B32, VGLOBALOp.GLOBAL_ATOMIC_INC_U32, VGLOBALOp.GLOBAL_ATOMIC_DEC_U32, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_ADD_U64, VGLOBALOp.GLOBAL_ATOMIC_SUB_U64, VGLOBALOp.GLOBAL_ATOMIC_MIN_I64, VGLOBALOp.GLOBAL_ATOMIC_MIN_U64, VGLOBALOp.GLOBAL_ATOMIC_MAX_I64, VGLOBALOp.GLOBAL_ATOMIC_MAX_U64, VGLOBALOp.GLOBAL_ATOMIC_AND_B64, VGLOBALOp.GLOBAL_ATOMIC_OR_B64, VGLOBALOp.GLOBAL_ATOMIC_XOR_B64, VGLOBALOp.GLOBAL_ATOMIC_INC_U64, VGLOBALOp.GLOBAL_ATOMIC_DEC_U64, VGLOBALOp.GLOBAL_WBINV, VGLOBALOp.GLOBAL_ATOMIC_COND_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_NUM_F32, VGLOBALOp.GLOBAL_ATOMIC_MAX_NUM_F32, VGLOBALOp.GLOBAL_LOAD_BLOCK, VGLOBALOp.GLOBAL_STORE_BLOCK, VGLOBALOp.GLOBAL_ATOMIC_ADD_F32, VGLOBALOp.GLOBAL_LOAD_TR_B128, VGLOBALOp.GLOBAL_LOAD_TR_B64, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_F16, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_BF16, VGLOBALOp.GLOBAL_ATOMIC_ORDERED_ADD_B64})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
@@ -114,20 +114,7 @@ class VGLOBAL(Inst):
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VSCRATCH(Inst):
encoding = FixedBitField(31, 24, 0b11101100)
op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
saddr = SGPRField(6, 0, default=NULL)
nv = BitField(7, 7)
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = BitField(62, 55)
vsrc = VGPRField(62, 55)
ioffset = BitField(95, 72)
class VIMAGE(Inst):
@@ -253,6 +240,19 @@ class VSAMPLE(Inst):
vaddr2 = BitField(87, 80)
vaddr3 = BitField(95, 88)
class VSCRATCH(Inst):
encoding = FixedBitField(31, 24, 0b11101101)
op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
saddr = SGPRField(6, 0, default=NULL)
nv = BitField(7, 7)
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
ioffset = BitField(95, 72)
class SOP1_LIT(SOP1):
op = EnumBitField(15, 8, SOP1Op, {SOP1Op.S_MOV_B32, SOP1Op.S_MOV_B64, SOP1Op.S_CMOV_B32, SOP1Op.S_CMOV_B64, SOP1Op.S_BREV_B32, SOP1Op.S_BREV_B64, SOP1Op.S_CTZ_I32_B32, SOP1Op.S_CTZ_I32_B64, SOP1Op.S_CLZ_I32_U32, SOP1Op.S_CLZ_I32_U64, SOP1Op.S_CLS_I32, SOP1Op.S_CLS_I32_I64, SOP1Op.S_SEXT_I32_I8, SOP1Op.S_SEXT_I32_I16, SOP1Op.S_BITSET0_B32, SOP1Op.S_BITSET0_B64, SOP1Op.S_BITSET1_B32, SOP1Op.S_BITSET1_B64, SOP1Op.S_BITREPLICATE_B64_B32, SOP1Op.S_ABS_I32, SOP1Op.S_BCNT0_I32_B32, SOP1Op.S_BCNT0_I32_B64, SOP1Op.S_BCNT1_I32_B32, SOP1Op.S_BCNT1_I32_B64, SOP1Op.S_QUADMASK_B32, SOP1Op.S_QUADMASK_B64, SOP1Op.S_WQM_B32, SOP1Op.S_WQM_B64, SOP1Op.S_NOT_B32, SOP1Op.S_NOT_B64, SOP1Op.S_AND_SAVEEXEC_B32, SOP1Op.S_AND_SAVEEXEC_B64, SOP1Op.S_OR_SAVEEXEC_B32, SOP1Op.S_OR_SAVEEXEC_B64, SOP1Op.S_XOR_SAVEEXEC_B32, SOP1Op.S_XOR_SAVEEXEC_B64, SOP1Op.S_NAND_SAVEEXEC_B32, SOP1Op.S_NAND_SAVEEXEC_B64, SOP1Op.S_NOR_SAVEEXEC_B32, SOP1Op.S_NOR_SAVEEXEC_B64, SOP1Op.S_XNOR_SAVEEXEC_B32, SOP1Op.S_XNOR_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_SAVEEXEC_B32, SOP1Op.S_AND_NOT0_SAVEEXEC_B64, SOP1Op.S_OR_NOT0_SAVEEXEC_B32, SOP1Op.S_OR_NOT0_SAVEEXEC_B64, SOP1Op.S_AND_NOT1_SAVEEXEC_B32, SOP1Op.S_AND_NOT1_SAVEEXEC_B64, SOP1Op.S_OR_NOT1_SAVEEXEC_B32, SOP1Op.S_OR_NOT1_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_WREXEC_B32, SOP1Op.S_AND_NOT0_WREXEC_B64, SOP1Op.S_AND_NOT1_WREXEC_B32, SOP1Op.S_AND_NOT1_WREXEC_B64, SOP1Op.S_MOVRELS_B32, SOP1Op.S_MOVRELS_B64, SOP1Op.S_MOVRELD_B32, SOP1Op.S_MOVRELD_B64, SOP1Op.S_MOVRELSD_2_B32, SOP1Op.S_GETPC_B64, SOP1Op.S_SETPC_B64, SOP1Op.S_SWAPPC_B64, SOP1Op.S_RFE_B64, SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64, SOP1Op.S_BARRIER_SIGNAL, SOP1Op.S_BARRIER_SIGNAL_ISFIRST, SOP1Op.S_GET_BARRIER_STATE, SOP1Op.S_BARRIER_INIT, SOP1Op.S_BARRIER_JOIN, SOP1Op.S_ALLOC_VGPR, SOP1Op.S_SLEEP_VAR, SOP1Op.S_CEIL_F32, SOP1Op.S_FLOOR_F32, SOP1Op.S_TRUNC_F32, SOP1Op.S_RNDNE_F32, SOP1Op.S_CVT_F32_I32, SOP1Op.S_CVT_F32_U32, SOP1Op.S_CVT_I32_F32, SOP1Op.S_CVT_U32_F32, SOP1Op.S_CVT_F16_F32, SOP1Op.S_CVT_F32_F16, SOP1Op.S_CVT_HI_F32_F16, SOP1Op.S_CEIL_F16, SOP1Op.S_FLOOR_F16, SOP1Op.S_TRUNC_F16, SOP1Op.S_RNDNE_F16})
literal = BitField(63, 32)
+30 -22
View File
@@ -1,6 +1,7 @@
# RDNA3/RDNA4/CDNA disassembler
from __future__ import annotations
import re, struct
from typing import Callable
from extra.assembly.amd.dsl import Inst, Reg
# Special register mappings for disassembly
@@ -84,10 +85,11 @@ from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as
VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4)
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, HWREG as HWREG_CDNA
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
def _is_r4(inst: Inst) -> bool: return 'rdna4' in inst.__class__.__module__
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
@@ -100,7 +102,7 @@ def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1
def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n) # accumulator registers for GFX90a
def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _ttmp(b, n: int = 1) -> str | None: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
v = _unwrap(v)
@@ -225,7 +227,7 @@ NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
def _disasm_sopp(inst: SOPP) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
is_rdna4 = _is_r4(inst)
# Ops that have no argument when simm16 == 0
no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
's_endpgm_ordered_ps_done', 's_code_end'}
@@ -257,14 +259,14 @@ def _disasm_sopp(inst: SOPP) -> str:
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
return f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if name in ('s_gl1_inv', 's_dcache_inv', 's_dcache_inv_vol', 's_dcache_wb', 's_dcache_wb_vol', 's_icache_inv'): return name
soe, imm = getattr(inst, 'soe', 0) or getattr(inst, 'soffset_en', 0), getattr(inst, 'imm', 1)
is_rdna4 = 'rdna4' in inst.__class__.__module__
is_rdna4 = _is_r4(inst)
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}"
@@ -303,27 +305,30 @@ def _disasm_smem(inst: SMEM) -> str:
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
def _disasm_flat(inst: FLAT) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(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'
if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
# Global/scratch uses 13-bit signed offset
offset = inst.ioffset if r4 else inst.offset
if seg != 'flat':
if cdna:
# CDNA: bit 12 is sign bit but not in offset field
raw = int.from_bytes(inst.to_bytes(), 'little')
off_val = inst.offset | ((raw >> 12) & 1) << 12 # get bit 12
off_val = offset | ((raw >> 12) & 1) << 12 # get bit 12
else:
off_val = inst.offset
off_val = offset
off_val = off_val if off_val < 4096 else off_val - 8192 # sign extend 13-bit
else:
off_val = inst.offset
off_val = offset
# Use get_field_bits: data for stores/atomics, d for loads
regs = inst.canonical_op_regs
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
off_s = f" offset:{off_val}" if off_val else ""
if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
if seg == 'flat': saddr_s = ""
elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
@@ -332,18 +337,21 @@ def _disasm_flat(inst: FLAT) -> str:
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
# RDNA4: vaddr instead of addr, vsrc instead of data
addr = inst.vaddr if r4 else inst.addr
data = inst.vsrc if r4 else inst.data
# load_lds_* instructions: vaddr, saddr (no vdst, data goes to LDS)
if 'load_lds' in name:
addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
return f"{instr} {addr_s}{saddr_s}{mods}"
if seg == 'flat': addr_w = 2 # flat always uses 64-bit vaddr
elif cdna: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(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
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
if 'atomic' in name:
glc_or_sc0 = inst.sc0 if cdna else inst.glc
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}"
@@ -458,16 +466,15 @@ def _disasm_vop3sd(inst: VOP3SD) -> str:
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal
is_rdna4 = 'rdna4' in inst.__class__.__module__
op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
vdst_y, nx, ny = (_unwrap(inst.vdsty) << 1) | ((_unwrap(inst.vdstx) & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
op_enum = R4_VOPDOp if _is_r4(inst) else VOPDOp
nx, ny = op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
def half(n, vd, s0, vs1):
vd, vs1 = _vi(vd), _vi(vs1)
if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, inst.vdsty, inst.srcy0, inst.vsrcy1)}"
def _disasm_vop3p(inst: VOP3P) -> str:
name = inst.op_name.lower()
@@ -545,7 +552,7 @@ _HWREG_BLACKLIST_CDNA = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_R
'HW_REG_SQ_PERF_SNAPSHOT_PC_LO', 'HW_REG_SQ_PERF_SNAPSHOT_PC_HI', 'HW_REG_XCC_ID'}
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__
is_rdna4 = _is_r4(inst)
hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
def fmt_hwreg(hid, hoff, hsz):
@@ -569,7 +576,7 @@ def _disasm_vinterp(inst: VINTERP) -> str:
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
DISASM_HANDLERS: dict[type, callable] = {
DISASM_HANDLERS: dict[type, Callable[..., str]] = {
VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
@@ -582,6 +589,7 @@ DISASM_HANDLERS: dict[type, callable] = {
R4_VOP2: _disasm_vop2, R4_VOP2_LIT: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOPC_LIT: _disasm_vopc,
R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3_SDST_LIT: _disasm_vop3, R4_VOP3_LIT: _disasm_vop3,
R4_VOP3SD: _disasm_vop3sd, R4_VOP3SD_LIT: _disasm_vop3sd, R4_VOP3P: _disasm_vop3p, R4_VOP3P_LIT: _disasm_vop3p,
R4_FLAT: _disasm_flat, R4_GLOBAL: _disasm_flat, R4_SCRATCH: _disasm_flat,
R4_VOPD: _disasm_vopd, R4_VOPD_LIT: _disasm_vopd, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem, R4_DS: _disasm_ds,
R4_SOP1: _disasm_sop1, R4_SOP1_LIT: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOP2_LIT: _disasm_sop2,
R4_SOPC: _disasm_sopc, R4_SOPC_LIT: _disasm_sopc, R4_SOPK: _disasm_sopk, R4_SOPK_LIT: _disasm_sopk}
@@ -596,7 +604,7 @@ from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as
VOP1_SDWA as CDNA_VOP1_SDWA, VOP1_DPP16 as CDNA_VOP1_DPP16,
VOP2 as CDNA_VOP2, VOP2_LIT as CDNA_VOP2_LIT, VOP2_SDWA as CDNA_VOP2_SDWA, VOP2_DPP16 as CDNA_VOP2_DPP16,
VOPC as CDNA_VOPC, VOPC_LIT as CDNA_VOPC_LIT, VOPC_SDWA_SDST as CDNA_VOPC_SDWA_SDST,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3PX2 as CDNA_VOP3PX2,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3P_MFMA as CDNA_VOP3P_MFMA, VOP3PX2 as CDNA_VOP3PX2,
SOP1 as CDNA_SOP1, SOP1_LIT as CDNA_SOP1_LIT, SOP2 as CDNA_SOP2, SOP2_LIT as CDNA_SOP2_LIT,
SOPC as CDNA_SOPC, SOPC_LIT as CDNA_SOPC_LIT, SOPK as CDNA_SOPK, SOPK_LIT as CDNA_SOPK_LIT,
SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
@@ -895,5 +903,5 @@ DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP1_LIT: _disasm_vop1,
CDNA_SOP1: _disasm_sop1, CDNA_SOP1_LIT: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOP2_LIT: _disasm_sop2,
CDNA_SOPC: _disasm_sopc, CDNA_SOPC_LIT: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPK_LIT: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_GLOBAL: _disasm_flat, CDNA_SCRATCH: _disasm_flat,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})
+71 -45
View File
@@ -1,27 +1,24 @@
# dsl.py - clean DSL for AMD assembly
from typing import Any
# ══════════════════════════════════════════════════════════════
# Registers - unified src encoding space (0-511)
# ══════════════════════════════════════════════════════════════
def _reg_size(t: str | None) -> int: return {'b64': 2, 'f64': 2, 'u64': 2, 'i64': 2, 'b128': 4}.get(t, 1)
class Reg:
_NAMES = {106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
# Register names vary by arch: RDNA has NULL@124/M0@125, CDNA has M0@124/reserved@125
# RDNA4 has DPP8@233, CDNA has SDWA@249/DPP@250/VCCZ@251/EXECZ@252
_NAMES = {102: "FLAT_SCRATCH_LO", 103: "FLAT_SCRATCH_HI", 104: "XNACK_MASK_LO", 105: "XNACK_MASK_HI",
106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
233: "DPP8", 234: "DPP8FI", 235: "SHARED_BASE", 236: "SHARED_LIMIT", 237: "PRIVATE_BASE", 238: "PRIVATE_LIMIT",
240: "0.5", 241: "-0.5", 242: "1.0", 243: "-1.0", 244: "2.0", 245: "-2.0", 246: "4.0", 247: "-4.0",
248: "INV_2PI", 250: "DPP16", 253: "SCC", 255: "LIT"}
248: "INV_2PI", 249: "SDWA", 250: "DPP", 251: "VCCZ", 252: "EXECZ", 253: "SCC", 254: "SRC_LDS_DIRECT", 255: "LIT"}
_PAIRS = {106: "VCC", 126: "EXEC"}
def __init__(self, offset: int = 0, sz: int = 512, *, neg: bool = False, abs_: bool = False, hi: bool = False):
self.offset, self.sz = offset, sz
self.neg, self.abs_, self.hi = neg, abs_, hi
# TODO: remove these legacy aliases
@property
def count(self): return self.sz
@property
def idx(self): return self.offset
def __hash__(self): return hash((self.offset, self.sz, self.neg, self.abs_, self.hi))
def __getitem__(self, key):
if isinstance(key, slice):
@@ -78,9 +75,13 @@ EXEC = src[126:127]
# 128: 0, 129-192: integers 1-64, 193-208: integers -1 to -16
# 240-248: float constants (0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0, -4.0, 1/(2*PI))
INV_2PI = src[248]
DPP16 = src[250]
SDWA = src[249]
DPP = DPP16 = src[250]
VCCZ = src[251]
EXECZ = src[252]
SCC = src[253]
# 255: literal constant
SRC_LDS_DIRECT = src[254]
LIT = src[255] # literal constant marker
v = src[256:511] # VGPR0-255
# ══════════════════════════════════════════════════════════════
@@ -93,12 +94,13 @@ class _Bits:
bits = _Bits()
class BitField:
name: str | None
def __init__(self, hi: int, lo: int, default: int = 0):
self.hi, self.lo, self.default, self.name, self.mask = hi, lo, default, None, (1 << (hi - lo + 1)) - 1
def __set_name__(self, owner, name): self.name = name
def __eq__(self, other) -> 'FixedBitField':
def __set_name__(self, owner, name: str): self.name = name
def __eq__(self, other) -> 'FixedBitField': # type: ignore[override]
if isinstance(other, int): return FixedBitField(self.hi, self.lo, other)
return NotImplemented
raise TypeError(f"BitField.__eq__ expects int, got {type(other).__name__}")
def enum(self, enum_cls) -> 'EnumBitField': return EnumBitField(self.hi, self.lo, enum_cls)
def encode(self, val) -> int:
assert isinstance(val, int), f"BitField.encode expects int, got {type(val).__name__}"
@@ -107,11 +109,14 @@ class BitField:
def set(self, raw: int, val) -> int:
if val is None: val = self.default
encoded = self.encode(val)
# Handle signed values: convert negative to 2's complement
if encoded < 0: encoded = encoded & self.mask
if encoded < 0 or encoded > self.mask: raise RuntimeError(f"field '{self.name}': value {encoded} doesn't fit in {self.hi - self.lo + 1} bits")
return (raw & ~(self.mask << self.lo)) | (encoded << self.lo)
def __get__(self, obj, objtype=None):
if obj is None: return self
return self.decode((obj._raw >> self.lo) & self.mask)
def __set__(self, obj, val): obj._raw = self.set(obj._raw, val)
class FixedBitField(BitField):
def set(self, raw: int, val=None) -> int:
@@ -168,7 +173,10 @@ class SrcField(BitField):
# Resize register based on operand info (skip non-resizable special registers)
# VCC/EXEC pairs (106, 126), NULL (124), M0 (125), float constants (240-255)
if reg.offset not in (124, 125) and not 240 <= reg.offset <= 255:
if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
# Map variant field names (vsrc0->src0, vsrc1->src1, etc.) for DPP/SDWA classes
assert self.name is not None
name = self.name[1:] if self.name.startswith('v') and self.name[1:] in obj.op_regs else self.name
if sz := obj.op_regs.get(name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
return reg
class VGPRField(SrcField):
@@ -211,7 +219,12 @@ class VDSTYField(BitField):
if not isinstance(val, Reg): raise TypeError(f"VDSTYField requires Reg, got {type(val).__name__}")
if not (256 <= val.offset < 512): raise ValueError(f"VDSTYField requires VGPR, got offset {val.offset}")
return (val.offset - 256) >> 1
def decode(self, raw): return raw # raw value, actual vdsty = (raw << 1) | ((vdstx & 1) ^ 1)
def __get__(self, obj, objtype=None):
if obj is None: return self
raw = (obj._raw >> self.lo) & self.mask
vdstx_bit0 = (obj.vdstx.offset - 256) & 1
vgpr_idx = (raw << 1) | (vdstx_bit0 ^ 1)
return Reg(256 + vgpr_idx, 1)
# ══════════════════════════════════════════════════════════════
# Operand info from XML
@@ -240,6 +253,15 @@ def _get_variant(cls, suffix: str):
module = sys.modules.get(cls.__module__)
return getattr(module, f"{cls.__name__}{suffix}", None) if module else None
def _canonical_name(name: str) -> str | None:
"""Map operand name to canonical name."""
if name in ('src0', 'vsrc0', 'ssrc0'): return 's0'
if name in ('src1', 'vsrc1', 'ssrc1'): return 's1'
if name == 'src2': return 's2'
if name in ('vdst', 'sdst', 'sdata'): return 'd'
if name in ('data', 'vdata', 'data0', 'vsrc'): return 'data'
return None
class Inst:
_fields: list[tuple[str, BitField]]
_base_size: int
@@ -255,31 +277,30 @@ class Inst:
cls._base_size = (max(f.hi for _, f in cls._fields) + 8) // 8
def __new__(cls, *args, **kwargs):
# Auto-upgrade to _LIT variant if needed (only for base classes, not variants)
# Auto-upgrade to variant if needed (only for base classes, not variants)
if not any(cls.__name__.endswith(sfx) for sfx in ('_LIT', '_DPP16', '_DPP8', '_SDWA', '_SDWA_SDST', '_MFMA')):
lit_cls = _get_variant(cls, '_LIT')
if lit_cls is not None:
# Check if any src field needs a literal
# Map positional args to field names to find src values
args_iter = iter(args)
for name, field in cls._fields:
if isinstance(field, FixedBitField): continue
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
if isinstance(field, SrcField) and _needs_literal(val):
return lit_cls(*args, **kwargs)
args_iter = iter(args)
for name, field in cls._fields:
if isinstance(field, FixedBitField): continue
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
if not isinstance(field, SrcField): continue
if isinstance(val, Reg) and val.offset == 255 and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
if isinstance(val, Reg) and val.offset == 249:
if (sdwa_cls := _get_variant(cls, '_SDWA') or _get_variant(cls, '_SDWA_SDST')): return sdwa_cls(*args, **kwargs)
if isinstance(val, Reg) and val.offset == 250 and (dpp_cls := _get_variant(cls, '_DPP16')): return dpp_cls(*args, **kwargs)
if _needs_literal(val) and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
return object.__new__(cls)
def __init__(self, *args, **kwargs):
self._raw = 0
# Map positional args to field names (skip FixedBitFields)
args_iter = iter(args)
vals = {}
vals: dict[str, Any] = {}
for name, field in self._fields:
if isinstance(field, FixedBitField): vals[name] = None
elif name in kwargs: vals[name] = kwargs[name]
else: vals[name] = next(args_iter, None)
remaining = list(args_iter)
assert not remaining, f"too many positional args: {remaining}"
assert not (remaining := list(args_iter)), f"too many positional args: {remaining}"
# Extract modifiers from Reg objects and merge into neg/abs/opsel
neg_bits, abs_bits, opsel_bits = 0, 0, 0
for name, bit in [('src0', 0), ('src1', 1), ('src2', 2)]:
@@ -304,16 +325,16 @@ class Inst:
# Set all field values
for name, field in self._fields:
self._raw = field.set(self._raw, vals[name])
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC)
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC, SDWA/DPP markers)
for name, expected in self.op_regs.items():
if (val := vals.get(name)) is None: continue
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or val.offset == 253):
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or 249 <= val.offset <= 255):
raise TypeError(f"{name} expects {expected} register(s), got {val.sz}")
@property
def op_name(self) -> str: return self.op.name
def op_name(self) -> str: return getattr(self, 'op').name
@property
def operands(self) -> dict: return OPERANDS.get(self.op, {}) if hasattr(self, 'op') else {}
def operands(self) -> dict: return OPERANDS.get(getattr(self, 'op'), {}) if hasattr(self, 'op') else {}
def _is_cdna(self) -> bool: return 'cdna' in type(self).__module__
@functools.cached_property
@@ -325,9 +346,9 @@ class Inst:
if not self._is_cdna():
name = self.op_name.lower()
if 'cndmask' in name and 'src2' in bits: bits['src2'] = 32
if '_co_ci_' in name:
if 'src2' in bits: bits['src2'] = 32
if 'sdst' in bits: bits['sdst'] = 32
if '_co_ci_' in name and 'src2' in bits: bits['src2'] = 32 # carry-in source
# VOP3SD: sdst is always wavefront-size dependent (carry-out or condition mask)
if 'VOP3SD' in type(self).__name__ and 'sdst' in bits: bits['sdst'] = 32
if 'cmp' in name and 'vdst' in bits: bits['vdst'] = 32
# GLOBAL/FLAT: addr is 32-bit if saddr is valid SGPR, 64-bit if saddr is NULL
# SCRATCH: addr is always 32-bit (offset from scratch base, not absolute address)
@@ -341,8 +362,8 @@ class Inst:
# VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
if 'f8f6f4' in getattr(self, 'op_name', '').lower():
# Use explicit fields if available (VOP3PX2), else extract from VOP3P-MAI bit positions
cbsz = self.cbsz if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
blgp = self.blgp if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
cbsz = getattr(self, 'cbsz') if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
blgp = getattr(self, 'blgp') if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
bits['src0'], bits['src1'] = vgprs.get(cbsz, 8) * 32, vgprs.get(blgp, 8) * 32
return bits
@@ -356,12 +377,17 @@ class Inst:
"""Get bit widths with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
bits = {'d': 32, 's0': 32, 's1': 32, 's2': 32, 'data': 32}
for name, val in self.op_bits.items():
if name in ('src0', 'vsrc0', 'ssrc0'): bits['s0'] = val
elif name in ('src1', 'vsrc1', 'ssrc1'): bits['s1'] = val
elif name == 'src2': bits['s2'] = val
elif name in ('vdst', 'sdst', 'sdata'): bits['d'] = val
elif name in ('data', 'vdata', 'data0'): bits['data'] = val
if (cn := _canonical_name(name)): bits[cn] = val
return bits
@functools.cached_property
def canonical_operands(self) -> dict:
"""Get operands with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
result = {}
for name, val in self.operands.items():
if (cn := _canonical_name(name)): result[cn] = val
return result
@property
def canonical_op_regs(self) -> dict[str, int]:
"""Get register counts with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
+1200 -449
View File
File diff suppressed because it is too large Load Diff
@@ -22,6 +22,8 @@ FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_
"rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}},
"cdna": {"DS": {152: "DS_GWS_SEMA_RELEASE_ALL", 154: "DS_GWS_SEMA_V", 156: "DS_GWS_SEMA_P"},
"VOP3P": {44: "V_MFMA_LD_SCALE_B32", 62: "V_MFMA_F32_16X16X8_XF32", 63: "V_MFMA_F32_32X32X4_XF32"}}}
# Fields missing from XML but present in hardware (format: {arch: {encoding: [(name, hi, lo), ...]}})
FIELD_FIXES = {"cdna": {"VOP3P": [("opsel_hi2", 14, 14)]}}
# Encoding suffixes to strip (variants we don't generate separate classes for)
_ENC_SUFFIXES = ("_NSA1",)
# Encoding suffix to class suffix mapping (for variants we DO generate)
@@ -124,13 +126,13 @@ def parse_xml(filename: str):
if fmt and fmt not in fmts: fmts[fmt] = 0
if otype: op_types_set.add(otype)
if op_info: types[(name, base_enum)] = op_info
# Find opcodes that only exist in _LIT encoding (no base format version)
lit_only_ops: dict[str, set[int]] = {}
# Find opcodes that only exist in a specific variant encoding (no base format version)
suffix_only_ops: dict[str, dict[str, set[int]]] = {} # {suffix: {base_fmt: {opcodes}}}
for base_fmt, opcodes in opcode_encs.items():
for opcode, encs in opcodes.items():
if all("_LIT" in e for e in encs):
lit_only_ops.setdefault(base_fmt, set()).add(opcode)
return encodings, enums, types, fmts, op_types_set, lit_only_ops
suffix = next((s for s in _ENC_SUFFIX_MAP.values() if all(s in e for e in encs)), None)
if suffix is not None: suffix_only_ops.setdefault(suffix, {}).setdefault(base_fmt, set()).add(opcode)
return encodings, enums, types, fmts, op_types_set, suffix_only_ops
# ═══════════════════════════════════════════════════════════════════════════════
# PDF parsing
@@ -220,6 +222,10 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
def write_common(all_fmts, all_op_types, path):
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
lines.append("class ReprEnum(Enum):")
lines.append(' """Enum with clean repr that roundtrips with eval()."""')
lines.append(' def __repr__(self): return f"{type(self).__name__}.{self.name}"')
lines.append("")
lines.append("class Fmt(Enum):")
for fmt in sorted(all_fmts.keys()): lines.append(f" {fmt} = auto()")
lines.append("")
@@ -232,11 +238,11 @@ def write_common(all_fmts, all_op_types, path):
with open(path, "w") as f: f.write("\n".join(lines))
def write_enum(enums, path):
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum", "from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401", ""]
lines = ["# autogenerated from AMD ISA XML - do not edit", "from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401", ""]
for name, ops in sorted(enums.items()):
if not ops: continue
suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
lines.append(f"class {name}(Enum):" if name in ("HWREG", "MSG") else f"class {name}Op(Enum):")
lines.append(f"class {name}(ReprEnum):" if name in ("HWREG", "MSG") else f"class {name}Op(ReprEnum):")
aliases = []
for op, mem in sorted(ops.items()):
msuf = suffix if name != "VOP3" or op < 512 else ""
@@ -246,8 +252,8 @@ def write_enum(enums, path):
lines.append("")
with open(path, "w") as f: f.write("\n".join(lines))
def write_ins(encodings, enums, lit_only_ops, types, arch, path):
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr"}
def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr", "vsrc"}
_VARIANT_SUFFIXES = ("_LIT", "_DPP16", "_DPP8", "_SDWA_SDST", "_SDWA", "_MFMA")
def get_base_fmt(fmt):
for sfx in _VARIANT_SUFFIXES: fmt = fmt.replace(sfx, "")
@@ -267,6 +273,8 @@ def write_ins(encodings, enums, lit_only_ops, types, arch, path):
if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
# GLOBAL/SCRATCH: offset is 13-bit signed [12:0], FLAT: 12-bit unsigned (XML has 12-bit for all)
if name == "offset" and base_fmt in ("GLOBAL", "SCRATCH"): return f"BitField(12, {lo})"
if base_fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
if base_fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
return f"BitField({hi}, {lo})"
@@ -305,12 +313,15 @@ def write_ins(encodings, enums, lit_only_ops, types, arch, path):
# Generate base classes first
for enc_name, (fields, enc_bits) in sorted(base_encodings.items()):
# Get lit-only ops for this format (these can't be used in base class)
base_lit_ops = lit_only_ops.get(enc_name, set())
all_ops = set(enums.get(enc_name, {}).keys())
# Get suffix-only ops for this format (these can't be used in base class)
base_suffix_ops = set().union(*(d.get(enc_name, set()) for d in suffix_only_ops.values()))
# Exclude SDST ops from base class (they need VOP1_SDST/VOP3_SDST/VOP3B)
base_allowed = all_ops - base_lit_ops - sdst_opcodes.get(enc_name, set())
if enc_name in ("FLAT", "VFLAT"):
base_allowed = all_ops - base_suffix_ops - sdst_opcodes.get(enc_name, set())
# RDNA3 FLAT/GLOBAL/SCRATCH share encoding bits, differentiated by seg field
# RDNA4 VFLAT/VGLOBAL/VSCRATCH have distinct encoding bits, no seg field needed
has_seg_field = any(fn == "seg" for fn, _, _ in fields)
if enc_name in ("FLAT", "VFLAT") and has_seg_field:
prefix = "V" if enc_name == "VFLAT" else ""
for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
cls_ops = set(enums.get(cls, {}).keys())
@@ -320,7 +331,7 @@ def write_ins(encodings, enums, lit_only_ops, types, arch, path):
elif fn == "op": lines.append(f" op = EnumBitField({hi}, {lo}, {op_enum}, {fmt_allowed(op_enum, cls_ops)})")
else: lines.append(f" {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
lines.append("")
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "VGLOBAL", "VSCRATCH", "DPP", "SDWA"):
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "DPP", "SDWA"):
lines.append(f"class {enc_name}(Inst):")
for fn, hi, lo in sort_fields(fields):
if fn == "op":
@@ -336,15 +347,19 @@ def write_ins(encodings, enums, lit_only_ops, types, arch, path):
if base not in base_encodings: continue # skip if no base class
base_fields = {f[0] for f in base_encodings[base][0]}
extra_fields = [(fn, hi, lo) for fn, hi, lo in fields if fn not in base_fields]
is_lit = enc_name.endswith("_LIT")
# Check if this is a suffix-only variant
variant_suffix = next((sfx for sfx in _VARIANT_SUFFIXES if enc_name.endswith(sfx)), None)
is_suffix_variant = variant_suffix in suffix_only_ops
all_ops = set(enums.get(base, {}).keys())
if extra_fields or is_lit:
if extra_fields or is_suffix_variant:
lines.append(f"class {enc_name}({base}):")
op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
# _LIT classes: override op to allow all opcodes (base excludes lit-only ops)
if op_field and is_lit:
# other classes override op to only suffix-only opcodes
if op_field and is_suffix_variant:
_, hi, lo = op_field
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', all_ops)})")
allowed_ops = all_ops if variant_suffix == "_LIT" else suffix_only_ops[variant_suffix][base]
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', allowed_ops)})")
for fn, hi, lo in sort_fields(extra_fields):
lines.append(f" {fn} = {field_def(fn, hi, lo, enc_name)}")
lines.append("")
@@ -378,14 +393,14 @@ def write_ins(encodings, enums, lit_only_ops, types, arch, path):
for fmt, ops in sorted(enums.items()):
if fmt not in base_encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
lit_ops = lit_only_ops.get(fmt, set())
op_to_suffix = {op:suffix for suffix,ops in suffix_only_ops.items() for op in ops.get(fmt, set())}
fmt_sdst_ops = sdst_opcodes.get(fmt, set())
for op, name in sorted(ops.items()):
msuf = suffix if fmt != "VOP3" or op < 512 else ""
# Determine class: SDST variants, LIT-only instructions, or base
# Determine class: SDST variants, suffix-specific variants (e.g., _MFMA, _LIT), or base
if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
elif fmt == "VOP3" and (op in fmt_sdst_ops or op < 256): cls = "VOP3_SDST"
elif op in lit_ops: cls = f"{fmt}_LIT"
elif op_to_suffix.get(op): cls = f"{fmt}{op_to_suffix[op]}"
else: cls = fmt
lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
with open(path, "w") as f: f.write("\n".join(lines))
@@ -434,9 +449,11 @@ if __name__ == "__main__":
# First pass: parse XML for all architectures
for arch, cfg in ARCHS.items():
print(f"Parsing XML: {cfg['xml']} -> {arch}")
encodings, enums, types, fmts, op_types_set, lit_only_ops = parse_xml(cfg["xml"])
encodings, enums, types, fmts, op_types_set, suffix_only_ops = parse_xml(cfg["xml"])
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "lit_only_ops": lit_only_ops}
for fmt, fields in FIELD_FIXES.get(arch, {}).items():
if fmt in encodings: encodings[fmt] = (encodings[fmt][0] + fields, encodings[fmt][1])
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "suffix_only_ops": suffix_only_ops}
for fmt, bits in fmts.items():
assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
all_fmts[fmt] = bits
@@ -449,7 +466,7 @@ if __name__ == "__main__":
for arch, data in arch_data.items():
base = pathlib.Path(__file__).parent / "autogen" / arch
write_enum(data["enums"], base / "enum.py")
write_ins(data["encodings"], data["enums"], data["lit_only_ops"], data["types"], arch, base / "ins.py")
write_ins(data["encodings"], data["enums"], data["suffix_only_ops"], data["types"], arch, base / "ins.py")
write_operands(data["types"], data["enums"], arch, base / "operands.py")
print(f" {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
# Second pass: parse PDFs and write pcode
+1097 -793
View File
File diff suppressed because it is too large Load Diff
+144 -72
View File
@@ -26,7 +26,7 @@ class AluSrc(Enum):
VALU_SALU = 3
class InstOp(Enum):
"""SQTT instruction operation types.
"""SQTT instruction operation types for RDNA3 (gfx1100).
Memory ops appear in two ranges depending on which SIMD executes them:
- 0x1x-0x2x range: ops on traced SIMD
@@ -95,6 +95,21 @@ class InstOp(Enum):
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
VALU_CMPX = 0x73 # v_cmpx_*
class InstOpL4(Enum):
"""SQTT instruction operation types for RDNA4 (gfx1200). Different encoding from RDNA3."""
# TODO: we need to do discovery of all of these from instructions
SALU = 0x0
SMEM = 0x1
UNK_02 = 0x2
JUMP_NO = 0x4
UNK_06 = 0x6
VMEM = 0x10
UNK_11 = 0x11
VINTERP = 0x12
UNK_14 = 0x14
OTHER_VMEM = 0x5e
UNK_60 = 0x60
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE BASE CLASS
# ═══════════════════════════════════════════════════════════════════════════════
@@ -117,9 +132,64 @@ class PacketType:
return inst
def __repr__(self) -> str:
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_'))
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_') and k != 'encoding')
return f"{self.__class__.__name__}({fields_str})"
# ═══════════════════════════════════════════════════════════════════════════════
# TS PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
class TS_DELTA_S8_W3(PacketType):
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[63:11]
class TS_DELTA_S8_W3_L4(PacketType): # Layout 4: 64->72 bits
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[71:11]
class TS_DELTA_S5_W3(PacketType):
encoding = bits[4:0] == 0b00110
delta = bits[7:5]
_padding = bits[51:8]
class TS_DELTA_S5_W3_L4(PacketType): # Layout 4: 52->56 bits
encoding = bits[4:0] == 0b00110
delta = bits[9:7]
_padding = bits[55:10]
class TS_DELTA_SHORT(PacketType):
encoding = bits[3:0] == 0b1000
delta = bits[7:4]
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 TS_DELTA_OR_MARK_L4(PacketType): # Layout 4: 48->64 bits
encoding = bits[6:0] == 0b0000001
delta = bits[63:12]
bit7 = bits[7:7]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool((self.bit9 and not self.bit8) or self.bit7)
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[47:7]
class TS_DELTA_S5_W2_L4(PacketType): # Layout 4: 48->40 bits
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[39:7]
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
@@ -155,11 +225,6 @@ class WAVERDY(PacketType): # exclude: 1 << 3
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]
@@ -181,29 +246,36 @@ class WAVESTART(PacketType): # exclude: 1 << 4
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
class WAVESTART_L4(PacketType): # Layout 4 has wave field at different position
encoding = bits[4:0] == 0b01100
delta = bits[6:5]
_padding = bits[47:7]
flag7 = bits[7:7]
simd = bits[9:8]
cu_lo = bits[12:10]
wave = bits[19:15]
id7 = bits[31:20]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
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
class WAVEALLOC_L4(PacketType): # Layout 4: 20->24 bits
encoding = bits[4:0] == 0b00101
delta = bits[7:5]
_padding = bits[51:8]
_padding = bits[23: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 PERF_L4(PacketType): # Layout 4: 28->32 bits
encoding = bits[4:0] == 0b10110
delta = bits[9:7]
arg = bits[31:10]
class NOP(PacketType):
encoding = bits[3:0] == 0b0000
@@ -244,14 +316,6 @@ class SNAPSHOT(PacketType):
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
@@ -271,56 +335,54 @@ class INST(PacketType):
wave = bits[12:8]
op = bits[19:13].enum(InstOp)
class INST_L4(PacketType): # Layout 4: different delta position and InstOp encoding
encoding = bits[2:0] == 0b010
delta = bits[5:3]
flag1 = bits[6:6]
flag2 = bits[7:7]
wave = bits[12:8]
op = bits[19:13].enum(InstOpL4)
class UTILCTR(PacketType):
encoding = bits[6:0] == 0b0110001
delta = bits[8:7]
ctr = bits[47:9]
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
PACKET_TYPES: list[type[PacketType]] = [
EVENT, EVENT_BIG,
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
VALUINST, INST,
NOP,
]
# Packet types with rocprof type IDs as keys
PACKET_TYPES_L3: dict[int, type[PacketType]] = {
1: VALUINST, 2: VMEMEXEC, 3: ALUEXEC, 4: IMMEDIATE, 5: IMMEDIATE_MASK, 6: WAVERDY, 7: TS_DELTA_S8_W3, 8: WAVEEND,
9: WAVESTART, 10: TS_DELTA_S5_W2, 11: WAVEALLOC, 12: TS_DELTA_S5_W3, 13: PERF, 14: UTILCTR, 15: TS_DELTA_SHORT,
16: NOP, 17: TS_WAVE_STATE, 18: EVENT, 19: EVENT_BIG, 20: REG, 21: SNAPSHOT, 22: TS_DELTA_OR_MARK, 23: LAYOUT_HEADER, 24: INST,
}
PACKET_TYPES_L4: dict[int, type[PacketType]] = {
**PACKET_TYPES_L3,
7: TS_DELTA_S8_W3_L4, 9: WAVESTART_L4, 10: TS_DELTA_S5_W2_L4, 11: WAVEALLOC_L4,
12: TS_DELTA_S5_W3_L4, 13: PERF_L4, 22: TS_DELTA_OR_MARK_L4, 24: INST_L4,
}
def _build_decode_tables(packet_types: dict[int, type[PacketType]]) -> tuple[dict[int, tuple], bytes]:
# Build state table: byte -> opcode. Sort by mask specificity (more bits first), NOP last
sorted_types = sorted(packet_types.items(), key=lambda x: (-bin(x[1].encoding.mask).count('1'), x[0] == 16))
state_table = bytes(next((op for op, cls in sorted_types if (b & cls.encoding.mask) == cls.encoding.default), 16) for b in range(256))
# Build decode info: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8)
decode_info = {}
for opcode, pkt_cls in packet_types.items():
delta_field = getattr(pkt_cls, 'delta', None)
special = {22: 1, 15: 2}.get(opcode, 0) # TS_DELTA_OR_MARK=22, TS_DELTA_SHORT=15
decode_info[opcode] = (pkt_cls, pkt_cls._size_nibbles, delta_field.lo if delta_field else 0, delta_field.mask if delta_field else 0, special)
return decode_info, state_table
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
for byte_val in range(256):
for opcode, pkt_cls in enumerate(PACKET_TYPES):
if (byte_val & pkt_cls.encoding.mask) == pkt_cls.encoding.default:
table[byte_val] = opcode
break
return bytes(table), opcode_to_class
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
# Precompute special case opcodes
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
_DECODE_INFO: dict[int, tuple] = {}
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
_delta_field = getattr(_pkt_cls, 'delta', None)
_delta_lo = _delta_field.lo if _delta_field else 0
_delta_mask = _delta_field.mask if _delta_field else 0
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
_DECODE_INFO_L3, _STATE_TABLE_L3 = _build_decode_tables(PACKET_TYPES_L3)
_DECODE_INFO_L4, _STATE_TABLE_L4 = _build_decode_tables(PACKET_TYPES_L4)
# ═══════════════════════════════════════════════════════════════════════════════
# DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def decode(data: bytes) -> Iterator[PacketType]:
"""Decode raw SQTT blob, yielding packet instances."""
"""Decode raw SQTT blob, yielding packet instances. Auto-detects layout from LAYOUT_HEADER."""
n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
decode_info, state_table = _DECODE_INFO_L3, _STATE_TABLE_L3 # default to layout 3, will update after seeing LAYOUT_HEADER
while pos + ((nib_count + nib_off + 1) >> 1) <= n:
need = nib_count - nib_off
@@ -333,13 +395,20 @@ def decode(data: bytes) -> Iterator[PacketType]:
# 3. if odd, read low nibble
if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
opcode = STATE_TO_OPCODE[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = _DECODE_INFO[opcode]
opcode = state_table[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
delta = (reg >> delta_lo) & delta_mask
if special == 1 and (reg >> 9) & 1 and not (reg >> 8) & 1: delta = 0 # TS_DELTA_OR_MARK marker
if special == 1: # TS_DELTA_OR_MARK
pkt = pkt_cls.from_raw(reg, 0) # create packet to check is_marker
if pkt.is_marker: delta = 0
elif special == 2: delta += 8 # TS_DELTA_SHORT
time += delta
yield pkt_cls.from_raw(reg, time)
pkt = pkt_cls.from_raw(reg, time)
# detect layout from first LAYOUT_HEADER and switch decode tables if needed
# NOTE: CDNA uses a completely different 16-bit header format, not nibbles - not supported here
if pkt_cls is LAYOUT_HEADER and pkt.layout == 4:
decode_info, state_table = _DECODE_INFO_L4, _STATE_TABLE_L4
yield pkt
# ═══════════════════════════════════════════════════════════════════════════════
# PRINTER
@@ -355,23 +424,26 @@ PACKET_COLORS = {
def format_packet(p) -> str:
from tinygrad.helpers import colored
name = type(p).__name__
if isinstance(p, INST):
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
if isinstance(p, (INST, INST_L4)):
op_name = p.op.name if isinstance(p.op, (InstOp, InstOpL4)) 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 isinstance(p, (WAVESTART, WAVESTART_L4, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
elif hasattr(p, '_fields'):
filt = {'delta', 'encoding'} if not isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_L4)) else {'encoding'}
fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
for k in p._fields if not k.startswith('_') and k not in {'delta', 'encoding'})
for k in p._fields if not k.startswith('_') and k not in filt)
else: fields = ""
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name.replace('_L4', ''), 'white'))} {fields}"
def print_packets(packets) -> None:
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
from tinygrad.helpers import getenv
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"} if not getenv("NOSKIP") else {"NOP"}
for p in packets:
if type(p).__name__ not in skip: print(format_packet(p))
if type(p).__name__.replace("_L4", "") not in skip: print(format_packet(p))
if __name__ == "__main__":
import sys, pickle
+161
View File
@@ -0,0 +1,161 @@
"""SQTT (SQ Thread Trace) packet decoder for CDNA/MI300 GPUs.
CDNA uses a completely different 16-bit header format from RDNA's nibble-based encoding.
"""
from __future__ import annotations
from typing import Iterator
from extra.assembly.amd.dsl import bits
from extra.assembly.amd.sqtt import PacketType
# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
class CDNA_DELTA(PacketType):
"""pkt_fmt=0: 16-bit timestamp delta packet"""
encoding = bits[3:0] == 0
delta = bits[11:4] # (data >> 4) & 0xff
unk_0 = bits[12:12] # (data >> 0xc) & 1
unk_1 = bits[15:13] # (data >> 0xd)
class CDNA_TIMESTAMP(PacketType):
"""pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
encoding = bits[3:0] == 1
unk_0 = bits[15:4]
timestamp = bits[63:16] # stored as (data_word >> 0x10) in low 46 bits of local_58
class CDNA_PKT_2(PacketType):
"""pkt_fmt=2: 64-bit packet (case 0x4)"""
encoding = bits[3:0] == 2
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[7:7] # (data >> 7) + 1 & 1
unk_padding = bits[63:8]
class CDNA_WAVESTART(PacketType):
"""pkt_fmt=3: 32-bit WAVESTART packet (case 0x8)"""
encoding = bits[3:0] == 3
unk_0 = bits[5:5] # (data >> 5) & 1
unk_1 = bits[9:6] # (data >> 6) & 0xf
wave = bits[13:10] # (data >> 10) & 0xf
simd = bits[15:14] # (data >> 0xe) & 3
cu = bits[17:16] # (data >> 0x10) & 3
unk_5 = bits[19:18] # (data >> 0x12) & 3
unk_6 = bits[28:22] # (data >> 0x16) & 0x7f
unk_padding = bits[31:29]
class CDNA_PKT_4(PacketType):
"""pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
encoding = bits[3:0] == 4
unk_0 = bits[5:5] # (data_word >> 5) & 1
unk_1 = bits[9:6] # (data_word >> 6) & 0xf
unk_2 = bits[13:10] # (data_word >> 10) & 0xf
unk_3 = bits[15:14] # (data_word >> 0xe)
class CDNA_PKT_5(PacketType):
"""pkt_fmt=5: 48-bit packet (case 0x10)"""
encoding = bits[3:0] == 5
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[7:7] # (data >> 7) + 1 & 1
unk_2 = bits[15:9] # (data >> 9) & 0x7f
unk_padding = bits[47:16]
class CDNA_WAVEEND(PacketType):
"""pkt_fmt=6: 16-bit WAVEEND packet (case 0x14, same as 0x8/0xc)"""
encoding = bits[3:0] == 6
unk_0 = bits[5:5] # (data_word >> 5) & 1
unk_1 = bits[9:6] # (data_word >> 6) & 0xf
wave = bits[13:10] # (data_word >> 10) & 0xf
simd = bits[15:14] # (data_word >> 0xe)
class CDNA_EXEC(PacketType):
"""pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
encoding = bits[3:0] == 10
unk_0 = bits[8:5] # (data_word >> 5) & 0xf
unk_1 = bits[10:9] # (data_word >> 9) & 3
unk_2 = bits[15:11] # (data_word >> 0xb)
class CDNA_PKT_11(PacketType):
"""pkt_fmt=11: 64-bit packet (case 0x28)"""
encoding = bits[3:0] == 11
unk_0 = bits[8:5] # (data_word >> 5) & 0xf
unk_1 = bits[10:9] # (data_word >> 9) & 3
unk_2 = bits[15:15] # (data_word >> 0xf) & 1
unk_padding = bits[63:16]
class CDNA_INST(PacketType):
"""pkt_fmt=13: 32-bit INST packet (case 0x30)"""
encoding = bits[3:0] == 13
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[9:8] # (data >> 8) & 3
unk_2 = bits[11:10] # (data >> 10) & 3
unk_3 = bits[13:12] # (data >> 0xc) & 3
unk_4 = bits[15:14] # (data >> 0xe) & 3
unk_5 = bits[19:18] # (data >> 0x12) & 3
unk_6 = bits[21:20] # (data >> 0x14) & 3
unk_7 = bits[23:22] # (data >> 0x16) & 3
unk_8 = bits[25:24] # (data >> 0x18) & 3
unk_9 = bits[27:26] # (data >> 0x1a) & 3
unk_padding = bits[31:28]
class CDNA_PKT_14(PacketType):
"""pkt_fmt=14: 64-bit packet (case 0x34)"""
encoding = bits[3:0] == 14
unk_0 = bits[5:5] # (data >> 5) & 1
unk_1 = bits[9:6] # (data >> 6) & 0xf
unk_2 = bits[11:10] # (data >> 10) & 3
unk_3 = bits[24:12] # (data >> 0xc) & 0x1fff
unk_4 = bits[37:25] # (data >> 0x19) & 0x1fff
unk_5 = bits[50:38] # (data >> 0x26) & 0x1fff
unk_6 = bits[51:51] # (data >> 0x33) & 1
unk_padding = bits[63:52]
class CDNA_PKT_15(PacketType):
"""pkt_fmt=15: 48-bit packet (case 0x38, same as 0x10)"""
encoding = bits[3:0] == 15
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[7:7] # (data >> 7) + 1 & 1
unk_2 = bits[15:9] # (data >> 9) & 0x7f
unk_padding = bits[47:16]
CDNA_PKT_TYPES: dict[int, type[PacketType]] = {
0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4,
5: CDNA_PKT_5, 6: CDNA_WAVEEND, 10: CDNA_EXEC, 11: CDNA_PKT_11, 13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
}
# Validate CDNA packet definitions
for pkt_fmt, pkt_cls in CDNA_PKT_TYPES.items():
assert pkt_cls.encoding.default == pkt_fmt, f"{pkt_cls.__name__} encoding {pkt_cls.encoding.default} != pkt_fmt {pkt_fmt}"
assert CDNA_PKT_SIZES[pkt_fmt] * 2 == pkt_cls._size_nibbles, f"{pkt_cls.__name__} size {pkt_cls._size_nibbles//2} != {CDNA_PKT_SIZES[pkt_fmt]}"
def decode(data: bytes) -> Iterator[PacketType]:
"""Decode CDNA SQTT blob using 16-bit header format."""
pos, time, ts_offset = 0, 0, None
while pos + 2 <= len(data):
header = int.from_bytes(data[pos:pos+2], 'little')
pkt_fmt = header & 0xf
pkt_size = CDNA_PKT_SIZES[pkt_fmt]
if pos + pkt_size > len(data): break
raw = int.from_bytes(data[pos:pos+pkt_size], 'little')
# pkt_fmt=0 has delta in bits[11:4], accumulate it
if pkt_fmt == 0: time += ((raw >> 4) & 0xff) * 4
# pkt_fmt=1 with unk_0=0 is absolute timestamp - use it to anchor time
if pkt_fmt == 1 and ((raw >> 4) & 0xfff) == 0:
abs_ts = raw >> 16
if ts_offset is None: ts_offset = abs_ts - time # first timestamp: save offset
else: time = ((abs_ts - ts_offset) & ~3) - 4 # subsequent: compute time, align to 4, subtract 4
pkt_cls = CDNA_PKT_TYPES[pkt_fmt]
yield pkt_cls.from_raw(raw, time)
pos += pkt_size
if __name__ == "__main__":
import sys, pickle
if len(sys.argv) < 2:
print("Usage: python sqtt_cdna.py <pkl_file>")
sys.exit(1)
with open(sys.argv[1], "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
for i, event in enumerate(sqtt_events):
print(f"\n=== event {i} ===")
for pkt in decode(event.blob):
print(f"{pkt._time:8}: {pkt}")
+14 -24
View File
@@ -2,11 +2,8 @@
from dataclasses import dataclass
from typing import Iterator
from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd.sqtt import decode, print_packets, INST, VALUINST, IMMEDIATE, WAVESTART, WAVEEND, InstOp, PacketType, IMMEDIATE_MASK
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP, s_endpgm
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
@@ -16,19 +13,11 @@ class InstructionInfo:
wave: int
inst: Inst
def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
def map_insts(data:bytes, lib:bytes, target:int) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
"""maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
# map pcs to insts
pc_map:dict[int, Inst] = {}
image, sections, _ = elf_loader(lib)
text = next((sh for sh in sections if sh.name == ".text"), None)
assert text is not None, "no .text section found"
text_off, text_size = text.header.sh_addr, text.header.sh_size
offset = text_off
while offset < text_off + text_size:
inst = decode_inst(image[offset:])
pc_map[offset-text_off] = inst
offset += inst.size()
from tinygrad.viz.serve import amd_decode
pc_map = amd_decode(lib, target)
wave_pc:dict[int, int] = {}
# only processing packets on one [CU, SIMD] unit
@@ -37,7 +26,7 @@ def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionIn
if not simd_select(p): continue
if isinstance(p, WAVESTART):
assert p.wave not in wave_pc, "only one inflight wave per unit"
wave_pc[p.wave] = 0
wave_pc[p.wave] = next(iter(pc_map))
continue
if isinstance(p, WAVEEND):
pc = wave_pc.pop(p.wave)
@@ -80,22 +69,22 @@ def map_insts(data:bytes, lib:bytes) -> Iterator[tuple[PacketType, InstructionIn
# test to compare every packet with the rocprof decoder
def test_rocprof_inst_traces_match(sqtt, prg, target):
from tinygrad.viz.serve import llvm_disasm
from tinygrad.viz.serve import amd_decode
from extra.sqtt.roc import decode as roc_decode, InstExec
disasm = {addr+prg.base:inst_disasm for addr, inst_disasm in llvm_disasm(target, prg.lib).items()}
rctx = roc_decode([sqtt], {prg.name:disasm})
rwaves = rctx.inst_execs[(sqtt.kern, sqtt.exec_tag)]
addr_table = amd_decode(prg.lib, target)
disasm = {addr+prg.base:(inst.disasm(), inst.size()) for addr,inst in addr_table.items()}
rctx = roc_decode([sqtt], {prg.tag:disasm})
rwaves = rctx.inst_execs.get((sqtt.kern, sqtt.exec_tag), [])
rwaves_iter:dict[int, list[Iterator[InstExec]]] = {} # wave unit (0-15) -> list of inst trace iterators for all executions on that unit
for w in rwaves: rwaves_iter.setdefault(w.wave_id, []).append(w.unpack_insts())
rwaves_base = next(iter(disasm)) # base program counter
passed_insts = 0
for pkt, info in map_insts(sqtt.blob, prg.lib):
for pkt, info in map_insts(sqtt.blob, prg.lib, target):
if DEBUG >= 2: print_packets([pkt])
if info is None: continue
if DEBUG >= 2: print(f"{' '*29}{info.inst.disasm()}")
rocprof_inst = next(rwaves_iter[info.wave][0])
ref_pc = rocprof_inst.pc-rwaves_base
ref_pc = rocprof_inst.pc-prg.base
# always check pc matches
assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm[rocprof_inst.pc][0]} != {info.pc}:{info.inst.disasm()}"
# special handling for s_endpgm, it marks the wave completion.
@@ -110,7 +99,8 @@ def test_rocprof_inst_traces_match(sqtt, prg, target):
for k,v in rwaves_iter.items():
assert len(v) == 0, f"incomplete wave {k}"
print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")
if len(rwaves):
print(f"passed for {passed_insts} instructions across {len(rwaves)} waves scheduled on {len(rwaves_iter)} wave units")
if __name__ == "__main__":
import argparse, pickle, pathlib
@@ -123,7 +113,7 @@ if __name__ == "__main__":
with open(args.profile, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
kern_events = {e.name:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
kern_events = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
target = next((e for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.device.startswith("AMD"))).props["gfx_target_version"]
for e in sqtt_events:
if args.kernel is not None and args.kernel != e.kern: continue
+97 -23
View File
@@ -6,7 +6,9 @@ from pathlib import Path
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
from extra.assembly.amd.emu import run_asm as python_run_asm, decode_program
from extra.assembly.amd import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP, SOPPOp
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
@@ -64,6 +66,55 @@ def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size
return sum(times) / len(times)
def profile_instructions(kernel: bytes):
"""Profile individual instruction compile times."""
from extra.assembly.amd.emu import _get_runner, _canonical_runner_cache
from tinygrad.helpers import Context
_get_runner.cache_clear()
_canonical_runner_cache.clear()
results = []
i = 0
while i < len(kernel):
inst = decode_inst(kernel[i:])
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
inst_bytes = bytes(kernel[i:i + inst.size() + 4])
try: inst_str = repr(inst)
except Exception: inst_str = f"<{type(inst).__name__}>"
# Time the full compile (sink + render + compile)
start = time.perf_counter()
with Context(CCACHE=0):
runner, is_new = _get_runner(inst_bytes)
compile_time = time.perf_counter() - start
results.append({
'inst_str': inst_str + ('' if is_new else ' [CACHED]'),
'compile_ms': compile_time * 1000 if is_new else 0,
})
i += inst.size()
return sorted(results, key=lambda x: x['compile_ms'], reverse=True)
def benchmark_python_split(kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
"""Benchmark Python emulator with compile and execution times."""
from extra.assembly.amd.emu import _get_runner, _canonical_runner_cache
from tinygrad.helpers import Context
_get_runner.cache_clear()
_canonical_runner_cache.clear()
decode_program.cache_clear()
# Measure compile time (decode_program builds sinks, renders, and compiles)
compile_start = time.perf_counter()
with Context(CCACHE=0):
program = decode_program(kernel)
compile_time = time.perf_counter() - compile_start
n_compiled = len(_canonical_runner_cache)
# Execution time
exec_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, iterations)
return compile_time, exec_time, len(program), n_compiled
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:
@@ -119,14 +170,36 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
print(f" Error getting kernel: {e}")
return None
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "sin", "gelu", "matmul_small"]
def main():
import argparse
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
parser.add_argument("--profile", type=str, default=None, help="Profile instructions for a specific kernel (e.g. 'sin')")
parser.add_argument("--top", type=int, default=20, help="Number of top instructions to show in profile")
args = parser.parse_args()
# Profile mode: show individual instruction timing
if args.profile:
kernel_info = get_tinygrad_kernel(args.profile)
if kernel_info is None:
print(f"Failed to get kernel for '{args.profile}'")
return
kernel = kernel_info[0]
print(f"Profiling instructions for '{args.profile}' kernel...")
print("=" * 110)
results = profile_instructions(kernel)
print(f"{'Instruction':<90} {'Compile(ms)':>12}")
print("-" * 110)
for r in results[:args.top]:
inst = r['inst_str'][:87] + "..." if len(r['inst_str']) > 90 else r['inst_str']
print(f"{inst:<90} {r['compile_ms']:>12.3f}")
print("-" * 110)
total = sum(r['compile_ms'] for r in results)
print(f"{'TOTAL':<90} {total:>12.3f}")
return
rust_remu = get_rust_remu()
if rust_remu is None:
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
@@ -149,44 +222,45 @@ def main():
continue
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
n_insts = count_instructions(kernel)
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
# Benchmark Python emulator (must be first to measure compile time before cache is populated)
py_compile, py_exec, n_insts, n_compiled = benchmark_python_split(kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
print(f"{n_insts} insts ({n_compiled} unique) × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if py_compile is not None:
py_exec_rate = total_work / py_exec / 1e6
print(f" Compile: {py_compile*1000:8.3f} ms ({n_compiled} unique)")
print(f" Exec: {py_exec*1000:8.3f} ms ({py_exec_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]")
speedup = py_exec / rust_time if py_exec else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
results.append((op_name, n_insts, n_compiled, n_workgroups, py_compile, py_exec, rust_time))
# Summary table
print("\n" + "=" * 90)
print("\n" + "=" * 110)
print("SUMMARY")
print("=" * 90)
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print("-" * 90)
print("=" * 110)
print(f"{'Name':<16} {'Insts':<6} {'Unique':<6} {'WGs':<5} {'Compile (ms)':<14} {'Exec (ms)':<12} {'Rust (ms)':<12} {'Speedup':<10}")
print("-" * 110)
for name, n_insts, n_wgs, py_time, rust_time in results:
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
for name, n_insts, n_compiled, n_wgs, py_compile, py_exec, rust_time in results:
compile_ms = f"{py_compile*1000:.3f}" if py_compile else "error"
exec_ms = f"{py_exec*1000:.3f}" if py_exec 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"
speedup = f"{py_exec/rust_time:.1f}x" if py_exec else "N/A"
else:
rust_ms, speedup = "N/A", "N/A"
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
print(f"{name:<16} {n_insts:<6} {n_compiled:<6} {n_wgs:<5} {compile_ms:<14} {exec_ms:<12} {rust_ms:<12} {speedup:<10}")
if __name__ == "__main__":
main()
@@ -1,196 +0,0 @@
# Usability tests for the RDNA3 ASM DSL
# These tests demonstrate how the DSL *should* work for a good user experience
# Currently many of these tests fail - they document desired behavior
import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
class TestReprReadability(unittest.TestCase):
"""
Issue: repr() leaks internal RawImm type and omits zero-valued fields.
When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
VOP1(op=1, src0=RawImm(257))
Problems:
1. vdst=v[0] is omitted because 0 is treated as "default"
2. src0 shows RawImm(257) instead of v[1]
3. User sees encoded values (257 = 256 + 1) instead of register names
Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
"""
def test_repr_shows_registers_not_raw_imm(self):
inst = v_mov_b32_e32(v[0], v[1])
# Should show v[1], not RawImm(257)
self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
self.assertIn("v[1]", repr(inst), "repr should show register name")
def test_repr_includes_zero_dst(self):
inst = v_mov_b32_e32(v[0], v[1])
# v[0] is a valid destination register, should be shown
self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
def test_repr_roundtrip(self):
# repr should produce something that can be eval'd back
inst = v_mov_b32_e32(v[0], v[1])
# This would require repr to output valid Python, e.g.:
# "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
r = repr(inst)
# At minimum, it should be human-readable
self.assertIn("v[", r, "repr should show register syntax")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
class TestVOPDHelperSignature(unittest.TestCase):
"""
Issue: VOPD helper functions have confusing semantics.
v_dual_mul_f32 is defined as:
v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
which is 'opx'. So v_dual_mul_f32 sets the X operation.
But then test_dual_mul in test_handwritten.py does:
v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
X and Y operations the same. This is confusing because:
1. The function name suggests it handles the X operation
2. But you still pass an opcode as the first arg (which becomes opy)
Expected: Either make the helper fully specify both ops, or make the
signature clearer about what the positional arg means.
"""
def test_vopd_helper_opy_should_be_required(self):
# Using only keyword args "works" but opy silently defaults to 0
inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
# Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
# This test documents the bug - it should fail once fixed
self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
def test_vopd_helper_positional_arg_is_opy(self):
# The first positional arg after the partial becomes opy, not a second opx
inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32) # From partial
self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32) # From first positional arg
class TestFieldAccessPreservesType(unittest.TestCase):
"""
Issue: Field access loses type information.
After creating an instruction, accessing fields returns encoded int values:
inst = v_mov_b32_e32(v[0], v[1])
inst.vdst # returns 0, not VGPR(0)
This makes it impossible to round-trip register types through field access.
"""
def test_vdst_returns_register(self):
inst = v_mov_b32_e32(v[5], v[1])
vdst = inst.vdst
# Should return a VGPR, not an int
self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
# Ideally: self.assertIsInstance(vdst, VGPR)
def test_src_returns_register_for_vgpr_source(self):
inst = v_mov_b32_e32(v[0], v[1])
# src0 is encoded as 257 (256 + 1 for v1)
# Ideally it should decode back to v[1]
src0_raw = inst._values.get('src0')
# Currently returns RawImm(257), should return VGPR(1) or similar
self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
class TestArgumentDiscoverability(unittest.TestCase):
"""
Issue: No clear signature for positional arguments.
inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
Users have no way to know the argument order without reading source code.
The order is implicitly defined by the class field definition order.
Possible fixes:
1. Add explicit parameter names to functools.partial
2. Generate type stubs with proper signatures
3. Add docstrings listing the expected arguments
"""
def test_signature_has_named_params(self):
import inspect
sig = inspect.signature(s_load_b128)
params = list(sig.parameters.keys())
# Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
# Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
self.assertIn('sdata', params, "signature should show field names")
class TestSpecialConstants(unittest.TestCase):
"""
Issue: NULL and other constants are IntEnum values that might be confusing.
NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
that clearly represents "no register" rather than a magic number.
"""
def test_null_has_clear_repr(self):
# NULL should have a clear string representation
self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
def test_null_is_distinguishable_from_int(self):
# NULL should be distinguishable from the raw integer 124
self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
if __name__ == "__main__":
unittest.main()
+2 -43
View File
@@ -5,6 +5,7 @@ from dataclasses import dataclass
@dataclass
class KernelInfo:
code: bytes
src: str
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
@@ -34,46 +35,4 @@ TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items
def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
def get_mattr(arch:str) -> str:
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION CONTEXT (for testing compiled pseudocode)
# ═══════════════════════════════════════════════════════════════════════════════
class ExecContext:
"""Context for running compiled pseudocode in tests."""
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
self.D0, self.D1 = Reg(d0), Reg(0)
self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
self.lane, self.laneId, self.literal = lane, lane, literal
self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
self.VGPR = vgprs if vgprs is not None else {}
self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
def run(self, code: str):
"""Execute compiled code."""
import extra.assembly.amd.pcode as pcode
ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
# Also include underscore-prefixed helpers that compiled pseudocode uses
for k in ['_pack', '_pack32']:
if hasattr(pcode, k): ns[k] = getattr(pcode, k)
ns.update({
'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
'tmp': self.tmp, 'saveexec': self.saveexec,
'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
})
exec(code, ns)
def _sync(ctx_reg, ns_val):
if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
else: ctx_reg._val = int(ns_val) & self._MASK64
for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
+36 -6
View File
@@ -13,7 +13,7 @@ def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
# f16 conversion helpers
def _f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
def f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
def f32_to_f16(f: float) -> int:
f = float(f)
if math.isnan(f): return 0x7e00
@@ -43,6 +43,23 @@ VCC = VCC_LO # For VOP3SD sdst field (VCC_LO is exported from dsl)
USE_HW = os.environ.get("USE_HW", "0") == "1"
FLOAT_TOLERANCE = 1e-5
def get_gpu_target() -> tuple[int, int, int]:
"""Get the GPU target as (major, minor, stepping) tuple."""
if not USE_HW: return (0, 0, 0)
from tinygrad.device import Device
return Device["AMD"].target
def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
"""Skip test if GPU target is below the minimum required version."""
import unittest
def decorator(test_func):
if not USE_HW: return test_func
target = get_gpu_target()
if target[0] < min_major or (target[0] == min_major and target[1] < min_minor):
return unittest.skip(reason or f"requires gfx{min_major}{min_minor}0+")(test_func)
return test_func
return decorator
# Output buffer layout: vgpr[16][32], sgpr[16], vcc, scc, exec
N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
VGPR_BYTES = N_VGPRS * WAVE_SIZE * 4 # 16 regs * 32 lanes * 4 bytes = 2048
@@ -141,7 +158,7 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
rsrc2 = 0x19c | (128 << 15)
scratch_size = 0x10000 # 64KB per lane, matches .amdhsa_private_segment_fixed_size in run_program_hw
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2, scratch_size)
assert result == 0, f"run_asm failed with {result}"
return parse_output(bytes(out_buf), n_lanes)
@@ -204,6 +221,7 @@ amdhsa.kernels:
prg = AMDProgram(dev, "test", lib)
out_gpu = dev.allocator.alloc(OUT_BYTES)
assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
out_buf = bytearray(OUT_BYTES)
@@ -211,8 +229,12 @@ amdhsa.kernels:
return parse_output(bytes(out_buf), n_lanes)
def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS) -> list[str]:
"""Compare two WaveStates and return list of differences."""
def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS, ulp_tolerance: int = 0) -> list[str]:
"""Compare two WaveStates and return list of differences.
Args:
ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
"""
import math
diffs = []
for i in range(n_vgprs):
@@ -223,6 +245,11 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
emu_f, hw_f = _f32(emu_val), _f32(hw_val)
if math.isnan(emu_f) and math.isnan(hw_f):
continue
# Check ULP difference for floats (only for same-sign values)
if ulp_tolerance > 0 and (emu_val < 0x80000000) == (hw_val < 0x80000000):
ulp_diff = abs(int(emu_val) - int(hw_val))
if ulp_diff <= ulp_tolerance:
continue
diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
for i in range(N_SGPRS):
emu_val = emu_st.sgpr[i]
@@ -235,16 +262,19 @@ def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgp
diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
return diffs
def run_program(instructions: list, n_lanes: int = 1) -> WaveState:
def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) -> WaveState:
"""Run instructions and return WaveState.
If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
Otherwise, runs only on emulator.
Args:
ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
"""
emu_st = run_program_emu(instructions, n_lanes)
if USE_HW:
hw_st = run_program_hw(instructions, n_lanes)
diffs = compare_wave_states(emu_st, hw_st, n_lanes)
diffs = compare_wave_states(emu_st, hw_st, n_lanes, ulp_tolerance=ulp_tolerance)
if diffs:
raise AssertionError(f"Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
return hw_st
+86
View File
@@ -138,6 +138,50 @@ class TestDS2AddrMore(unittest.TestCase):
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSB96(unittest.TestCase):
"""Tests for DS_STORE_B96 and DS_LOAD_B96 (96-bit / 3 dwords)."""
def test_ds_store_load_b96(self):
"""DS_STORE_B96 stores 3 VGPRs, DS_LOAD_B96 loads them back."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
ds_store_b96(addr=v[10], data0=v[0:2]),
s_waitcnt(lgkmcnt=0),
ds_load_b96(addr=v[10], vdst=v[4:6]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
def test_ds_store_b96_with_offset(self):
"""DS_STORE_B96 with non-zero offset."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
DS(DSOp.DS_STORE_B96, addr=v[10], data0=v[0:2], offset0=12),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B96, addr=v[10], vdst=v[4:6], offset0=12),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
class TestDSB128(unittest.TestCase):
"""Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
@@ -675,5 +719,47 @@ class TestAtomicOrdering(unittest.TestCase):
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
class TestDsPermute(unittest.TestCase):
"""Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
def test_ds_permute_b32_identity(self):
"""DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
# For simplicity, test with single lane
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (lane 0)
v_mov_b32_e32(v[1], 0xDEADBEEF), # data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
def test_ds_bpermute_b32_identity(self):
"""DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (read from lane 0)
v_mov_b32_e32(v[1], 0xCAFEBABE), # data in lane 0
ds_bpermute_b32(v[2], v[0], v[1]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 reads from lane 0's v[1]
self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
def test_ds_permute_b32_broadcast(self):
"""DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
# With 4 lanes, all sending to lane 0, highest lane wins
instructions = [
v_mov_b32_e32(v[0], 0), # All lanes send to addr 0 (lane 0)
v_mov_b32_e32(v[1], 0x11111111), # All lanes send same data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=4)
# Lane 0 receives data (highest numbered active lane wins)
self.assertEqual(st.vgpr[0][2], 0x11111111)
if __name__ == '__main__':
unittest.main()
+152
View File
@@ -523,5 +523,157 @@ class TestD16HiLoads(unittest.TestCase):
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
class TestGlobalOffset(unittest.TestCase):
"""Tests for GLOBAL instructions with different offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load/store wrong data.
"""
def test_global_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
# Load from offset 100 -> should get 0xAAAAAAAA
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[3], saddr=SrcEnum.NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], v[4]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_store_different_offsets(self):
"""Store to two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0x11111111 at offset 300
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=300),
# Store 0x22222222 at offset 400
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b32(addr=v[0:1], data=v[3], saddr=SrcEnum.NULL, offset=400),
s_waitcnt(vmcnt=0),
# Load back to verify
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=300),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[5], saddr=SrcEnum.NULL, offset=400),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x11111111, f"offset 300: expected 0x11111111, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0x22222222, f"offset 400: expected 0x22222222, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_no_saddr(self):
"""Test negative offset without saddr (VGPR pair for address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from vaddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from vaddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
# vaddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[6], saddr=SrcEnum.NULL, offset=-100),
# vaddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[7], saddr=SrcEnum.NULL, offset=-100),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[4], 0),
v_mov_b32_e32(v[5], 0),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_with_saddr(self):
"""Test negative offset with saddr (SGPR pair for base address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from saddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from saddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
# saddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[6], saddr=s[4:5], offset=-100),
# saddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[7], saddr=s[4:5], offset=-100),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
if __name__ == '__main__':
unittest.main()
+355
View File
@@ -0,0 +1,355 @@
"""Tests for SCRATCH instructions - scratch (private) memory operations.
Includes: scratch_load_*, scratch_store_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestScratchStore(unittest.TestCase):
"""Tests for SCRATCH store instructions."""
def test_scratch_store_b32_basic(self):
"""SCRATCH_STORE_B32 stores 32-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
# Store via scratch
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Load back via scratch
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
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)
def test_scratch_store_b64_basic(self):
"""SCRATCH_STORE_B64 stores 64-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
s_mov_b32(s[5], 0xCAFEBABE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
scratch_store_b64(addr=v[0], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b64(addr=v[0], vdst=v[4:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
def test_scratch_store_b8_basic(self):
"""SCRATCH_STORE_B8 stores single byte to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# First store full word
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Store single byte
v_mov_b32_e32(v[2], 0x42),
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Load back
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Only byte 0 should change from 0xEF to 0x42
self.assertEqual(st.vgpr[0][0], 0xDEADBE42)
def test_scratch_store_b16_basic(self):
"""SCRATCH_STORE_B16 stores 16-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
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], 0xDEADCAFE)
class TestScratchLoad(unittest.TestCase):
"""Tests for SCRATCH load instructions."""
def test_scratch_load_b96(self):
"""SCRATCH_LOAD_B96 loads 96-bit value correctly."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[4]),
s_mov_b32(s[4], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[4]),
scratch_store_b96(addr=v[0], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b96(addr=v[0], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[5]),
v_mov_b32_e32(v[1], v[6]),
v_mov_b32_e32(v[2], v[7]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][2], 0xCCCCCCCC)
def test_scratch_load_b128(self):
"""SCRATCH_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
s_mov_b32(s[4], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[4]),
s_mov_b32(s[4], 0x12345678),
v_mov_b32_e32(v[4], s[4]),
s_mov_b32(s[4], 0x9ABCDEF0),
v_mov_b32_e32(v[5], s[4]),
scratch_store_b128(addr=v[0], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b128(addr=v[0], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[2], v[8]),
v_mov_b32_e32(v[3], v[9]),
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)
self.assertEqual(st.vgpr[0][2], 0x12345678)
self.assertEqual(st.vgpr[0][3], 0x9ABCDEF0)
def test_scratch_load_u8(self):
"""SCRATCH_LOAD_U8 loads unsigned byte with zero extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADBEAB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_u8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
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], 0xAB)
def test_scratch_load_i8(self):
"""SCRATCH_LOAD_I8 loads signed byte with sign extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0x80), # -128 as signed byte
v_mov_b32_e32(v[2], s[4]),
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_i8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
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], 0xFFFFFF80)
def test_scratch_load_u16(self):
"""SCRATCH_LOAD_U16 loads unsigned 16-bit with zero extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADCAFE),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_u16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
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], 0xCAFE)
def test_scratch_load_i16(self):
"""SCRATCH_LOAD_I16 loads signed 16-bit with sign extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0x8000), # -32768 as signed 16-bit
v_mov_b32_e32(v[2], s[4]),
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_i16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
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], 0xFFFF8000)
class TestScratchSVE(unittest.TestCase):
"""Tests for SCRATCH SVE (Scratch VGPR Enable) bit behavior."""
def test_scratch_sve_zero_ignores_vaddr(self):
"""With SVE=0, VADDR should be ignored in address calculation."""
TEST_OFFSET = 256
# Store a marker value at offset 256 (where SVE=0 should go)
# Then set v[0] to a non-zero value (100) and store via scratch with SVE=0
# If SVE=0 is handled correctly, the VADDR (100) should be IGNORED,
# and the store should go to offset 256, not 256+100=356
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt(vmcnt=0),
# Now set v[0] to 100 (non-zero) and store 0xBBBBBBBB with SVE=0
# With SVE=0, v[0] should be IGNORED, so this should overwrite offset 256
v_mov_b32_e32(v[0], 100),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt(vmcnt=0),
# Load back from offset 256 (with v[0]=0) - should get 0xBBBBBBBB
v_mov_b32_e32(v[0], 0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# If SVE=0 works correctly, v[0] should be 0xBBBBBBBB (the second store overwrote the first)
# If SVE=0 is wrong (VADDR used), v[0] would be 0xAAAAAAAA (stores went to different locations)
self.assertEqual(st.vgpr[0][0], 0xBBBBBBBB, "SVE=0 should ignore VADDR, both stores should go to same location")
def test_scratch_sve_one_uses_vaddr(self):
"""With SVE=1, VADDR should be used as offset in address calculation."""
TEST_OFFSET = 256
# Store at offset 256 with v[0]=0, then store at offset 256 with v[0]=100 and SVE=1
# With SVE=1, the second store should go to 256+100=356, not 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt(vmcnt=0),
# Now set v[0] to 100 and store 0xBBBBBBBB with SVE=1
# With SVE=1, v[0] IS used, so this should go to offset 256+100=356
v_mov_b32_e32(v[0], 100),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt(vmcnt=0),
# Load back from offset 256 (with v[0]=0) - should still be 0xAAAAAAAA
v_mov_b32_e32(v[0], 0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# If SVE=1 works correctly, v[0] should be 0xAAAAAAAA (stores went to different locations)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, "SVE=1 should use VADDR, stores should go to different locations")
class TestScratchMultiLane(unittest.TestCase):
"""Tests for SCRATCH operations with multiple lanes."""
def test_scratch_store_load_multi_lane(self):
"""SCRATCH store/load works correctly with multiple lanes (private per-lane memory)."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# Each lane stores its lane ID
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[2], v[255]), # v[255] has packed workitem IDs, low 10 bits = x
v_and_b32_e32(v[2], 0x3FF, v[2]), # extract lane ID
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Load back
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=4)
# Each lane should have loaded its own lane ID
for lane in range(4):
self.assertEqual(st.vgpr[lane][0], lane, f"Lane {lane} should have value {lane}")
if __name__ == '__main__':
unittest.main()
+448
View File
@@ -0,0 +1,448 @@
"""Tests for SMEM instructions - scalar memory operations.
Includes: s_load_b32, s_load_b64, s_load_b128, s_load_b256, s_load_b512
Tests both immediate and register offset addressing modes.
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
# Use offset into output buffer for test data (output buffer is 2124 bytes)
TEST_OFFSET = 2000
# Cache invalidation sequence for scalar loads after vector stores
# s_wait_idle waits for all outstanding memory operations including cache flushes
CACHE_INV = [s_gl1_inv(), s_dcache_inv(), s_wait_idle()]
class TestSLoadRegisterOffset(unittest.TestCase):
"""Tests for s_load with register offset (soffset field).
Bug: s_load_b32(s[dst], s[base:base+1], s[off]) ignores the register offset
and only uses the immediate offset field. This causes incorrect memory loads
when the offset comes from a register.
"""
def test_s_load_b32_register_offset_basic(self):
"""s_load_b32 with register offset should load from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values to output buffer: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[4], 0xAAAAAAAA),
s_mov_b32(s[5], 0xBBBBBBBB),
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_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Now test s_load with register offset
# Put offset value in s[4]: offset = 4 bytes (1 dword)
s_mov_b32(s[4], 4),
# Load from out_ptr + TEST_OFFSET + s[4] (should load 0xBBBBBBBB)
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB,
f"s_load with reg offset 4 should load 0xBBBBBBBB: s[5]=0x{st.sgpr[5]:08x}")
def test_s_load_b32_register_offset_different_from_immediate(self):
"""s_load_b32 with register offset loads different data than immediate offset 0."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[4], 0xAAAAAAAA),
s_mov_b32(s[5], 0xBBBBBBBB),
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_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with immediate offset 0
s_load_b32(s[5], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Load with register offset 4
s_mov_b32(s[4], 4),
s_load_b32(s[6], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# s[5] has dword at offset 0 (0xAAAAAAAA), s[6] has dword at offset 4 (0xBBBBBBBB)
self.assertEqual(st.sgpr[5], 0xAAAAAAAA)
self.assertEqual(st.sgpr[6], 0xBBBBBBBB)
self.assertNotEqual(st.sgpr[5], st.sgpr[6],
f"s_load with reg offset 4 should load different value than offset 0: "
f"s[5]=0x{st.sgpr[5]:08x}, s[6]=0x{st.sgpr[6]:08x}")
def test_s_load_b32_register_offset_same_as_dst(self):
"""s_load_b32 where soffset register is same as destination.
This is the exact pattern that exposes the bug:
s_load_b32(s[8], s[2:3], s[8])
The offset should be read BEFORE the destination is overwritten.
"""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[6], 0xAAAAAAAA),
s_mov_b32(s[7], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[6]),
v_mov_b32_e32(v[3], s[7]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Set up s[4] = 4 (offset in bytes)
s_mov_b32(s[4], 4),
# Load using s[4] as both offset and destination
# Should load from base + 4, then store result in s[4]
s_load_b32(s[4], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Also load with immediate offset 4 for comparison
s_load_b32(s[5], s[2:3], NULL, offset=TEST_OFFSET+4),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# s[4] and s[5] should have the same value (both loaded from offset 4 = 0xBBBBBBBB)
self.assertEqual(st.sgpr[4], 0xBBBBBBBB)
self.assertEqual(st.sgpr[4], st.sgpr[5],
f"s_load with reg offset s[4]=4 should match immediate offset=4: "
f"s[4]=0x{st.sgpr[4]:08x}, s[5]=0x{st.sgpr[5]:08x}")
def test_s_load_b32_register_offset_zero(self):
"""s_load_b32 with register offset = 0 should be same as immediate offset 0."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test value: 0xDEADBEEF at offset
s_mov_b32(s[7], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[7]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 0
s_mov_b32(s[4], 0),
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Load with immediate offset 0
s_load_b32(s[6], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[5], 0xDEADBEEF)
self.assertEqual(st.sgpr[5], st.sgpr[6],
f"s_load with reg offset 0 should match immediate offset 0: "
f"s[5]=0x{st.sgpr[5]:08x}, s[6]=0x{st.sgpr[6]:08x}")
def test_s_load_b32_register_plus_immediate_offset(self):
"""s_load_b32 with both register and immediate offset should add them."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[8], 0xAAAAAAAA),
s_mov_b32(s[9], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[8]),
v_mov_b32_e32(v[3], s[9]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# reg offset = 4, imm offset = 0 -> total offset = 4
s_mov_b32(s[4], 4),
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# reg offset = 0, imm offset = 4 -> total offset = 4
s_mov_b32(s[6], 0),
s_load_b32(s[7], s[2:3], s[6], offset=TEST_OFFSET+4),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Both should load from offset 4 (0xBBBBBBBB)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB)
self.assertEqual(st.sgpr[7], 0xBBBBBBBB)
self.assertEqual(st.sgpr[5], st.sgpr[7],
f"reg_off=4 + imm_off=0 should equal reg_off=0 + imm_off=4: "
f"s[5]=0x{st.sgpr[5]:08x}, s[7]=0x{st.sgpr[7]:08x}")
class TestSLoadMultiDword(unittest.TestCase):
"""Tests for multi-dword s_load with register offset."""
def test_s_load_b64_register_offset(self):
"""s_load_b64 with register offset should load 2 dwords from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA, 0xBBBBBBBB at offset
s_mov_b32(s[10], 0xAAAAAAAA),
s_mov_b32(s[11], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[10]),
v_mov_b32_e32(v[3], s[11]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 0
s_mov_b32(s[4], 0),
s_load_b64(s[6:7], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Compare with immediate offset
s_load_b64(s[8:9], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], 0xAAAAAAAA)
self.assertEqual(st.sgpr[7], 0xBBBBBBBB)
self.assertEqual(st.sgpr[6], st.sgpr[8])
self.assertEqual(st.sgpr[7], st.sgpr[9])
def test_s_load_b128_register_offset(self):
"""s_load_b128 with register offset should load 4 dwords from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA, 0xBBBBBBBB, 0xCCCCCCCC, 0xDDDDDDDD at offset
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[14], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[14], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_mov_b32(s[14], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+8),
s_mov_b32(s[14], 0xDDDDDDDD),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+12),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 0 (s_load_b128 requires 4-aligned dest: s[4], s[8], s[12], ...)
s_mov_b32(s[15], 0),
s_load_b128(s[4:7], s[2:3], s[15], offset=TEST_OFFSET),
s_waitcnt(0),
# Compare with immediate offset
s_load_b128(s[8:11], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB)
self.assertEqual(st.sgpr[6], 0xCCCCCCCC)
self.assertEqual(st.sgpr[7], 0xDDDDDDDD)
self.assertEqual(st.sgpr[4], st.sgpr[8])
self.assertEqual(st.sgpr[5], st.sgpr[9])
class TestSLoadLarge(unittest.TestCase):
"""Tests for large s_load operations (s_load_b256, s_load_b512)."""
def test_s_load_b256_basic(self):
"""s_load_b256 loads 8 consecutive dwords."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 8 test values
s_mov_b32(s[20], 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[20], 0x22222222),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_mov_b32(s[20], 0x33333333),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+8),
s_mov_b32(s[20], 0x44444444),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+12),
s_mov_b32(s[20], 0x55555555),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+16),
s_mov_b32(s[20], 0x66666666),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+20),
s_mov_b32(s[20], 0x77777777),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+24),
s_mov_b32(s[20], 0x88888888),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+28),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load all 8 dwords with s_load_b256
s_load_b256(s[4:11], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111)
self.assertEqual(st.sgpr[5], 0x22222222)
self.assertEqual(st.sgpr[6], 0x33333333)
self.assertEqual(st.sgpr[7], 0x44444444)
self.assertEqual(st.sgpr[8], 0x55555555)
self.assertEqual(st.sgpr[9], 0x66666666)
self.assertEqual(st.sgpr[10], 0x77777777)
self.assertEqual(st.sgpr[11], 0x88888888)
def test_s_load_b512_basic(self):
"""s_load_b512 loads 16 consecutive dwords."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 16 test values (use a pattern: 0x10, 0x20, ..., 0x100)
*[instr for i in range(16) for instr in [
s_mov_b32(s[20], (i + 1) * 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET + i * 4),
]],
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load all 16 dwords with s_load_b512
s_load_b512(s[64:79], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(lgkmcnt=0),
# Copy results to lower regs for verification (since st.sgpr only has 16 regs in test)
s_mov_b32(s[4], s[64]),
s_mov_b32(s[5], s[65]),
s_mov_b32(s[6], s[78]),
s_mov_b32(s[7], s[79]),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111, "first dword")
self.assertEqual(st.sgpr[5], 0x22222222, "second dword")
self.assertEqual(st.sgpr[6], 0xFFFFFFFF & (15 * 0x11111111), "15th dword")
self.assertEqual(st.sgpr[7], 0xFFFFFFFF & (16 * 0x11111111), "16th dword")
def test_s_load_b256_with_register_offset(self):
"""s_load_b256 with register offset should add reg offset to address."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store pattern at TEST_OFFSET+8: skip first 2 dwords
*[instr for i in range(8) for instr in [
s_mov_b32(s[20], (i + 1) * 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET + 8 + i * 4),
]],
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 8
s_mov_b32(s[20], 8),
s_load_b256(s[4:11], s[2:3], s[20], offset=TEST_OFFSET),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111, "first dword at offset+8")
self.assertEqual(st.sgpr[5], 0x22222222, "second dword at offset+8")
self.assertEqual(st.sgpr[11], 0x88888888, "last dword at offset+8")
class TestSLoadOffset(unittest.TestCase):
"""Tests for s_load with different immediate offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load wrong data.
"""
def test_s_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load from offset 100 -> should get 0xAAAAAAAA
s_load_b32(s[4], s[2:3], NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
s_load_b32(s[5], s[2:3], NULL, offset=200),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.sgpr[4]:08x}")
self.assertEqual(st.sgpr[5], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.sgpr[5]:08x}")
def test_s_load_negative_offset(self):
"""Test negative offset (21-bit signed).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from base+200 -> should get 0xAAAA.
Load with offset -100 from base+300 -> should get 0xBBBB."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[8], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[8]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[8], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[8]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# base+200, load with offset -100 -> should get value at 100
s_add_u32(s[6], s[2], 200),
s_addc_u32(s[7], s[3], 0),
s_load_b32(s[4], s[6:7], NULL, offset=-100),
# base+300, load with offset -100 -> should get value at 200
s_add_u32(s[6], s[2], 300),
s_addc_u32(s[7], s[3], 0),
s_load_b32(s[5], s[6:7], NULL, offset=-100),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[6], 0),
s_mov_b32(s[7], 0),
s_mov_b32(s[8], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.sgpr[4]:08x}")
self.assertEqual(st.sgpr[5], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.sgpr[5]:08x}")
if __name__ == '__main__':
unittest.main()
+313
View File
@@ -62,6 +62,7 @@ class TestBasicScalar(unittest.TestCase):
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 0x80000000)
@skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
def test_s_fmamk_f32(self):
"""S_FMAMK_F32: D = S0 * literal + S1."""
# 2.0 * 3.0 + 1.0 = 7.0
@@ -73,6 +74,7 @@ class TestBasicScalar(unittest.TestCase):
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], f2i(7.0))
@skip_unless_gfx(11, 5, "SALU FP ops require gfx1150+")
def test_s_fmamk_f32_negative(self):
"""S_FMAMK_F32 with negative values."""
# -2.0 * 4.0 + 10.0 = 2.0
@@ -85,6 +87,50 @@ class TestBasicScalar(unittest.TestCase):
self.assertEqual(st.sgpr[2], f2i(2.0))
class TestPack(unittest.TestCase):
"""Tests for S_PACK instructions."""
def test_s_pack_ll_b32_b16(self):
"""S_PACK_LL_B32_B16 packs low 16 bits of two sources into one 32-bit result."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_ll_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xBBBBAAAA)
def test_s_pack_lh_b32_b16(self):
"""S_PACK_LH_B32_B16: D0 = { S1[31:16], S0[15:0] }."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_lh_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xDEADAAAA)
def test_s_pack_hh_b32_b16(self):
"""S_PACK_HH_B32_B16: D0 = { S1[31:16], S0[31:16] }."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_hh_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xDEADDEAD)
def test_s_pack_hl_b32_b16(self):
"""S_PACK_HL_B32_B16: D0 = { S1[15:0], S0[31:16] }."""
instructions = [
s_mov_b32(s[0], 0xDEADAAAA),
s_mov_b32(s[1], 0xDEADBBBB),
s_pack_hl_b32_b16(s[2], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0xBBBBDEAD)
class TestQuadmaskWqm(unittest.TestCase):
"""Tests for S_QUADMASK_B32 and S_WQM_B32."""
@@ -619,5 +665,272 @@ class Test64BitCompare(unittest.TestCase):
self.assertEqual(st.sgpr[4], 1)
class TestSOPPNop(unittest.TestCase):
"""Tests for S_NOP and other SOPP instructions with expression-based for loops.
S_NOP's pcode uses 'for i in 0U : SIMM16.u16[3 : 0].u32 do' which requires
the parser to handle non-constant loop bounds.
"""
def test_s_nop_basic(self):
"""S_NOP executes without side effects."""
# S_NOP with immediate 0 should just do nothing
instructions = [
s_mov_b32(s[0], 42),
s_nop(0), # nop with simm16=0
s_mov_b32(s[1], 100),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 42)
self.assertEqual(st.sgpr[1], 100)
def test_s_nop_with_count(self):
"""S_NOP with count parameter executes multiple nops."""
# S_NOP with immediate 3 should execute 4 nops (0:3 inclusive)
instructions = [
s_mov_b32(s[0], 1),
s_nop(3), # nop with simm16=3 -> 4 iterations
s_add_u32(s[0], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 2)
class TestNullRegister(unittest.TestCase):
"""Tests for NULL register (124) behavior - writes should be discarded, reads return 0."""
def test_s_mov_b32_from_null(self):
"""S_MOV_B32 from NULL should read as 0."""
instructions = [
s_mov_b32(s[0], 0xDEADBEEF), # Set s[0] to sentinel
s_mov_b32(s[0], NULL), # Read from NULL - should be 0
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 0)
def test_s_add_u32_with_null_src(self):
"""S_ADD_U32 with NULL as source should use 0."""
instructions = [
s_mov_b32(s[0], 100),
s_add_u32(s[1], s[0], NULL), # 100 + 0 = 100
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 100)
def test_s_mov_b32_to_null(self):
"""S_MOV_B32 to NULL (sdst=124) should discard the write."""
instructions = [
s_mov_b32(s[0], 0xDEADBEEF), # Set s[0] to sentinel
s_mov_b32(NULL, 42), # Write to NULL - should be discarded
# s[0] should still be 0xDEADBEEF since NULL write doesn't affect it
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 0xDEADBEEF)
def test_s_add_u32_to_null(self):
"""S_ADD_U32 with sdst=NULL should discard result but still set SCC."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 1),
s_add_u32(NULL, s[0], s[1]), # overflow, write to NULL
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
# SCC should still be set from overflow even though result was discarded
self.assertEqual(st.sgpr[2], 1)
self.assertEqual(st.scc, 1)
def test_s_and_b32_to_null(self):
"""S_AND_B32 with sdst=NULL should discard result but still set SCC."""
instructions = [
s_mov_b32(s[0], 0xFF00FF00),
s_mov_b32(s[1], 0x0F0F0F0F),
s_and_b32(NULL, s[0], s[1]), # result=0x0F000F00, non-zero so SCC=1
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 1) # SCC=1 because result was non-zero
self.assertEqual(st.scc, 1)
def test_s_or_b32_to_null_zero_result(self):
"""S_OR_B32 with sdst=NULL and zero result should set SCC=0."""
instructions = [
s_mov_b32(s[0], 0),
s_mov_b32(s[1], 0),
s_or_b32(NULL, s[0], s[1]), # result=0, so SCC=0
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0) # SCC=0 because result was zero
self.assertEqual(st.scc, 0)
class Test64BitSOP1InlineConstants(unittest.TestCase):
"""Tests for 64-bit SOP1 instructions with inline constants.
Regression tests for bug where rsrc_dyn didn't properly handle 64-bit
inline constants, incorrectly duplicating lo bits to hi instead of
zero/sign-extending.
"""
def test_s_mov_b64_inline_0(self):
"""S_MOV_B64 with inline constant 0."""
instructions = [
s_mov_b64(s[0:1], 0),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_16(self):
"""S_MOV_B64 with inline constant 16 should set lo=16, hi=0."""
instructions = [
s_mov_b64(s[0:1], 16),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_64(self):
"""S_MOV_B64 with inline constant 64 (max positive)."""
instructions = [
s_mov_b64(s[0:1], 64),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 64)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_neg1(self):
"""S_MOV_B64 with inline constant -1 should sign-extend."""
instructions = [
s_mov_b64(s[0:1], -1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
def test_s_mov_b64_inline_neg16(self):
"""S_MOV_B64 with inline constant -16 should sign-extend."""
instructions = [
s_mov_b64(s[0:1], -16),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFF0)
self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
def test_s_mov_b64_float_const_1_0(self):
"""S_MOV_B64 with float inline constant 1.0 - casts F32 to F64."""
instructions = [
s_mov_b64(s[0:1], 1.0), # inline constant 242 (1.0f)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
# Hardware casts F32 to F64: 1.0f64 = 0x3FF0000000000000
self.assertEqual(st.vgpr[0][0], 0x00000000) # lo
self.assertEqual(st.vgpr[0][1], 0x3FF00000) # hi
def test_s_or_b64_inline_constant(self):
"""S_OR_B64 with 64-bit inline constant."""
instructions = [
s_mov_b64(s[0:1], 0),
s_or_b64(s[2:3], s[0:1], 16),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_and_b64_inline_constant(self):
"""S_AND_B64 with 64-bit inline constant."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 0xFFFFFFFF),
s_and_b64(s[2:3], s[0:1], 16),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
class Test64BitSOPLiterals(unittest.TestCase):
"""Tests for 64-bit SOP instructions with 32-bit literals.
Tests the behavior when a 64-bit SOP instruction uses a 32-bit literal
(offset 255 in instruction encoding). The literal is zero-extended to 64 bits.
"""
def test_s_mov_b64_literal(self):
"""S_MOV_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b64(s[0:1], 0x12345678), # literal > 64, uses literal encoding
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_or_b64_literal(self):
"""S_OR_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b64(s[0:1], 0),
s_or_b64(s[2:3], s[0:1], 0x12345678), # literal
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_and_b64_literal(self):
"""S_AND_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 0xFFFFFFFF),
s_and_b64(s[2:3], s[0:1], 0x12345678), # literal
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_literal_negative(self):
"""S_MOV_B64 with 0xFFFFFFFF literal - zero-extended (not sign-extended)."""
instructions = [
s_mov_b64(s[0:1], 0xFFFFFFFF), # -1 as 32-bit, but zero-extended to 64-bit
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][1], 0) # zero-extended, not sign-extended
def test_s_mov_b64_literal_high_bit(self):
"""S_MOV_B64 with 0x80000000 literal - zero-extended (not sign-extended)."""
instructions = [
s_mov_b64(s[0:1], 0x80000000), # high bit set, but zero-extended
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x80000000)
self.assertEqual(st.vgpr[0][1], 0) # zero-extended, not sign-extended
if __name__ == '__main__':
unittest.main()
+80 -12
View File
@@ -255,7 +255,6 @@ class TestF16Conversions(unittest.TestCase):
def test_v_cvt_f16_f32_small(self):
"""V_CVT_F16_F32 converts small f32 value."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
instructions = [
v_mov_b32_e32(v[0], 0.5),
v_cvt_f16_f32_e32(v[1], v[0]),
@@ -293,7 +292,6 @@ class TestF16Conversions(unittest.TestCase):
def test_v_cvt_f16_f32_reads_full_32bit_source(self):
"""V_CVT_F16_F32 must read full 32-bit f32 source."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x3fc00000), # f32 1.5
v_mov_b32_e32(v[0], s[0]),
@@ -302,7 +300,7 @@ class TestF16Conversions(unittest.TestCase):
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo_bits = result & 0xffff
self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({_f16(lo_bits)})")
self.assertEqual(lo_bits, 0x3e00, f"Expected f16(1.5)=0x3e00, got 0x{lo_bits:04x} ({f16(lo_bits)})")
def test_v_cvt_i16_f16_zero(self):
"""V_CVT_I16_F16 converts f16 zero to i16 zero."""
@@ -696,7 +694,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_abs_negative(self):
"""V_CVT_F32_F16 with |abs| on negative value."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg1 = f32_to_f16(-1.0) # 0xbc00
instructions = [
s_mov_b32(s[0], f16_neg1),
@@ -709,7 +706,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_abs_positive(self):
"""V_CVT_F32_F16 with |abs| on positive value (should stay positive)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0) # 0x4000
instructions = [
s_mov_b32(s[0], f16_2),
@@ -722,7 +718,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_neg_positive(self):
"""V_CVT_F32_F16 with neg on positive value."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0) # 0x4000
instructions = [
s_mov_b32(s[0], f16_2),
@@ -735,7 +730,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f32_f16_neg_negative(self):
"""V_CVT_F32_F16 with neg on negative value (double negative)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg2 = f32_to_f16(-2.0) # 0xc000
instructions = [
s_mov_b32(s[0], f16_neg2),
@@ -748,7 +742,6 @@ class TestCvtF16Modifiers(unittest.TestCase):
def test_v_cvt_f16_f32_then_pack_for_wmma(self):
"""CVT F32->F16 followed by pack (common WMMA pattern)."""
from extra.assembly.amd.test.hw.helpers import _f16
f32_val = 3.5
instructions = [
s_mov_b32(s[0], f2i(f32_val)),
@@ -757,8 +750,8 @@ class TestCvtF16Modifiers(unittest.TestCase):
v_pack_b32_f16(v[2], v[1], v[1]), # Pack same value
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][2] & 0xffff)
hi = _f16((st.vgpr[0][2] >> 16) & 0xffff)
lo = f16(st.vgpr[0][2] & 0xffff)
hi = f16((st.vgpr[0][2] >> 16) & 0xffff)
self.assertAlmostEqual(lo, f32_val, places=1)
self.assertAlmostEqual(hi, f32_val, places=1)
@@ -804,7 +797,6 @@ class TestConversionRounding(unittest.TestCase):
def test_f16_to_f32_precision(self):
"""F16 to F32 conversion precision."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_val = f32_to_f16(1.5)
instructions = [
s_mov_b32(s[0], f16_val),
@@ -816,7 +808,6 @@ class TestConversionRounding(unittest.TestCase):
def test_f16_denormal_to_f32(self):
"""F16 denormal converts to small positive f32."""
from extra.assembly.amd.test.hw.helpers import _f16
f16_denorm = 0x0001 # Smallest positive f16 denormal
instructions = [
v_mov_b32_e32(v[0], f16_denorm),
@@ -1512,5 +1503,82 @@ class TestReciprocalF16(unittest.TestCase):
self.assertAlmostEqual(result, 0.25, places=2, msg="1/4.0 should be 0.25")
class TestCvtNormF16(unittest.TestCase):
"""Tests for V_CVT_NORM_I16_F16 and V_CVT_NORM_U16_F16."""
def test_cvt_norm_i16_f16_positive(self):
"""V_CVT_NORM_I16_F16: f16 1.0 -> i16 max (32767)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(1.0)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_i16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 32767)
def test_cvt_norm_i16_f16_negative(self):
"""V_CVT_NORM_I16_F16: f16 -1.0 -> i16 -32767 (0x8001)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(-1.0)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_i16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x8001) # -32767, hardware uses symmetric range
def test_cvt_norm_i16_f16_zero(self):
"""V_CVT_NORM_I16_F16: f16 0.0 -> i16 0."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_cvt_norm_i16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0)
def test_cvt_norm_u16_f16_one(self):
"""V_CVT_NORM_U16_F16: f16 1.0 -> u16 max (65535)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(1.0)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_u16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 65535)
def test_cvt_norm_u16_f16_half(self):
"""V_CVT_NORM_U16_F16: f16 0.5 -> u16 ~32768."""
instructions = [
s_mov_b32(s[0], f32_to_f16(0.5)),
v_mov_b32_e32(v[0], s[0]),
v_cvt_norm_u16_f16_e32(v[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertAlmostEqual(result, 32768, delta=1)
class TestPermlane64(unittest.TestCase):
"""Tests for V_PERMLANE64_B32 instruction (wave64 cross-half swap)."""
def test_v_permlane64_b32_is_nop_in_wave32(self):
"""V_PERMLANE64_B32 is a NOP in wave32 mode.
Per AMD pcode: "if WAVE32 then s_nop(...) else ... endif"
The emulator runs in wave32 mode, so this instruction should not modify registers.
"""
instructions = [
v_mov_b32_e32(v[0], 0xCAFEBABE), # source
v_mov_b32_e32(v[1], 0x12345678), # dest (should be preserved)
v_permlane64_b32_e32(v[1], v[0]), # NOP in wave32
]
st = run_program(instructions, n_lanes=1)
# Dest register should be unchanged (NOP behavior in wave32)
self.assertEqual(st.vgpr[0][1], 0x12345678)
if __name__ == '__main__':
unittest.main()
+118
View File
@@ -237,6 +237,32 @@ class TestF16Ops(unittest.TestCase):
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
def test_v_max_f16_basic(self):
"""V_MAX_F16 returns the maximum of two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_max_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected 0x4000 (f16 2.0), got 0x{result:04x}")
def test_v_min_f16_basic(self):
"""V_MIN_F16 returns the minimum of two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_min_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected 0x3c00 (f16 1.0), got 0x{result:04x}")
def test_v_fmaak_f16_basic(self):
"""V_FMAAK_F16: d = a * b + K."""
instructions = [
@@ -810,6 +836,81 @@ class TestCarryOps(unittest.TestCase):
self.assertEqual(st.vgpr[0][2], 0) # Overflowed to 0
self.assertEqual(st.vcc, 1) # Carry out
def test_v_add_co_ci_u32_clears_carry(self):
"""V_ADD_CO_CI_U32: VCC must be updated even when no carry is generated.
This tests the case where VCC=1 going in (carry-in consumed) but the addition
does not overflow, so VCC must be cleared to 0.
Regression test for: VCC not being written by v_add_co_ci_u32_e32.
"""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 3) # 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out - VCC must be cleared
def test_v_add_co_ci_u32_multilane_clears_vcc(self):
"""V_ADD_CO_CI_U32 with multiple lanes: VCC bits must be updated per-lane.
When VCC has multiple bits set (one per active lane), and the addition doesn't
overflow for any lane, all VCC bits must be cleared.
Regression test for: VCC not being written by v_add_co_ci_u32_e32 in multi-lane case.
"""
instructions = [
s_mov_b32(VCC_LO, 0b11), # VCC = 0b11 (lanes 0,1 have carry-in)
v_mov_b32_e32(v[0], 1), # S0 = 1 for all lanes
v_mov_b32_e32(v[1], 1), # S1 = 1 for all lanes
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vgpr[0][2], 3) # lane 0: 1 + 1 + 1 = 3
self.assertEqual(st.vgpr[1][2], 3) # lane 1: 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out for any lane - all VCC bits must be cleared
def test_v_add_co_ci_u32_preserves_inactive_vcc_bits(self):
"""V_ADD_CO_CI_U32: VCC carry-out overwrites entire VCC register.
VOP2 carry instructions write ALL VCC bits based on carry-out, clearing
bits for lanes that don't overflow regardless of EXEC mask.
Note: This differs from VOPC which only writes active lane bits.
"""
instructions = [
s_mov_b32(VCC_LO, 0x00010000), # VCC bit 16 set
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 0 = 2 (no carry)
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vgpr[0][2], 2) # lane 0: 1 + 1 + 0 = 2
# VCC should be completely cleared (all lanes have no carry-out)
self.assertEqual(st.vcc, 0)
def test_v_add_co_ci_u32_all_lanes_same_result(self):
"""V_ADD_CO_CI_U32: all active lanes should produce the same result.
When the same constant inputs are used across all lanes, each lane should
compute the same result and write to its own VGPR slot.
Regression test for: VGPR writes not happening for all lanes.
"""
instructions = [
s_mov_b32(VCC_LO, 0), # No carry-in
v_mov_b32_e32(v[0], 3), # inline constant 3
v_mov_b32_e32(v[1], 5), # value 5
v_add_co_ci_u32_e32(v[1], 3, v[1]), # v[1] = 3 + v[1] + 0 = 3 + 5 = 8
]
st = run_program(instructions, n_lanes=4)
# All 4 lanes should have v[1] = 8
for lane in range(4):
self.assertEqual(st.vgpr[lane][1], 8, f"lane {lane} should have v[1]=8")
def test_v_sub_co_ci_u32_no_borrow(self):
"""V_SUB_CO_CI_U32: D0 = S0 - S1 - VCC_IN, when VCC_IN=0."""
instructions = [
@@ -860,6 +961,23 @@ class TestCarryOps(unittest.TestCase):
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.sgpr[10], 0) # No carry out
def test_v_add_co_ci_u32_vop3sd_null_sdst(self):
"""VOP3SD V_ADD_CO_CI_U32 with sdst=NULL: carry output is discarded.
When sdst=NULL (register 124), the carry-out should NOT be written anywhere.
We verify this by checking that VCC (which we set to a sentinel value) is unchanged.
"""
instructions = [
s_mov_b32(VCC_LO, 0xDEADBEEF), # Sentinel value in VCC
s_mov_b32(s[6], 0), # carry-in = 0
# VOP3SD with NULL sdst: carry-out should be discarded
# Uses 0xFFFFFFFF + 1 + 0 = 0 with carry-out=1, but carry should not be written
v_add_co_ci_u32(v[0], NULL, 0xFFFFFFFF, 1, s[6]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0) # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
self.assertEqual(st.vcc, 0xDEADBEEF) # VCC unchanged - carry was discarded
if __name__ == '__main__':
unittest.main()
+841 -14
View File
@@ -857,7 +857,6 @@ class TestF16Modifiers(unittest.TestCase):
def test_v_fma_f16_inline_const_1_0(self):
"""V_FMA_F16: a*b + 1.0 should use f16 inline constant."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
f16_a = f32_to_f16(0.325928) # ~0x3537
f16_b = f32_to_f16(-0.486572) # ~0xb7c9
instructions = [
@@ -868,13 +867,12 @@ class TestF16Modifiers(unittest.TestCase):
v_fma_f16(v[4], v[4], v[6], 1.0), # 1.0 is inline constant
]
st = run_program(instructions, n_lanes=1)
result = _f16(st.vgpr[0][4] & 0xffff)
result = f16(st.vgpr[0][4] & 0xffff)
expected = 0.325928 * (-0.486572) + 1.0
self.assertAlmostEqual(result, expected, delta=0.01)
def test_v_fma_f16_inline_const_0_5(self):
"""V_FMA_F16: a*b + 0.5 should use f16 inline constant."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
f16_a = f32_to_f16(2.0)
f16_b = f32_to_f16(3.0)
instructions = [
@@ -885,13 +883,12 @@ class TestF16Modifiers(unittest.TestCase):
v_fma_f16(v[2], v[0], v[1], 0.5), # 0.5 is inline constant
]
st = run_program(instructions, n_lanes=1)
result = _f16(st.vgpr[0][2] & 0xffff)
result = f16(st.vgpr[0][2] & 0xffff)
expected = 2.0 * 3.0 + 0.5
self.assertAlmostEqual(result, expected, delta=0.01)
def test_v_fma_f16_inline_const_neg_1_0(self):
"""V_FMA_F16: a*b + (-1.0) should use f16 inline constant."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
f16_a = f32_to_f16(2.0)
f16_b = f32_to_f16(3.0)
instructions = [
@@ -902,13 +899,12 @@ class TestF16Modifiers(unittest.TestCase):
v_fma_f16(v[2], v[0], v[1], -1.0), # -1.0 is inline constant
]
st = run_program(instructions, n_lanes=1)
result = _f16(st.vgpr[0][2] & 0xffff)
result = f16(st.vgpr[0][2] & 0xffff)
expected = 2.0 * 3.0 + (-1.0)
self.assertAlmostEqual(result, expected, delta=0.01)
def test_v_add_f16_abs_both(self):
"""V_ADD_F16 with abs on both operands."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
f16_neg2 = f32_to_f16(-2.0)
f16_neg3 = f32_to_f16(-3.0)
instructions = [
@@ -919,12 +915,11 @@ class TestF16Modifiers(unittest.TestCase):
v_add_f16_e64(v[2], abs(v[0]), abs(v[1])), # |-2| + |-3| = 5
]
st = run_program(instructions, n_lanes=1)
result = _f16(st.vgpr[0][2] & 0xffff)
result = f16(st.vgpr[0][2] & 0xffff)
self.assertAlmostEqual(result, 5.0, delta=0.01)
def test_v_mul_f16_neg_abs(self):
"""V_MUL_F16 with neg on one operand and abs on another."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16, _f16
f16_2 = f32_to_f16(2.0)
f16_neg3 = f32_to_f16(-3.0)
instructions = [
@@ -935,7 +930,7 @@ class TestF16Modifiers(unittest.TestCase):
v_mul_f16_e64(v[2], -v[0], abs(v[1])), # -(2) * |-3| = -6
]
st = run_program(instructions, n_lanes=1)
result = _f16(st.vgpr[0][2] & 0xffff)
result = f16(st.vgpr[0][2] & 0xffff)
self.assertAlmostEqual(result, -6.0, delta=0.01)
def test_v_fmac_f16_hi_dest(self):
@@ -943,7 +938,6 @@ class TestF16Modifiers(unittest.TestCase):
This tests the case from AMD_LLVM sin(0) where V_FMAC_F16 writes to v0.h.
"""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x38003c00), # v0 = {hi=0.5, lo=1.0}
v_mov_b32_e32(v[0], s[0]),
@@ -954,8 +948,8 @@ class TestF16Modifiers(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
v0 = st.vgpr[0][0]
result_hi = _f16((v0 >> 16) & 0xffff)
result_lo = _f16(v0 & 0xffff)
result_hi = f16((v0 >> 16) & 0xffff)
result_lo = f16(v0 & 0xffff)
self.assertAlmostEqual(result_hi, 0.5, delta=0.01, msg=f"Expected hi=0.5, got {result_hi}")
self.assertAlmostEqual(result_lo, 1.0, delta=0.01, msg=f"Expected lo=1.0, got {result_lo}")
@@ -1077,7 +1071,7 @@ class TestF64Ops(unittest.TestCase):
s_mov_b32(s[3], one_f64 >> 32),
v_mov_b32_e32(v[2], s[2]),
v_mov_b32_e32(v[3], s[3]),
VOP3SD(VOP3SDOp.V_DIV_SCALE_F64, vdst=v[4:5], sdst=s[10:11], src0=v[0:1], src1=v[0:1], src2=v[2:3]),
VOP3SD(VOP3SDOp.V_DIV_SCALE_F64, vdst=v[4:5], sdst=s[10], src0=v[0:1], src1=v[0:1], src2=v[2:3]),
]
st = run_program(instructions, n_lanes=1)
result = i642f(st.vgpr[0][4] | (st.vgpr[0][5] << 32))
@@ -1359,6 +1353,43 @@ class TestF64ToI64Conversion(unittest.TestCase):
self.assertEqual(result, 5000000000)
class TestB64VOPLiteral(unittest.TestCase):
"""Tests for B64 VOP operations with literal encoding.
B64 operations (like V_LSHLREV_B64) should zero-extend the literal to 64 bits,
NOT put it in the high 32 bits like F64 operations do.
"""
def test_v_lshlrev_b64_literal_shift_amount(self):
"""V_LSHLREV_B64 with literal shift amount (src0 is 32-bit)."""
# Shift 1 left by 100 (0x64) - uses literal encoding for src0
# Shift amount is 100 & 63 = 36, so 1 << 36 = 0x1000000000
instructions = [
s_mov_b32(s[0], 1),
s_mov_b32(s[1], 0),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_lshlrev_b64(v[2:3], 100, v[0:1]), # 100 > 64, uses literal encoding
]
st = run_program(instructions, n_lanes=1)
# lo = 0x00000000, hi = 0x00000010 = 1 << (36-32)
self.assertEqual(st.vgpr[0][2], 0x00000000)
self.assertEqual(st.vgpr[0][3], 0x00000010)
def test_v_lshlrev_b64_literal_value(self):
"""V_LSHLREV_B64 with literal as the 64-bit value being shifted (src1).
B64 literals are zero-extended (not shifted to high bits like F64).
0xDEADBEEF << 4 = 0xDEADBEEF0 = lo=0xEADBEEF0, hi=0x0000000D
"""
instructions = [
v_lshlrev_b64(v[0:1], 4, 0xDEADBEEF), # shift literal left by 4
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xEADBEEF0) # lo
self.assertEqual(st.vgpr[0][1], 0x0000000D) # hi
class TestWMMAMore(unittest.TestCase):
"""More WMMA tests."""
@@ -2811,5 +2842,801 @@ class TestMin3Max3Unsigned(unittest.TestCase):
self.assertEqual(st.vgpr[0][1] & 0xFFFF, 0)
class TestVOP3Clamp(unittest.TestCase):
"""Tests for VOP3 clamp modifier (clmp=1).
The clamp modifier restricts float outputs to [0.0, 1.0] range.
This is used by operations like clip(0, 1) which AMD LLVM compiles to
v_max_f32_e64 with clmp=1.
Regression test for: clip(0, 1) bug where emulator ignored clmp field.
"""
def test_v_max_f32_e64_clamp_positive(self):
"""V_MAX_F32_E64 with clamp: value > 1.0 should be clamped to 1.0."""
instructions = [
v_mov_b32_e32(v[0], 2.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 1.0, places=5)
def test_v_max_f32_e64_clamp_negative(self):
"""V_MAX_F32_E64 with clamp: value < 0.0 should be clamped to 0.0."""
instructions = [
v_mov_b32_e32(v[0], -1.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.0, places=5)
def test_v_max_f32_e64_clamp_in_range(self):
"""V_MAX_F32_E64 with clamp: value in [0,1] should pass through."""
instructions = [
v_mov_b32_e32(v[0], 0.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.5, places=5)
def test_v_max_f32_e64_no_clamp(self):
"""V_MAX_F32_E64 without clamp: value > 1.0 should pass through."""
instructions = [
v_mov_b32_e32(v[0], 2.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=0),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.5, places=5)
def test_v_min_f32_e64_clamp_negative(self):
"""V_MIN_F32_E64 with clamp: value < 0.0 should be clamped to 0.0."""
instructions = [
v_mov_b32_e32(v[0], -2.0),
VOP3(VOP3Op.V_MIN_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.0, places=5)
def test_v_add_f32_e64_clamp(self):
"""V_ADD_F32_E64 with clamp: 0.7 + 0.8 = 1.5 -> 1.0."""
instructions = [
v_mov_b32_e32(v[0], 0.7),
v_mov_b32_e32(v[1], 0.8),
VOP3(VOP3Op.V_ADD_F32_E64, vdst=v[2], src0=v[0], src1=v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 1.0, places=5)
def test_v_mul_f32_e64_clamp_underflow(self):
"""V_MUL_F32_E64 with clamp: 0.5 * -2.0 = -1.0 -> 0.0."""
instructions = [
v_mov_b32_e32(v[0], 0.5),
v_mov_b32_e32(v[1], -2.0),
VOP3(VOP3Op.V_MUL_F32_E64, vdst=v[2], src0=v[0], src1=v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 0.0, places=5)
def test_v_fma_f32_clamp(self):
"""V_FMA_F32 with clamp: 2*2+1 = 5 -> 1.0."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 2.0),
v_mov_b32_e32(v[2], 1.0),
VOP3(VOP3Op.V_FMA_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][3]), 1.0, places=5)
def test_v_max_f32_e64_clamp_multilane(self):
"""V_MAX_F32_E64 with clamp: test multiple lanes with different values."""
# lane 0: -0.5 -> 0.0
# lane 1: 0.5 -> 0.5
# lane 2: 1.5 -> 1.0
# lane 3: 2.5 -> 1.0
instructions = [
# Setup different values per lane using lane_id
s_mov_b32(s[0], f2i(0.5)),
v_cvt_f32_i32_e32(v[0], v[255]), # Convert lane_id to float
v_mov_b32_e32(v[2], s[0]), # v2 = 0.5
v_sub_f32_e32(v[0], v[0], v[2]), # Subtract 0.5: lane0=-0.5, lane1=0.5, lane2=1.5, lane3=2.5
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=4)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.0, places=5, msg="lane 0: -0.5 should clamp to 0.0")
self.assertAlmostEqual(i2f(st.vgpr[1][1]), 0.5, places=5, msg="lane 1: 0.5 should pass through")
self.assertAlmostEqual(i2f(st.vgpr[2][1]), 1.0, places=5, msg="lane 2: 1.5 should clamp to 1.0")
self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")
class TestVOP3ClampUint32(unittest.TestCase):
"""Tests for VOP3 clamp modifier on unsigned 32-bit integer operations."""
def test_v_sub_nc_u32_e64_clamp_underflow(self):
"""V_SUB_NC_U32_E64 with clamp: 0 - 1 should saturate to 0."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 1),
v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
def test_v_sub_nc_u32_e64_clamp_no_underflow(self):
"""V_SUB_NC_U32_E64 with clamp: 100 - 50 = 50 (no saturation needed)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 50),
v_sub_nc_u32_e64(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 50, f"expected 50, got {st.vgpr[0][2]}")
def test_v_add_nc_u32_e64_clamp_overflow(self):
"""V_ADD_NC_U32_E64 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFFFFFF),
v_mov_b32_e32(v[1], 1),
v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
def test_v_add_nc_u32_e64_clamp_no_overflow(self):
"""V_ADD_NC_U32_E64 with clamp: 100 + 50 = 150 (no saturation needed)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 50),
v_add_nc_u32_e64(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 150, f"expected 150, got {st.vgpr[0][2]}")
class TestVOP3ClampUint16(unittest.TestCase):
"""Tests for VOP3 clamp modifier on unsigned 16-bit integer operations."""
def test_v_sub_nc_u16_clamp_underflow(self):
"""V_SUB_NC_U16 with clamp: 0 - 1 should saturate to 0."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 1),
v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0, f"expected 0, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
def test_v_sub_nc_u16_clamp_no_underflow(self):
"""V_SUB_NC_U16 with clamp: 100 - 50 = 50 (no saturation needed)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 50),
v_sub_nc_u16(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2] & 0xFFFF, 50, f"expected 50, got {st.vgpr[0][2] & 0xFFFF}")
def test_v_add_nc_u16_clamp_overflow(self):
"""V_ADD_NC_U16 with clamp: 0xFFFF + 1 should saturate to 0xFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFF),
v_mov_b32_e32(v[1], 1),
v_add_nc_u16(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][2] & 0xFFFF:04x}")
def test_v_add_nc_u16_clamp_no_overflow(self):
"""V_ADD_NC_U16 with clamp: 100 + 50 = 150 (no saturation needed)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 50),
v_add_nc_u16(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2] & 0xFFFF, 150, f"expected 150, got {st.vgpr[0][2] & 0xFFFF}")
class TestVOP3ClampInt32(unittest.TestCase):
"""Tests for VOP3 clamp modifier on signed 32-bit integer operations."""
def test_v_add_nc_i32_clamp_overflow(self):
"""V_ADD_NC_I32 with clamp: INT_MAX + 1 should saturate to INT_MAX."""
instructions = [
v_mov_b32_e32(v[0], 0x7FFFFFFF), # S0 = INT_MAX
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_nc_i32(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
def test_v_add_nc_i32_clamp_underflow(self):
"""V_ADD_NC_I32 with clamp: INT_MIN + (-1) should saturate to INT_MIN."""
instructions = [
v_mov_b32_e32(v[0], 0x80000000), # S0 = INT_MIN
v_mov_b32_e32(v[1], 0xFFFFFFFF), # S1 = -1
v_add_nc_i32(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
def test_v_sub_nc_i32_clamp_underflow(self):
"""V_SUB_NC_I32 with clamp: INT_MIN - 1 should saturate to INT_MIN."""
instructions = [
v_mov_b32_e32(v[0], 0x80000000), # S0 = INT_MIN
v_mov_b32_e32(v[1], 1), # S1 = 1
v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x80000000, f"expected 0x80000000, got 0x{st.vgpr[0][2]:08x}")
def test_v_sub_nc_i32_clamp_overflow(self):
"""V_SUB_NC_I32 with clamp: INT_MAX - (-1) should saturate to INT_MAX."""
instructions = [
v_mov_b32_e32(v[0], 0x7FFFFFFF), # S0 = INT_MAX
v_mov_b32_e32(v[1], 0xFFFFFFFF), # S1 = -1
v_sub_nc_i32(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF, f"expected 0x7FFFFFFF, got 0x{st.vgpr[0][2]:08x}")
def test_v_add_nc_i32_no_saturation_positive(self):
"""V_ADD_NC_I32 with clamp: 100 + 200 = 300 (no saturation needed)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 200),
v_add_nc_i32(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
def test_v_add_nc_i32_no_saturation_negative(self):
"""V_ADD_NC_I32 with clamp: -100 + -200 = -300 (no saturation needed)."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFFFF9C), # -100
v_mov_b32_e32(v[1], 0xFFFFFF38), # -200
v_add_nc_i32(v[2], v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
expected = 0xFFFFFED4 # -300
self.assertEqual(st.vgpr[0][2], expected, f"expected 0x{expected:08x}, got 0x{st.vgpr[0][2]:08x}")
class TestVOP3ClampCarry(unittest.TestCase):
"""Tests for VOP3 clamp modifier on carry operations (VOP3SD)."""
def test_v_add_co_u32_clamp_overflow(self):
"""V_ADD_CO_U32 with clamp: 0xFFFFFFFF + 1 should saturate to 0xFFFFFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFFFFFF),
v_mov_b32_e32(v[1], 1),
v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
def test_v_add_co_u32_clamp_no_overflow(self):
"""V_ADD_CO_U32 with clamp: 100 + 200 = 300 (no saturation)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 200),
v_add_co_u32(v[2], VCC, v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 300, f"expected 300, got {st.vgpr[0][2]}")
def test_v_sub_co_u32_clamp_underflow(self):
"""V_SUB_CO_U32 with clamp: 0 - 1 should saturate to 0."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 1),
v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
def test_v_sub_co_u32_clamp_no_underflow(self):
"""V_SUB_CO_U32 with clamp: 300 - 100 = 200 (no saturation)."""
instructions = [
v_mov_b32_e32(v[0], 300),
v_mov_b32_e32(v[1], 100),
v_sub_co_u32(v[2], VCC, v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 200, f"expected 200, got {st.vgpr[0][2]}")
def test_v_subrev_co_u32_clamp_underflow(self):
"""V_SUBREV_CO_U32 with clamp: 1 - 0 reversed = 0 - 1 should saturate to 0."""
instructions = [
v_mov_b32_e32(v[0], 1), # This becomes the subtrahend
v_mov_b32_e32(v[1], 0), # This becomes the minuend (0 - 1)
v_subrev_co_u32(v[2], VCC, v[0], v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
def test_v_add_co_ci_u32_clamp_overflow(self):
"""V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFF + 1 + 0 should saturate to 0xFFFFFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFFFFFF),
v_mov_b32_e32(v[1], 1),
s_mov_b64(VCC, 0), # No carry in
v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
def test_v_add_co_ci_u32_clamp_overflow_with_carry(self):
"""V_ADD_CO_CI_U32 with clamp: 0xFFFFFFFE + 1 + 1 should saturate to 0xFFFFFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFFFFFE),
v_mov_b32_e32(v[1], 1),
s_mov_b64(VCC, 1), # Carry in = 1
v_add_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFFFFFFFF, f"expected 0xFFFFFFFF, got 0x{st.vgpr[0][2]:08x}")
def test_v_sub_co_ci_u32_clamp_underflow(self):
"""V_SUB_CO_CI_U32 with clamp: 0 - 1 - 0 should saturate to 0."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 1),
s_mov_b64(VCC, 0), # No borrow in
v_sub_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
def test_v_subrev_co_ci_u32_clamp_underflow(self):
"""V_SUBREV_CO_CI_U32 with clamp: reversed 1 - 0 - 0 = 0 - 1 should saturate to 0."""
instructions = [
v_mov_b32_e32(v[0], 1),
v_mov_b32_e32(v[1], 0),
s_mov_b64(VCC, 0),
v_subrev_co_ci_u32(v[2], VCC, v[0], v[1], VCC, clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, f"expected 0, got 0x{st.vgpr[0][2]:08x}")
class TestVOP3ClampMAD(unittest.TestCase):
"""Tests for VOP3 clamp modifier on MAD (multiply-add) operations."""
def test_v_mad_u16_clamp_overflow(self):
"""V_MAD_U16 with clamp: 0xFFFF * 2 + 0 should saturate to 0xFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFF),
v_mov_b32_e32(v[1], 2),
v_mov_b32_e32(v[2], 0),
v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
def test_v_mad_u16_clamp_overflow_with_add(self):
"""V_MAD_U16 with clamp: 0x8000 * 2 + 0x1000 should saturate to 0xFFFF."""
instructions = [
v_mov_b32_e32(v[0], 0x8000), # 32768
v_mov_b32_e32(v[1], 2), # * 2 = 65536
v_mov_b32_e32(v[2], 0x1000), # + 4096 = 69632 > 0xFFFF
v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFF, f"expected 0xFFFF, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
def test_v_mad_u16_no_overflow(self):
"""V_MAD_U16 with clamp: 100 * 100 + 50 = 10050 (no saturation)."""
instructions = [
v_mov_b32_e32(v[0], 100),
v_mov_b32_e32(v[1], 100),
v_mov_b32_e32(v[2], 50),
v_mad_u16(v[3], v[0], v[1], v[2], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3] & 0xFFFF, 10050, f"expected 10050, got {st.vgpr[0][3] & 0xFFFF}")
def test_v_mad_u16_no_clamp(self):
"""V_MAD_U16 without clamp: 0xFFFF * 2 + 0 should wrap to 0xFFFE."""
instructions = [
v_mov_b32_e32(v[0], 0xFFFF),
v_mov_b32_e32(v[1], 2),
v_mov_b32_e32(v[2], 0),
v_mad_u16(v[3], v[0], v[1], v[2], clmp=0),
]
st = run_program(instructions, n_lanes=1)
# 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
class TestCvtPkF16(unittest.TestCase):
"""Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
def test_cvt_pk_rtz_f16_f32_basic(self):
"""V_CVT_PK_RTZ_F16_F32: basic pack of two f32 values."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], 2.0),
v_cvt_pk_rtz_f16_f32_e64(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo_f16 = f16(result & 0xffff)
hi_f16 = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo_f16, 1.0, delta=0.01)
self.assertAlmostEqual(hi_f16, 2.0, delta=0.01)
class TestCvtPkNorm(unittest.TestCase):
"""Tests for V_CVT_PK_NORM_I16_F32 and V_CVT_PK_NORM_U16_F32."""
def test_cvt_pk_norm_i16_f32_basic(self):
"""V_CVT_PK_NORM_I16_F32: pack two f32 to normalized i16."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], -1.0),
v_cvt_pk_norm_i16_f32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
self.assertEqual(lo, 32767)
self.assertEqual(hi, 0x8001) # -32767, hardware uses symmetric range
def test_cvt_pk_norm_u16_f32_basic(self):
"""V_CVT_PK_NORM_U16_F32: pack two f32 to normalized u16."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], 0.5),
v_cvt_pk_norm_u16_f32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
self.assertEqual(lo, 65535)
self.assertAlmostEqual(hi, 32768, delta=1)
class TestCvtPkInt(unittest.TestCase):
"""Tests for V_CVT_PK_I16_I32, V_CVT_PK_U16_U32, V_CVT_PK_I16_F32, V_CVT_PK_U16_F32."""
def test_cvt_pk_i16_i32_basic(self):
"""V_CVT_PK_I16_I32: pack two i32 to i16."""
instructions = [
s_mov_b32(s[0], 100),
s_mov_b32(s[1], -100 & 0xffffffff),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cvt_pk_i16_i32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
lo_signed = lo if lo < 32768 else lo - 65536
hi_signed = hi if hi < 32768 else hi - 65536
self.assertEqual(lo_signed, 100)
self.assertEqual(hi_signed, -100)
def test_cvt_pk_u16_u32_basic(self):
"""V_CVT_PK_U16_U32: pack two u32 to u16."""
instructions = [
s_mov_b32(s[0], 1000),
s_mov_b32(s[1], 2000),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cvt_pk_u16_u32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
self.assertEqual(lo, 1000)
self.assertEqual(hi, 2000)
def test_cvt_pk_i16_f32_basic(self):
"""V_CVT_PK_I16_F32: convert two f32 to packed i16."""
instructions = [
v_mov_b32_e32(v[0], 100.5),
v_mov_b32_e32(v[1], -50.7),
v_cvt_pk_i16_f32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
lo_signed = lo if lo < 32768 else lo - 65536
hi_signed = hi if hi < 32768 else hi - 65536
self.assertEqual(lo_signed, 100)
self.assertEqual(hi_signed, -50)
def test_cvt_pk_u16_f32_basic(self):
"""V_CVT_PK_U16_F32: convert two f32 to packed u16."""
instructions = [
v_mov_b32_e32(v[0], 100.9),
v_mov_b32_e32(v[1], 200.1),
v_cvt_pk_u16_f32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
self.assertEqual(lo, 100)
self.assertEqual(hi, 200)
def test_cvt_pk_u8_f32_basic(self):
"""V_CVT_PK_U8_F32: convert f32 to u8 and pack at byte position."""
instructions = [
v_mov_b32_e32(v[0], 128.5),
v_mov_b32_e32(v[1], 0),
v_mov_b32_e32(v[2], 0),
v_cvt_pk_u8_f32(v[2], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
byte0 = result & 0xff
self.assertEqual(byte0, 128)
class TestDotProduct(unittest.TestCase):
"""Tests for dot product instructions V_DOT4_U32_U8, V_DOT8_U32_U4."""
def test_v_dot4_u32_u8_basic(self):
"""V_DOT4_U32_U8: 4-element dot product of u8 vectors."""
src0 = 0x04030201 # {4, 3, 2, 1}
src1 = 0x01010101 # {1, 1, 1, 1}
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_u32_u8(v[2], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 10)
def test_v_dot4_u32_u8_with_accumulator(self):
"""V_DOT4_U32_U8 with non-zero accumulator."""
src0 = 0x02020202 # {2, 2, 2, 2}
src1 = 0x03030303 # {3, 3, 3, 3}
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 100),
v_dot4_u32_u8(v[2], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 124)
def test_v_dot8_u32_u4_basic(self):
"""V_DOT8_U32_U4: 8-element dot product of u4 vectors."""
# src0 = 8 nibbles: {1,2,3,4,5,6,7,8} packed as 0x87654321
# src1 = 8 nibbles: {1,1,1,1,1,1,1,1} packed as 0x11111111
# result = 1+2+3+4+5+6+7+8 = 36
src0 = 0x87654321
src1 = 0x11111111
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot8_u32_u4(v[2], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
self.assertEqual(result, 36)
class TestMinMaxF16Vop3(unittest.TestCase):
"""Tests for V_MIN3_F16, V_MAX3_F16, V_MED3_F16, V_MINMAX_F16, V_MAXMIN_F16."""
def test_v_min3_f16_basic(self):
"""V_MIN3_F16: minimum of three f16 values."""
instructions = [
s_mov_b32(s[0], f32_to_f16(3.0)),
s_mov_b32(s[1], f32_to_f16(1.0)),
s_mov_b32(s[2], f32_to_f16(2.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_min3_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 1.0, delta=0.01)
def test_v_max3_f16_basic(self):
"""V_MAX3_F16: maximum of three f16 values."""
instructions = [
s_mov_b32(s[0], f32_to_f16(1.0)),
s_mov_b32(s[1], f32_to_f16(3.0)),
s_mov_b32(s[2], f32_to_f16(2.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_max3_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 3.0, delta=0.01)
def test_v_med3_f16_basic(self):
"""V_MED3_F16: median of three f16 values."""
instructions = [
s_mov_b32(s[0], f32_to_f16(3.0)),
s_mov_b32(s[1], f32_to_f16(1.0)),
s_mov_b32(s[2], f32_to_f16(2.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_med3_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 2.0, delta=0.01)
def test_v_minmax_f16_basic(self):
"""V_MINMAX_F16: clamp(src0, min=src1, max=src2)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(2.5)),
s_mov_b32(s[1], f32_to_f16(1.0)),
s_mov_b32(s[2], f32_to_f16(2.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_minmax_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 2.0, delta=0.01)
def test_v_maxmin_f16_basic(self):
"""V_MAXMIN_F16: clamp(src0, min=src2, max=src1)."""
instructions = [
s_mov_b32(s[0], f32_to_f16(0.5)),
s_mov_b32(s[1], f32_to_f16(2.0)),
s_mov_b32(s[2], f32_to_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_maxmin_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 1.0, delta=0.01)
def test_v_min3_f16_with_neg(self):
"""V_MIN3_F16 with neg modifier: min(-3, 1, 2) = -3."""
instructions = [
s_mov_b32(s[0], f32_to_f16(3.0)),
s_mov_b32(s[1], f32_to_f16(1.0)),
s_mov_b32(s[2], f32_to_f16(2.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_min3_f16(v[3], -v[0], v[1], v[2]), # neg on first operand
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, -3.0, delta=0.01)
def test_v_max3_f16_with_abs(self):
"""V_MAX3_F16 with abs modifier: max(|-3|, 1, 2) = 3."""
instructions = [
s_mov_b32(s[0], f32_to_f16(-3.0)),
s_mov_b32(s[1], f32_to_f16(1.0)),
s_mov_b32(s[2], f32_to_f16(2.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_max3_f16(v[3], abs(v[0]), v[1], v[2]), # abs on first operand
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 3.0, delta=0.01)
def test_v_med3_f16_opsel_hi(self):
"""V_MED3_F16 with opsel reading from hi half."""
# Pack two f16 values: hi=5.0, lo=1.0
packed = (f32_to_f16(5.0) << 16) | f32_to_f16(1.0)
instructions = [
s_mov_b32(s[0], packed),
s_mov_b32(s[1], f32_to_f16(3.0)),
s_mov_b32(s[2], f32_to_f16(4.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]),
# Read hi half of v[0] (5.0), med3(5, 3, 4) = 4
v_med3_f16(v[3], v[0].h, v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 4.0, delta=0.01)
class TestSadHi(unittest.TestCase):
"""Tests for V_SAD_HI_U8 instruction."""
def test_v_sad_hi_u8_basic(self):
"""V_SAD_HI_U8: (sad << 16) + acc."""
# |1-5| + |2-6| + |3-7| + |4-8| = 16, << 16 = 0x100000, + 100 = 0x100064
instructions = [
v_mov_b32_e32(v[0], 0x04030201),
v_mov_b32_e32(v[1], 0x08070605),
v_mov_b32_e32(v[2], 100),
v_sad_hi_u8(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], (16 << 16) + 100)
def test_v_sad_hi_u8_zero_diff(self):
"""V_SAD_HI_U8: identical inputs gives acc only."""
instructions = [
v_mov_b32_e32(v[0], 0x12345678),
v_mov_b32_e32(v[2], 50),
v_sad_hi_u8(v[3], v[0], v[0], v[2]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 50)
class TestPermlane(unittest.TestCase):
"""Tests for V_PERMLANE16_B32 and V_PERMLANEX16_B32 instructions."""
def test_v_permlane16_b32_identity(self):
"""V_PERMLANE16_B32 with identity permutation (lane i reads from lane i within row)."""
# lanesel encodes 4 bits per position: position i gets lanesel[i*4+3:i*4]
# Identity: position 0->0, 1->1, ..., 15->15
# lanesel = 0xFEDCBA9876543210 (positions 15-0 in nibbles)
instructions = [
v_mov_b32_e32(v[0], 0xDEADBEEF), # source data
s_mov_b32(s[0], 0x76543210), # lanesel low (positions 0-7)
s_mov_b32(s[1], 0xFEDCBA98), # lanesel high (positions 8-15)
v_permlane16_b32(v[1], v[0], s[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 reads from lane 0 (position 0 -> lanesel[3:0] = 0)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
def test_v_permlane16_b32_broadcast(self):
"""V_PERMLANE16_B32 broadcast lane 0 to all lanes in row."""
# lanesel = all zeros -> all positions read from lane 0 within row
instructions = [
v_mov_b32_e32(v[0], 0xCAFEBABE), # source data
s_mov_b32(s[0], 0), # lanesel low = 0 (all read lane 0)
s_mov_b32(s[1], 0), # lanesel high = 0
v_permlane16_b32(v[1], v[0], s[0], s[1]),
]
st = run_program(instructions, n_lanes=4)
# All lanes read from lane 0 of their row
for lane in range(4):
self.assertEqual(st.vgpr[lane][1], 0xCAFEBABE)
def test_v_permlanex16_b32_identity(self):
"""V_PERMLANEX16_B32 cross-row read with identity selection."""
# In wave32: row 0 (lanes 0-15) reads from row 1 (lanes 16-31) and vice versa
# With single lane in row 0, it reads from lane 0 of row 1 (lane 16)
# But lane 16 doesn't exist in 1-lane test, so use 32 lanes
instructions = [
v_mov_b32_e32(v[0], 0x11111111), # All lanes have this initially
s_mov_b32(s[0], 0x76543210), # lanesel low
s_mov_b32(s[1], 0xFEDCBA98), # lanesel high
v_permlanex16_b32(v[1], v[0], s[0], s[1]),
]
st = run_program(instructions, n_lanes=32)
# Lane 0 in row 0 reads from lane 0 of row 1 (lane 16)
self.assertEqual(st.vgpr[0][1], 0x11111111)
# Lane 16 in row 1 reads from lane 0 of row 0 (lane 0)
self.assertEqual(st.vgpr[16][1], 0x11111111)
if __name__ == '__main__':
unittest.main()
+463 -31
View File
@@ -149,7 +149,6 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_src2_f16_lo(self):
"""V_FMA_MIX_F32 with src2 as f16 from lo bits."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], f2i(1.0)),
@@ -166,7 +165,6 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_src2_f16_hi(self):
"""V_FMA_MIX_F32 with src2 as f16 from hi bits."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_2 = f32_to_f16(2.0)
val = (f16_2 << 16) | 0
instructions = [
@@ -199,7 +197,6 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_with_abs_f16_src2_lo(self):
"""V_FMA_MIX_F32 with abs modifier on f16 src2 (lo half). Regression test for sin(1.0) bug."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg1 = f32_to_f16(-1.0) # 0xbc00
instructions = [
s_mov_b32(s[0], f2i(0.0)), # src0 = 0.0 (f32)
@@ -217,7 +214,6 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_with_neg_f16_src2_lo(self):
"""V_FMA_MIX_F32 with neg modifier on f16 src2 (lo half)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_1 = f32_to_f16(1.0) # 0x3c00
instructions = [
s_mov_b32(s[0], f2i(0.0)), # src0 = 0.0 (f32)
@@ -235,7 +231,6 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mix_f32_with_abs_f16_src2_hi(self):
"""V_FMA_MIX_F32 with abs modifier on f16 src2 (hi half)."""
from extra.assembly.amd.test.hw.helpers import f32_to_f16
f16_neg1 = f32_to_f16(-1.0) # 0xbc00
val = (f16_neg1 << 16) | 0 # -1.0 in hi, 0 in lo
instructions = [
@@ -254,7 +249,6 @@ class TestFmaMix(unittest.TestCase):
def test_v_fma_mixlo_f16(self):
"""V_FMA_MIXLO_F16 writes to low 16 bits of destination."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], f2i(2.0)),
v_mov_b32_e32(v[0], s[0]),
@@ -267,14 +261,13 @@ class TestFmaMix(unittest.TestCase):
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
lo = f16(st.vgpr[0][3] & 0xffff)
hi = (st.vgpr[0][3] >> 16) & 0xffff
self.assertAlmostEqual(lo, 7.0, places=1)
self.assertEqual(hi, 0xdead, f"hi should be preserved, got 0x{hi:04x}")
def test_v_fma_mixlo_f16_all_f32_sources(self):
"""V_FMA_MIXLO_F16 with all f32 sources."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], f2i(1.0)),
v_mov_b32_e32(v[0], s[0]),
@@ -286,13 +279,12 @@ class TestFmaMix(unittest.TestCase):
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
lo = f16(st.vgpr[0][3] & 0xffff)
# 1*2+3 = 5
self.assertAlmostEqual(lo, 5.0, places=1)
def test_v_fma_mixlo_f16_sin_case(self):
"""V_FMA_MIXLO_F16 case from sin kernel."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x3f800000), # f32 1.0
v_mov_b32_e32(v[3], s[0]),
@@ -305,7 +297,7 @@ class TestFmaMix(unittest.TestCase):
VOP3P(VOP3POp.V_FMA_MIXLO_F16, vdst=v[3], src0=v[3], src1=s[6], src2=v[5], opsel=0, opsel_hi=0, opsel_hi2=0),
]
st = run_program(instructions, n_lanes=1)
lo = _f16(st.vgpr[0][3] & 0xffff)
lo = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(lo, -3.14159, delta=0.01)
@@ -314,7 +306,6 @@ class TestVOP3P(unittest.TestCase):
def test_v_pk_add_f16_basic(self):
"""V_PK_ADD_F16 adds two packed f16 values."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0, lo=1.0
s_mov_b32(s[1], 0x44004200), # hi=4.0, lo=3.0
@@ -324,14 +315,13 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 4.0, places=2)
self.assertAlmostEqual(hi, 6.0, places=2)
def test_v_pk_mul_f16_basic(self):
"""V_PK_MUL_F16 multiplies two packed f16 values."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x42004000), # hi=3.0, lo=2.0
s_mov_b32(s[1], 0x45004400), # hi=5.0, lo=4.0
@@ -341,14 +331,13 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 8.0, places=1)
self.assertAlmostEqual(hi, 15.0, places=1)
def test_v_pk_fma_f16_basic(self):
"""V_PK_FMA_F16: D = A * B + C for packed f16."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x42004000), # A: hi=3.0, lo=2.0
s_mov_b32(s[1], 0x45004400), # B: hi=5.0, lo=4.0
@@ -360,8 +349,8 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 9.0, places=1) # 2*4+1
self.assertAlmostEqual(hi, 16.0, places=0) # 3*5+1
@@ -370,7 +359,6 @@ class TestVOP3P(unittest.TestCase):
Inline constants for VOP3P are f16 values in the low 16 bits only.
hi half of inline constant is 0, so hi result = v0.hi + 0 = 1.0.
"""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x3c003c00), # packed f16: hi=1.0, lo=1.0
v_mov_b32_e32(v[0], s[0]),
@@ -378,8 +366,8 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
# lo = 1.0 + 1.0 = 2.0, hi = 1.0 + 0.0 = 1.0 (inline const hi half is 0)
self.assertAlmostEqual(lo, 2.0, places=2)
self.assertAlmostEqual(hi, 1.0, places=2)
@@ -388,7 +376,6 @@ class TestVOP3P(unittest.TestCase):
"""V_PK_MUL_F16 with inline constant POS_TWO (2.0).
Inline constant has value only in low 16 bits, hi is 0.
"""
from extra.assembly.amd.test.hw.helpers import _f16
# v0 = packed (3.0, 4.0), multiply by POS_TWO
# lo = 3.0 * 2.0 = 6.0, hi = 4.0 * 0.0 = 0.0 (inline const hi is 0)
instructions = [
@@ -398,14 +385,120 @@ class TestVOP3P(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 6.0, places=1)
self.assertAlmostEqual(hi, 0.0, places=1)
def test_v_pk_add_u16_float_inline_const_opsel(self):
"""V_PK_ADD_U16 with float inline constant 2.0
Regression test: for integer packed ops, do not perform the f32->f16 conversion.
"""
# src1 = inline float constant 2.0
instructions = [
s_mov_b32(s[0], 0x00030005), # packed u16: hi=3, lo=5
v_mov_b32_e32(v[0], s[0]),
v_pk_add_u16(v[1], v[0], SrcEnum.POS_TWO, opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1]
lo = result & 0xffff
hi = (result >> 16) & 0xffff
# lo = 5 + 0x0000 = 0x0005, hi = 3 + 0x4000 = 0x4003
self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")
class TestWMMAF16(unittest.TestCase):
"""Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
Note: RDNA3 WMMA F16 uses 8 VGPRs for accumulator/output (same as F32 variant),
but values are packed as f16. This differs from RDNA4 which uses 4 VGPRs.
"""
def test_v_wmma_f16_16x16x16_f16_all_ones(self):
"""V_WMMA_F16_16X16X16_F16 with all ones produces 16.0 in f16."""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[16:23] (8 regs)
for i in range(16, 24):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[24:31] (8 regs)
for i in range(24, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[0:7] to zero (8 regs for RDNA3 WMMA F16)
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA F16: D = A @ B + C
instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 in f16, stored in lo 16 bits of each VGPR (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = f16(result & 0xffff)
self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
"""V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
instructions.append(s_mov_b32(s[1], 0x4500)) # f16 5.0 in lo bits only
# Initialize A matrix in v[16:23] (8 regs)
for i in range(16, 24):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[24:31] (8 regs)
for i in range(24, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[0:7] to 5.0 in lo bits (8 regs for RDNA3 WMMA F16)
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
# WMMA F16: D = A @ B + C
instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 + 5.0 = 21.0 in f16, stored in lo 16 bits (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = f16(result & 0xffff)
self.assertAlmostEqual(lo, 21.0, places=0, msg=f"v[{reg}] lane {lane}: expected 21.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
def test_v_wmma_f16_16x16x16_f16_high_registers(self):
"""V_WMMA_F16_16X16X16_F16 with high register indices.
Regression test: WMMA was using static register indices instead of dynamic.
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
"""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[64:71] (8 regs)
for i in range(64, 72):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[80:87] (8 regs)
for i in range(80, 88):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[96:103] to zero (8 regs for RDNA3 WMMA F16)
for i in range(96, 104):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA F16: D = A @ B + C, result in v[96:103]
instructions.append(v_wmma_f16_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
# Copy results to v[0:7] for checking
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 in f16, stored in lo 16 bits (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = f16(result & 0xffff)
self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions with F32 output."""
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
@@ -440,6 +533,75 @@ class TestWMMA(unittest.TestCase):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_f16_high_registers(self):
"""V_WMMA_F32_16X16X16_F16 with high register indices.
Regression test: WMMA was using static register indices instead of dynamic,
causing incorrect results when registers weren't at the default positions.
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
"""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[64:71]
for i in range(64, 72):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[80:87]
for i in range(80, 88):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[96:103] to zero
for i in range(96, 104):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA: D = A @ B + C, result in v[96:103]
instructions.append(v_wmma_f32_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
# Copy results to v[0:7] for checking
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
class TestWMMABF16(unittest.TestCase):
"""Tests for WMMA BF16 instructions."""
def test_v_wmma_f32_16x16x16_bf16_all_ones(self):
"""V_WMMA_F32_16X16X16_BF16 with all ones produces 16.0."""
instructions = []
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
instructions.append(s_mov_b32(s[0], 0x3f803f80))
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
instructions.append(v_wmma_f32_16x16x16_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_bf16_with_accumulator(self):
"""V_WMMA_F32_16X16X16_BF16 with non-zero accumulator."""
instructions = []
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
instructions.append(s_mov_b32(s[0], 0x3f803f80))
instructions.append(s_mov_b32(s[1], f2i(5.0)))
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
instructions.append(v_wmma_f32_16x16x16_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
expected = f2i(21.0) # 16 + 5
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
class TestSpecialOps(unittest.TestCase):
"""Tests for special operations (SAD, PERM, DOT2)."""
@@ -553,7 +715,6 @@ class TestPackedMixedSigns(unittest.TestCase):
def test_pk_add_f16_mixed_signs(self):
"""V_PK_ADD_F16 with mixed positive/negative values."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0xc0003c00), # packed: hi=-2.0, lo=1.0
s_mov_b32(s[1], 0x3c003c00), # packed: hi=1.0, lo=1.0
@@ -563,14 +724,13 @@ class TestPackedMixedSigns(unittest.TestCase):
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = _f16(result & 0xffff)
hi = _f16((result >> 16) & 0xffff)
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 2.0, places=2) # 1.0 + 1.0
self.assertAlmostEqual(hi, -1.0, places=2) # -2.0 + 1.0
def test_pk_mul_f16_zero(self):
"""V_PK_MUL_F16 with zero."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = [
s_mov_b32(s[0], 0x40004000), # packed: 2.0, 2.0
s_mov_b32(s[1], 0x00000000), # packed: 0.0, 0.0
@@ -583,5 +743,277 @@ class TestPackedMixedSigns(unittest.TestCase):
self.assertEqual(result, 0x00000000, "2.0 * 0.0 should be 0.0")
class TestDot2F32F16(unittest.TestCase):
"""Tests for V_DOT2_F32_F16 - dot product of f16 pairs producing f32."""
def test_v_dot2_f32_f16_basic(self):
"""V_DOT2_F32_F16: dot product of two packed f16 pairs -> f32."""
# src0 = {hi=2.0, lo=1.0}, src1 = {hi=4.0, lo=3.0}
# result = 1.0*3.0 + 2.0*4.0 + 0 = 3 + 8 = 11.0
src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(4.0) << 16) | f32_to_f16(3.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 11.0, places=2)
def test_v_dot2_f32_f16_with_accumulator(self):
"""V_DOT2_F32_F16 with non-zero f32 accumulator."""
# src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 5.0
# result = 1.0*1.0 + 1.0*1.0 + 5.0 = 7.0
src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], f2i(5.0)),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]), # same as src0
v_mov_b32_e32(v[2], s[1]),
v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 7.0, places=2)
def test_v_dot2_f32_f16_negative_values(self):
"""V_DOT2_F32_F16 with negative f16 values."""
# src0 = {hi=-2.0, lo=3.0}, src1 = {hi=1.0, lo=2.0}
# result = 3.0*2.0 + (-2.0)*1.0 + 0 = 6 - 2 = 4.0
# NOTE: Hardware DOT2 may have up to 1 ULP difference due to internal implementation
src0 = (f32_to_f16(-2.0) << 16) | f32_to_f16(3.0)
src1 = (f32_to_f16(1.0) << 16) | f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot2_f32_f16(v[3], v[0], v[1], v[2], opsel_hi=3, opsel_hi2=1),
]
st = run_program(instructions, n_lanes=1, ulp_tolerance=1)
result = i2f(st.vgpr[0][3])
self.assertAlmostEqual(result, 4.0, places=2)
class TestDot2F16F16(unittest.TestCase):
"""Tests for V_DOT2_F16_F16 - dot product of f16 pairs producing f16."""
def test_v_dot2_f16_f16_basic(self):
"""V_DOT2_F16_F16: dot product of two packed f16 pairs -> f16."""
# src0 = {hi=2.0, lo=1.0}, src1 = {hi=3.0, lo=2.0}
# result = 1.0*2.0 + 2.0*3.0 + 0 = 2 + 6 = 8.0 (f16)
src0 = (f32_to_f16(2.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(3.0) << 16) | f32_to_f16(2.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 8.0, places=1)
def test_v_dot2_f16_f16_with_accumulator(self):
"""V_DOT2_F16_F16 with non-zero f16 accumulator."""
# src0 = {hi=1.0, lo=1.0}, src1 = {hi=1.0, lo=1.0}, acc = 3.0 (f16)
# result = 1.0*1.0 + 1.0*1.0 + 3.0 = 5.0 (f16)
src0 = (f32_to_f16(1.0) << 16) | f32_to_f16(1.0)
acc = f32_to_f16(3.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[2], acc),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]), # same as src0
v_mov_b32_e32(v[2], s[2]),
v_dot2_f16_f16(v[3], v[0], v[1], v[2]),
]
st = run_program(instructions, n_lanes=1)
result = f16(st.vgpr[0][3] & 0xffff)
self.assertAlmostEqual(result, 5.0, places=1)
class TestSignedDotProducts(unittest.TestCase):
"""Tests for V_DOT4_I32_IU8 and V_DOT8_I32_IU4 with signed inputs."""
def test_v_dot4_i32_iu8_signed_both(self):
"""V_DOT4_I32_IU8 with both inputs signed (neg=0b011)."""
# src0 = {-1, -2, 3, 4} as i8 = {0xff, 0xfe, 0x03, 0x04}
# src1 = {1, 1, 1, 1} as i8
# result = (-1)*1 + (-2)*1 + 3*1 + 4*1 = -1 - 2 + 3 + 4 = 4
src0 = (0xff << 24) | (0xfe << 16) | (0x03 << 8) | 0x04 # -1, -2, 3, 4
src1 = 0x01010101 # 1, 1, 1, 1
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b011), # both signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
# Result is i32, interpret as signed
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, 4)
def test_v_dot4_i32_iu8_src0_signed(self):
"""V_DOT4_I32_IU8 with only src0 signed (neg=0b001)."""
# src0 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
# src1 = {2, 2, 2, 2} as u8
# result = (-1)*2 + (-1)*2 + (-1)*2 + (-1)*2 = -8
src0 = 0xffffffff # -1, -1, -1, -1 (as i8)
src1 = 0x02020202 # 2, 2, 2, 2 (as u8)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b001), # src0 signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, -8)
def test_v_dot4_i32_iu8_src1_signed(self):
"""V_DOT4_I32_IU8 with only src1 signed (neg=0b010)."""
# src0 = {2, 2, 2, 2} as u8
# src1 = {-1, -1, -1, -1} as i8 = {0xff, 0xff, 0xff, 0xff}
# result = 2*(-1) + 2*(-1) + 2*(-1) + 2*(-1) = -8
src0 = 0x02020202 # 2, 2, 2, 2 (as u8)
src1 = 0xffffffff # -1, -1, -1, -1 (as i8)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0b010), # src1 signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, -8)
def test_v_dot4_i32_iu8_unsigned_as_reference(self):
"""V_DOT4_I32_IU8 with both unsigned (neg=0) - same as V_DOT4_U32_U8."""
# src0 = {0xff, 0xff, 0xff, 0xff} = 255 each as u8
# src1 = {1, 1, 1, 1}
# result = 255*1 + 255*1 + 255*1 + 255*1 = 1020
src0 = 0xffffffff
src1 = 0x01010101
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot4_i32_iu8(v[3], v[0], v[1], v[2], neg=0), # both unsigned
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 1020)
def test_v_dot8_i32_iu4_signed_both(self):
"""V_DOT8_I32_IU4 with both inputs signed (neg=0b011)."""
# src0 = 8 nibbles: {-1, -2, 3, 4, -1, -2, 3, 4} as i4
# i4 -1 = 0xf, -2 = 0xe, 3 = 0x3, 4 = 0x4
# src0 = 0xfe34fe34
# src1 = {1, 1, 1, 1, 1, 1, 1, 1} as i4 = 0x11111111
# result = 2 * ((-1)*1 + (-2)*1 + 3*1 + 4*1) = 2 * 4 = 8
src0 = 0xfe34fe34
src1 = 0x11111111
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011), # both signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, 8)
def test_v_dot8_i32_iu4_all_negative(self):
"""V_DOT8_I32_IU4 with all negative signed values."""
# src0 = 8 nibbles all -1 (0xf) = 0xffffffff
# src1 = 8 nibbles all 1 = 0x11111111
# result = 8 * ((-1)*1) = -8
src0 = 0xffffffff # all -1 as i4
src1 = 0x11111111 # all 1
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], 0),
v_dot8_i32_iu4(v[3], v[0], v[1], v[2], neg=0b011), # both signed
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][3]
if result >= 0x80000000:
result = result - 0x100000000
self.assertEqual(result, -8)
class TestPkMinMaxF16(unittest.TestCase):
"""Tests for V_PK_MIN_F16 and V_PK_MAX_F16."""
def test_v_pk_min_f16_basic(self):
"""V_PK_MIN_F16: packed min of two f16 pairs."""
# src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
# result = {min(3,2)=2, min(1,4)=1}
src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_min_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 1.0, delta=0.01)
self.assertAlmostEqual(hi, 2.0, delta=0.01)
def test_v_pk_max_f16_basic(self):
"""V_PK_MAX_F16: packed max of two f16 pairs."""
# src0 = {hi=3.0, lo=1.0}, src1 = {hi=2.0, lo=4.0}
# result = {max(3,2)=3, max(1,4)=4}
src0 = (f32_to_f16(3.0) << 16) | f32_to_f16(1.0)
src1 = (f32_to_f16(2.0) << 16) | f32_to_f16(4.0)
instructions = [
s_mov_b32(s[0], src0),
s_mov_b32(s[1], src1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_pk_max_f16(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2]
lo = f16(result & 0xffff)
hi = f16((result >> 16) & 0xffff)
self.assertAlmostEqual(lo, 4.0, delta=0.01)
self.assertAlmostEqual(hi, 3.0, delta=0.01)
if __name__ == '__main__':
unittest.main()
+105
View File
@@ -731,6 +731,111 @@ class TestVCCBehavior(unittest.TestCase):
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
class TestCmpNge(unittest.TestCase):
"""Tests for V_CMP_NGE (not-greater-or-equal) with NaN semantics.
NGE = !(a >= b). With NaN inputs:
- If either input is NaN, a >= b is false, so !(false) = true
- This differs from a < b which returns false for NaN inputs
"""
def test_v_cmp_nge_f32_normal_values(self):
"""v_cmp_nge_f32: basic comparison with normal floats."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(2.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 2.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= 2.0) should be true")
def test_v_cmp_nge_f32_equal_values(self):
"""v_cmp_nge_f32: equal values should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(1.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_greater_value(self):
"""v_cmp_nge_f32: greater value should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(2.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(2.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_neg_inf(self):
"""v_cmp_nge_f32: -inf compared to normal value."""
neg_inf = 0xff800000 # -inf
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(-inf >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(-inf >= 1.0) should be true")
def test_v_cmp_nge_f32_clears_inactive_vcc_bits(self):
"""v_cmp_nge_f32 with partial EXEC clears inactive VCC bits (hardware behavior)."""
neg_inf = 0xff800000 # -inf
instructions = [
# Set VCC to all 1s first
s_mov_b32(VCC_LO, 0xFFFFFFFF),
# Set EXEC to only lane 0
s_mov_b32(EXEC_LO, 0x00000001),
# v0 = 1.0 for lane 0
v_mov_b32_e32(v[0], f2i(1.0)),
# Compare: !(-inf >= 1.0) = true for lane 0
v_cmp_nge_f32_e32(neg_inf, v[0]),
]
st = run_program(instructions, n_lanes=16)
# Hardware clears inactive lane bits, only active lane results remain
# Lane 0 result = 1 (true), lanes 1-15 = 0 (cleared)
self.assertEqual(st.vcc, 0x00000001, "VCC should only have active lane results")
def test_v_cmp_nge_f32_nan_src0(self):
"""v_cmp_nge_f32: NaN in src0 should return true (NaN >= x is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= 1.0) should be true")
def test_v_cmp_nge_f32_nan_src1(self):
"""v_cmp_nge_f32: NaN in src1 should return true (x >= NaN is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= NaN) should be true")
def test_v_cmp_nge_f32_both_nan(self):
"""v_cmp_nge_f32: both NaN should return true."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= NaN) should be true")
class TestCmpxPartialWavefront(unittest.TestCase):
"""Tests for V_CMPX with partial wavefronts (fewer than 32 active lanes).
+135 -29
View File
@@ -1,18 +1,14 @@
# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
import unittest, ctypes, os
import unittest, ctypes
from dataclasses import dataclass
from pathlib import Path
# Set environment before any tinygrad imports to use MOCKGPU
# This allows generating AMD GPU kernels without requiring real hardware
os.environ["AMD"] = "1"
os.environ["MOCKGPU"] = "1"
os.environ["PYTHON_REMU"] = "1"
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
from extra.assembly.amd import decode_inst
from extra.assembly.amd.test.helpers import KernelInfo
from extra.assembly.amd.test.bench_emu import REMU_PATH
def set_valid_mem_ranges(ranges): pass # emu2 doesn't need this
def _is_f32_nan(bits: int) -> bool:
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
@@ -91,33 +87,61 @@ class PythonEmulator:
def __init__(self):
self.state: WaveState | None = None
self.program: dict | None = None
self.vmem_buf = None
self.lds_buf = None
self.kernel_buf = None # Keep kernel bytes alive
self.lib_addr = 0 # Base address of kernel code
def create(self, kernel: bytes, n_lanes: int):
self.program = decode_program(kernel)
self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
self.state.exec_mask = (1 << n_lanes) - 1
import ctypes
from tinygrad.device import Buffer, BufferSpec
from tinygrad.dtype import dtypes
# Store kernel in a ctypes buffer so generic instructions can read from vmem at actual PC address
self.kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
self.lib_addr = ctypes.addressof(self.kernel_buf)
# Remap program dict to use actual addresses (like run_asm does)
program_raw = decode_program(kernel)
self.program = {self.lib_addr + offset: val for offset, val in program_raw.items()}
self.state = WaveState(n_lanes)
self.state.pc = self.lib_addr # Set PC to code base address
self.vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
self.lds_buf = Buffer('CPU', 65536 // 4, dtypes.uint32).ensure_allocated()
def step(self) -> int:
import ctypes
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])
pc = self.state.pc
if pc == 0xFFFFFFFFFFFFFFFF or pc not in self.program: return -1
name, fxn, globals_list, _runner = self.program[pc]
if fxn is None: return 1 # unsupported instruction
buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,
2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}
# Direct ctypes call - bypasses HCQ overhead
fxn(*[ctypes.c_uint64(buf_addrs[g]) for g in globals_list], ctypes.c_int32(0))
return -1 if self.state.pc == 0xFFFFFFFFFFFFFFFF else 0
def set_sgpr(self, idx: int, val: int):
assert self.state is not None
self.state.sgpr[idx] = val & 0xffffffff
self.state._write_sgpr(idx, val)
def set_vgpr(self, lane: int, idx: int, val: int):
assert self.state is not None
self.state.vgpr[lane][idx] = val & 0xffffffff
self.state._write_vgpr(idx, lane, val)
def get_snapshot(self) -> StateSnapshot:
assert self.state is not None
return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
sgpr = [self.state._read_sgpr(i) for i in range(128)]
vgpr = [[self.state._read_vgpr(reg, lane) for reg in range(256)] for lane in range(WAVE_SIZE)]
# Convert actual PC address to word offset for comparison with Rust emulator
pc_offset = (self.state.pc - self.lib_addr) // 4 if self.state.pc != 0xFFFFFFFFFFFFFFFF else 0xFFFFFFFFFFFFFFFF
return StateSnapshot(pc=pc_offset, scc=self.state._read_sgpr(SCC.offset), vcc=sgpr[VCC_LO.offset],
exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
max_workgroups: int = 8) -> tuple[bool, str, int]:
local_size: tuple[int, int, int], program, max_steps: int, debug: bool, trace_len: int,
kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, int]:
"""Run a single kernel through both emulators. Returns (success, message, total_steps)."""
gx, gy, gz = global_size
lx, ly, lz = local_size
total_steps = 0
wg_count = 0
@@ -140,28 +164,52 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
emu.set_sgpr(13, gidx)
emu.set_sgpr(14, gidy)
emu.set_sgpr(15, gidz)
# Initialize v[0] with packed workitem IDs for each lane
for lane in range(n_lanes):
tid = lane
z, y, x = tid // (lx * ly), (tid // lx) % ly, tid % lx
emu.set_vgpr(lane, 0, (z << 20) | (y << 10) | x)
step = 0
trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
prev_sync_after = False # Track if previous instruction had known Rust bugs
try:
while step < max_steps:
rust_before = rust.get_snapshot()
python_before = python.get_snapshot()
inst = program.get(python_before.pc)
inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
inst_info = python.program.get(python.lib_addr + python_before.pc * 4) # Convert word offset to actual address
inst_hex_name = inst_info[0] if inst_info else f"unknown at PC={python_before.pc}"
# Decode the instruction to get mnemonic for sync_after checks
try:
# Format is mnemonic_hexbytes, e.g. v_exp_f32_e32_014b027e -> hex is 014b027e
parts = inst_hex_name.rsplit('_', 1)
inst_bytes_hex = parts[1] if len(parts) == 2 else ""
inst_bytes = bytes.fromhex(inst_bytes_hex) if inst_bytes_hex else b''
decoded = decode_inst(inst_bytes) if inst_bytes else None
inst_mnemonic = repr(decoded).split('(')[0] if decoded else ""
except:
inst_mnemonic = ""
# For generic instructions, use function name for sync_after check
if not inst_mnemonic: inst_mnemonic = inst_hex_name
inst_str = inst_hex_name
trace.append((step, python_before.pc, inst_str, rust_before, python_before))
if len(trace) > trace_len: trace.pop(0)
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
# Instructions with known Rust emulator bugs or precision differences - sync Python to Rust after execution
# v_div_scale/v_div_fixup: Rust has different VCC handling
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
# s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
diffs = rust_before.diff(python_before, n_lanes)
# v_exp_f32/v_log_f32/v_ldexp_f32: precision differences in transcendental functions
# s_delay_alu: Rust handles differently
# v_add_co_ci_u32/v_sub_co_ci_u32/v_subrev_co_ci_u32: Rust preserves inactive VCC bits, but hardware clears all bits
sync_after = any(x in inst_mnemonic.lower() for x in ('v_div_scale', 'v_div_fixup', 'v_cvt_f16_f32', 's_add_i32', 's_sub_i32',
'v_exp_f32', 'v_log_f32', 'v_ldexp_f32', 's_delay_alu',
'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'))
# Skip comparison if previous instruction had known Rust bugs (states were synced but may still differ slightly)
diffs = rust_before.diff(python_before, n_lanes) if not prev_sync_after else []
if diffs:
trace_lines = []
for idx, (s, pc, d, rb, pb) in enumerate(trace):
@@ -200,7 +248,12 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
for lane in range(n_lanes):
for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
assert python.state is not None
python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
# Convert Rust's word-based PC to Python's actual address
python.state.pc = python.lib_addr + rust_after.pc * 4
python.state._write_sgpr(SCC.offset, rust_after.scc)
python.state._write_sgpr(VCC_LO.offset, rust_after.vcc)
python.state._write_sgpr(EXEC_LO.offset, rust_after.exec_mask)
prev_sync_after = sync_after
if rust_result == -1:
total_steps += step + 1
@@ -254,7 +307,7 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
ok, msg, steps = run_single_kernel(
kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
program, max_steps, debug, trace_len, ki
kernel.local_size, program, max_steps, debug, trace_len, ki
)
total_steps += steps
if not ok:
@@ -281,7 +334,8 @@ def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list,
set_valid_mem_ranges(ranges)
program = decode_program(kernel)
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
# Legacy wrapper assumes local_size = (n_lanes, 1, 1)
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), program, max_steps, debug, trace_len)
return ok, msg
def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
@@ -324,6 +378,7 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int],
buf_sizes.append(b.nbytes)
kernels.append(KernelInfo(
code=bytes(sec.content),
src=lowered.prg.p.src,
global_size=tuple(lowered.prg.p.global_size),
local_size=tuple(lowered.prg.p.local_size),
buf_idxs=buf_idxs,
@@ -386,6 +441,7 @@ class TestTinygradKernels(unittest.TestCase):
from tinygrad import dtypes
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
def test_exp(self): self._test_kernel(lambda T: T.empty(1024).exp())
def test_cross_entropy(self):
import numpy as np
np.random.seed(0)
@@ -397,5 +453,55 @@ class TestTinygradKernels(unittest.TestCase):
from tinygrad import dtypes
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
def test_sin_large_f32(self):
"""Test sin with large values that trigger Payne-Hanek range reduction."""
# Values around 859240 trigger the Payne-Hanek algorithm
# This tests the integer multiply-high instructions used in range reduction
self._test_kernel(lambda T: T([859240.0, 1000000.0, 100594688.0]).sin())
def test_clip_zero_one(self):
"""Test clip(0, 1) - regression for binary_crossentropy failure."""
import numpy as np
np.random.seed(0)
x_np = np.random.uniform(-2, 2, (32, 10)).astype(np.float32).tolist()
self._test_kernel(lambda T: T(x_np).clip(0, 1))
def test_mod_int64(self):
"""Test int64 modulo, especially edge cases like 1 % -1."""
from tinygrad import dtypes
self._test_kernel(lambda T: T([1, 10, -10, 7], dtype=dtypes.int64) % T([-1, 3, 3, -3], dtype=dtypes.int64))
def test_expand_flatten_sum(self):
"""Test flatten of expanded tensor followed by sum.
Bug: flatten() of an expanded tensor produces wrong results for certain sizes.
Sizes that are multiples of 32 work (32, 48, 64), but sizes like 33, 49, 50 fail.
This breaks masked_select and nonzero operations.
"""
import numpy as np
np.random.seed(0)
x_np = np.random.uniform(-2, 2, (33,)).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).unsqueeze(-1).expand(33, 3).flatten().sum())
@unittest.skip("slow and broken with AMD_LLVM=1")
def test_nonzero(self):
"""Test nonzero operation - counts and gathers indices of non-zero elements."""
import numpy as np
np.random.seed(42)
x_np = np.random.rand(10, 5, 3).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).nonzero())
@unittest.skip("Precision differences in v_exp/v_log accumulate across kernels, causing memory divergence")
def test_softmax_argmax_fused(self):
"""Test fused softmax+argmax - tracks exp2 precision issue.
The fused kernel recomputes softmax inline and Python emulator's exp2 polynomial
has up to 1 ULP error vs native exp2f, causing accumulated differences.
"""
import torch
torch.manual_seed(0)
x_np = torch.rand(4, 10).numpy()
self._test_kernel(lambda T: T(x_np.tolist()).softmax(1).argmax())
if __name__ == "__main__":
unittest.main()
+30 -1
View File
@@ -34,15 +34,44 @@ def custom_add_one(A:UOp, arch:str) -> UOp:
sink = UOp.sink(A.base, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))
def custom_add_var(A:UOp, B:UOp, arch:str) -> UOp:
A,B = A.flatten(), B.flatten()
assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
threads = UOp.special(A.size, "lidx0")
var = UOp.variable("var", 0, 10)
insts = [
s_load_b128(s[4:7], s[0:1]),
s_load_b32(s[8], s[0:1], offset=0x10), # all threads load the same variable
s_waitcnt(lgkmcnt=0),
v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset, different per thread
global_load_b32(v[1], v[0], saddr=s[6:7]),
s_waitcnt(vmcnt=0),
v_add_nc_u32_e32(v[1], s[8], v[1]),
global_store_b32(addr=v[0], data=v[1], saddr=s[4:5]),
s_endpgm(),
]
sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)),
*assemble_insts(insts, name, arch, kernarg_size=16)))
class TestCustomKernel(unittest.TestCase):
def test_simple(self):
a = Tensor.full((16, 16), 1.).contiguous().realize()
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].arch))[0]
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].renderer.arch))[0]
ei = a.schedule()[-1].lower()
self.assertEqual(ei.prg.estimates.ops, a.numel())
self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
ei.run()
self.assertTrue((a.numpy() == 2.).all())
def test_variable(self):
b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
a = Tensor.zeros_like(b).contiguous().realize()
a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_var, arch=Device[Device.DEFAULT].renderer.arch))[0]
ei = a.schedule()[-1].lower()
for i in range(4):
ei.run({"var":i})
self.assertTrue((a.numpy() == 1+i).all())
if __name__ == "__main__":
unittest.main()
+332
View File
@@ -0,0 +1,332 @@
"""Tests for the pcode parser."""
import unittest
from collections import defaultdict
from tinygrad.helpers import DEBUG
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, Ops
from extra.assembly.amd.emu import parse_pcode
from extra.assembly.amd.pcode import parse_expr
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op, VOP3Op, SOP1Op, SOP2Op, DSOp
def _srcs():
"""Create minimal source variables for pcode parsing."""
u32 = lambda v=0: UOp.const(dtypes.uint32, v)
return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}
class TestBasicParsing(unittest.TestCase):
"""Test basic pcode parsing for common instruction patterns."""
def test_v_add_f32(self):
"""Test parsing V_ADD_F32 pcode."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], _srcs())
self.assertEqual(len(assigns), 1)
dest, _ = assigns[0]
self.assertTrue(dest.startswith('D0'))
def test_v_lshlrev_b32(self):
"""Test parsing V_LSHLREV_B32 pcode."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_LSHLREV_B32_E32], _srcs())
self.assertEqual(len(assigns), 1)
def test_s_cselect_b32(self):
"""Test parsing S_CSELECT_B32 pcode with ternary."""
_, assigns = parse_pcode(PCODE[SOP2Op.S_CSELECT_B32], _srcs())
self.assertEqual(len(assigns), 1)
def test_v_add_co_ci_u32(self):
"""Test parsing V_ADD_CO_CI_U32 with carry."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_CO_CI_U32_E32], _srcs())
self.assertGreaterEqual(len(assigns), 1)
class TestWithSources(unittest.TestCase):
"""Test pcode parsing with actual source operand values."""
def test_v_add_f32_with_sources(self):
"""Test V_ADD_F32 with actual float constants."""
s0 = UOp.const(dtypes.uint32, 0x3f800000) # 1.0f
s1 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be an ADD operation
self.assertEqual(val.op, Ops.ADD)
def test_v_mul_f32_with_sources(self):
"""Test V_MUL_F32 with actual float constants."""
s0 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f
s1 = UOp.const(dtypes.uint32, 0x40400000) # 3.0f
_, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertEqual(val.op, Ops.MUL)
class TestParseExpr(unittest.TestCase):
"""Test the parse_expr function directly."""
def test_integer_literals(self):
"""Test parsing integer literals."""
self.assertEqual(parse_expr('0', {}).arg, 0)
self.assertEqual(parse_expr('42', {}).arg, 42)
self.assertEqual(parse_expr('42U', {}).arg, 42)
def test_negative_integers(self):
"""Test parsing negative integer literals."""
result = parse_expr('-1', {})
self.assertEqual(result.arg, -1)
self.assertEqual(result.dtype, dtypes.int)
def test_float_literals(self):
"""Test parsing float literals."""
result = parse_expr('1.0F', {})
self.assertEqual(result.arg, 1.0)
self.assertEqual(result.dtype, dtypes.float32)
def test_hex_literals(self):
"""Test parsing hex literals."""
result = parse_expr('0xFF', {})
self.assertEqual(result.arg, 255)
def test_variable_lookup(self):
"""Test variable lookup in parse_expr."""
vars = {'x': UOp.const(dtypes.uint32, 42)}
result = parse_expr('x', vars)
self.assertEqual(result.arg, 42)
def test_binary_ops(self):
"""Test parsing binary operations."""
vars = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
# Addition
result = parse_expr('a + b', vars)
self.assertEqual(result.op, Ops.ADD)
# Subtraction with constant folding
result = parse_expr('10 - 5', {})
self.assertEqual(result.op, Ops.CONST)
self.assertEqual(result.arg, 5)
def test_ternary(self):
"""Test parsing ternary expressions."""
vars = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
result = parse_expr('cond ? a : b', vars)
self.assertEqual(result.op, Ops.WHERE)
class TestForLoopParsing(unittest.TestCase):
"""Test for loop parsing (CLZ/CTZ patterns)."""
def test_clz_pcode_exists(self):
"""Verify CLZ pcode is available."""
pcode = PCODE.get(VOP1Op.V_CLZ_I32_U32_E32)
self.assertIsNotNone(pcode)
self.assertIn('for', pcode.lower())
def test_clz_parsing(self):
"""Test CLZ pcode parsing produces correct structure."""
pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF) # All ones - CLZ should be 0
vars, assigns = parse_pcode(pcode, {'S0': S0})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be a nested WHERE structure
self.assertEqual(val.op, Ops.WHERE)
def test_clz_with_zero(self):
"""Test CLZ with input 0 - should return -1."""
pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
S0 = UOp.const(dtypes.uint32, 0)
vars, assigns = parse_pcode(pcode, {'S0': S0})
# Check that the innermost value (default) is -1 (may be wrapped in CAST)
val = assigns[0][1]
# Traverse to innermost WHERE
while val.op == Ops.WHERE:
val = val.src[2] # false branch
# Unwrap CAST if present
while val.op == Ops.CAST:
val = val.src[0]
self.assertEqual(val.arg, -1)
def test_ctz_parsing(self):
"""Test CTZ pcode parsing."""
pcode = PCODE.get(VOP1Op.V_CTZ_I32_B32_E32)
if pcode is None:
self.skipTest("V_CTZ_I32_B32_E32 pcode not available")
S0 = UOp.const(dtypes.uint32, 1) # LSB set - CTZ should be 0
vars, assigns = parse_pcode(pcode, {'S0': S0})
self.assertEqual(len(assigns), 1)
class TestDSPcodePatterns(unittest.TestCase):
"""Test DS instruction pcode patterns."""
def test_ds_load_b32_pcode(self):
"""Test DS_LOAD_B32 pcode is parseable."""
pcode = PCODE.get(DSOp.DS_LOAD_B32)
self.assertIsNotNone(pcode)
self.assertIn('RETURN_DATA', pcode)
self.assertIn('MEM[', pcode)
def test_ds_store_b32_pcode(self):
"""Test DS_STORE_B32 pcode is parseable."""
pcode = PCODE.get(DSOp.DS_STORE_B32)
self.assertIsNotNone(pcode)
self.assertIn('MEM[', pcode)
self.assertIn('DATA', pcode)
def test_mem_read_parsing(self):
"""Test MEM[addr].type read expression parsing."""
# Create a mock LDS buffer
lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
addr = UOp.const(dtypes.uint32, 0)
vars = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
result = parse_expr('MEM[ADDR + OFFSET].b32', vars)
# Should be an INDEX operation into LDS
self.assertIsNotNone(result)
def test_ds_store_2addr_b32_parsing(self):
"""Test DS_STORE_2ADDR_B32 pcode parsing produces MEM writes."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
self.assertIsNotNone(pcode)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Should have 2 MEM write assignments
self.assertEqual(len(assigns), 2)
for dest, val in assigns:
self.assertTrue(dest.startswith('MEM['))
# val should be (addr, write_val) tuple
self.assertIsInstance(val, tuple)
self.assertEqual(len(val), 2)
def test_ds_load_2addr_b32_parsing(self):
"""Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
self.assertIsNotNone(pcode)
lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'_lds': lds,
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Should have 2 RETURN_DATA assignments
self.assertEqual(len(assigns), 2)
self.assertEqual(assigns[0][0], 'RETURN_DATA[31:0]')
self.assertEqual(assigns[1][0], 'RETURN_DATA[63:32]')
def test_ds_store_address_calculation(self):
"""Test DS_STORE_2ADDR_B32 calculates correct addresses (offset * 4)."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 100),
'OFFSET0': UOp.const(dtypes.uint32, 2),
'OFFSET1': UOp.const(dtypes.uint32, 5),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
addr0, _ = assigns[0][1]
addr1, _ = assigns[1][1]
self.assertEqual(addr0.simplify().arg, 108)
self.assertEqual(addr1.simplify().arg, 120)
def test_ds_store_data_values(self):
"""Test DS_STORE_2ADDR_B32 uses correct data values."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
_, val0 = assigns[0][1]
_, val1 = assigns[1][1]
# DATA[31:0] should preserve the value
self.assertEqual(val0.simplify().arg, 0xAAAAAAAA)
self.assertEqual(val1.simplify().arg, 0xBBBBBBBB)
class TestConditionalParsing(unittest.TestCase):
"""Test conditional (if/elsif/else) pcode parsing."""
def test_ternary_in_assignment(self):
"""Test parsing ternary expression (which becomes WHERE)."""
# S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32
pcode = PCODE[SOP2Op.S_CSELECT_B32]
s0 = UOp.const(dtypes.uint32, 10)
s1 = UOp.const(dtypes.uint32, 20)
scc = UOp.const(dtypes.uint32, 1)
vars, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be a WHERE (ternary becomes WHERE)
self.assertEqual(val.op, Ops.WHERE)
class TestAllPcode(unittest.TestCase):
"""Test that all pcode from all architectures can be parsed."""
def _make_srcs(self):
"""Create dummy source variables for pcode parsing."""
u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
'SIMM16': u32(), 'SIMM32': u32(), 'OFFSET': u32(), 'OFFSET0': u32(), 'OFFSET1': u32(), 'offset1': u32(),
'ADDR': u32(), 'ADDR_BASE': u32(), 'TADDR': u32(), 'DATA': u32(), 'DATA0': u32(), 'DATA1': u32(), 'DATA2': u32(),
'VDATA': u32(), 'VDATA0': u32(), 'VDATA1': u32(), 'VDATA2': u32(), 'VDATA3': u32(),
'OPSEL': u32(), 'OPSEL_HI': u32(), 'NEG': u32(), 'NEG_HI': u32(), 'CLAMP': u32(),
'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(), 'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
'_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}
def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float):
"""Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs."""
srcs = self._make_srcs()
passed, skipped, errors = 0, 0, defaultdict(list)
for op, pcode in pcode_dict.items():
try:
parse_pcode(pcode, srcs)
passed += 1
except RuntimeError as e: skipped += 1; errors[str(e)].append(op.name)
except Exception as e: self.fail(f"[{arch}] {op.name}: {e}\nPcode: {pcode[:200]}")
total = len(pcode_dict)
pct = 100 * passed / total
print(f"{arch}: {passed}/{total} ({pct:.1f}%) parsed, {skipped} skipped")
if DEBUG >= 2:
for err, ops in sorted(errors.items(), key=lambda x: -len(x[1])):
print(f" {err}: {', '.join(ops[:5])}{'...' if len(ops) > 5 else ''} ({len(ops)})")
self.assertGreaterEqual(pct, min_pct, f"[{arch}] {pct:.1f}% < {min_pct}% threshold")
def test_parse_all_cdna_pcode(self):
from extra.assembly.amd.autogen.cdna.str_pcode import PCODE as CDNA_PCODE
self._parse_all_pcode(CDNA_PCODE, "CDNA", min_pct=60)
def test_parse_all_rdna3_pcode(self):
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE as RDNA3_PCODE
self._parse_all_pcode(RDNA3_PCODE, "RDNA3", min_pct=90)
def test_parse_all_rdna4_pcode(self):
from extra.assembly.amd.autogen.rdna4.str_pcode import PCODE as RDNA4_PCODE
self._parse_all_pcode(RDNA4_PCODE, "RDNA4", min_pct=65)
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -7,7 +7,7 @@ import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import VCC_HI, EXEC_LO, NULL
OFF = NULL # OFF is alias for NULL
from extra.assembly.amd.decode import detect_format
from extra.assembly.amd import detect_format
class TestDS(unittest.TestCase):
+17 -2
View File
@@ -6,17 +6,21 @@ from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.test.test_roundtrip import compile_asm
class TestIntegration(unittest.TestCase):
class IntegrationTestBase(unittest.TestCase):
inst: Inst
arch: str
def tearDown(self):
if not hasattr(self, 'inst'): return
b = self.inst.to_bytes()
st = self.inst.disasm()
# Test that the instruction can be compiled by LLVM and produces the same bytes
desc = f"{st:25s} {self.inst} {b!r}"
self.assertEqual(b, compile_asm(st), desc)
self.assertEqual(b, compile_asm(st, arch=self.arch), desc)
print(desc)
class TestIntegration(IntegrationTestBase):
arch: str = "rdna3"
def test_wmma(self):
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[184:191], v[136:143], v[0:7])
@@ -124,6 +128,17 @@ class TestIntegration(unittest.TestCase):
int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
self.assertEqual(self.inst, int_inst)
class TestIntegrationCDNA(IntegrationTestBase):
arch = "cdna"
def test_mfma(self):
from extra.assembly.amd.autogen.cdna.ins import v_mfma_f32_16x16x16_f16
self.inst = v_mfma_f32_16x16x16_f16(v[0:3], v[0:1], v[0:1], 0)
def test_mfma_fp8(self):
from extra.assembly.amd.autogen.cdna.ins import v_mfma_f32_16x16x128_f8f6f4
self.inst = v_mfma_f32_16x16x128_f8f6f4(v[0:3], v[0:5], v[0:5], 1, cbsz=2, blgp=2)
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
+24 -1
View File
@@ -11,7 +11,7 @@ Only compute-relevant instruction formats are tested. Graphics-only formats not
import unittest, re, subprocess, functools
from tinygrad.helpers import fetch
from extra.assembly.amd.disasm import disasm
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_target, get_mattr
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
@@ -125,6 +125,26 @@ def _make_test(f: str, arch: str, test_type: str):
except ValueError: skipped += 1 # skip invalid opcodes not in enum
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "repr":
# Test that eval(repr(inst)) reproduces the instruction
if arch == "rdna3": import extra.assembly.amd.autogen.rdna3.ins as ins
elif arch == "rdna4": import extra.assembly.amd.autogen.rdna4.ins as ins
elif arch == "cdna": import extra.assembly.amd.autogen.cdna.ins as ins
ns = {k: getattr(ins, k) for k in dir(ins) if not k.startswith('_')}
passed, skipped = 0, 0
for _, data in tests:
try:
decoded = detect_format(data, arch).from_bytes(data)
if decoded.to_bytes()[:len(data)] != data: skipped += 1; continue # skip if binary roundtrip fails
r = repr(decoded)
try:
decoded2 = eval(r, ns) # noqa: S307
if decoded == decoded2: passed += 1
else: skipped += 1
except Exception: skipped += 1
except ValueError: skipped += 1
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "disasm":
to_test = []
for _, data in tests:
@@ -149,12 +169,15 @@ class TestLLVM(unittest.TestCase): pass
for f in RDNA_FILES:
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
setattr(TestLLVM, f"test_rdna3_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "repr"))
for f in CDNA_FILES:
setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
setattr(TestLLVM, f"test_cdna_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "repr"))
for f in RDNA4_FILES:
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
setattr(TestLLVM, f"test_rdna4_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "repr"))
if __name__ == "__main__":
unittest.main()
-403
View File
@@ -1,403 +0,0 @@
#!/usr/bin/env python3
"""Tests for the RDNA3 pseudocode DSL."""
import unittest
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
from extra.assembly.amd.test.helpers import ExecContext
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
# Compile pseudocode functions on demand for regression tests
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', PCODE[VOP3SDOp.V_DIV_SCALE_F32])
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', PCODE[VOPCOp.V_CMP_CLASS_F32_E32])
class TestReg(unittest.TestCase):
def test_u32_read(self):
r = Reg(0xDEADBEEF)
self.assertEqual(int(r.u32), 0xDEADBEEF)
def test_u32_write(self):
r = Reg(0)
r.u32 = 0x12345678
self.assertEqual(r._val, 0x12345678)
def test_f32_read(self):
r = Reg(0x40400000) # 3.0f
self.assertAlmostEqual(float(r.f32), 3.0)
def test_f32_write(self):
r = Reg(0)
r.f32 = 3.0
self.assertEqual(r._val, 0x40400000)
def test_i32_signed(self):
r = Reg(0xFFFFFFFF) # -1 as signed
self.assertEqual(int(r.i32), -1)
def test_u64(self):
r = Reg(0xDEADBEEFCAFEBABE)
self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
def test_f64(self):
r = Reg(0x4008000000000000) # 3.0 as f64
self.assertAlmostEqual(float(r.f64), 3.0)
class TestTypedView(unittest.TestCase):
def test_bit_slice(self):
r = Reg(0xDEADBEEF)
# Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
self.assertEqual(r.u32[7:0].u32, 0xEF)
self.assertEqual(r.u32[15:8].u32, 0xBE)
self.assertEqual(r.u32[23:16].u32, 0xAD)
self.assertEqual(r.u32[31:24].u32, 0xDE)
# Also works with int() for arithmetic
self.assertEqual(int(r.u32[7:0]), 0xEF)
def test_single_bit_read(self):
r = Reg(0b11010101)
self.assertEqual(r.u32[0], 1)
self.assertEqual(r.u32[1], 0)
self.assertEqual(r.u32[2], 1)
self.assertEqual(r.u32[3], 0)
def test_single_bit_write(self):
r = Reg(0)
r.u32[5] = 1
r.u32[3] = 1
self.assertEqual(r._val, 0b00101000)
def test_nested_bit_access(self):
# S0.u32[S1.u32[4:0]] - access bit at position from another register
s0 = Reg(0b11010101)
s1 = Reg(3)
bit_pos = s1.u32[4:0] # TypedView, int value = 3
bit_val = s0.u32[int(bit_pos)] # bit 3 of s0 = 0
self.assertEqual(int(bit_pos), 3)
self.assertEqual(bit_val, 0)
def test_arithmetic(self):
r1 = Reg(0x40400000) # 3.0f
r2 = Reg(0x40800000) # 4.0f
result = r1.f32 + r2.f32
self.assertAlmostEqual(result, 7.0)
def test_comparison(self):
r1 = Reg(5)
r2 = Reg(3)
self.assertTrue(r1.u32 > r2.u32)
self.assertFalse(r1.u32 < r2.u32)
self.assertTrue(r1.u32 != r2.u32)
class TestTypedView(unittest.TestCase):
def test_slice_read(self):
r = Reg(0x56781234)
self.assertEqual(r[15:0].u16, 0x1234)
self.assertEqual(r[31:16].u16, 0x5678)
def test_slice_write(self):
r = Reg(0)
r[15:0].u16 = 0x1234
r[31:16].u16 = 0x5678
self.assertEqual(r._val, 0x56781234)
def test_slice_f16(self):
r = Reg(0)
r[15:0].f16 = 3.0
self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
class TestCompiler(unittest.TestCase):
def test_ternary(self):
result = _expr("a > b ? 1 : 0")
self.assertIn("if", result)
self.assertIn("else", result)
def test_type_prefix_strip(self):
self.assertEqual(_expr("1'0U"), "0")
self.assertEqual(_expr("32'1"), "1")
self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
def test_suffix_strip(self):
self.assertEqual(_expr("0ULL"), "0")
self.assertEqual(_expr("1LL"), "1")
self.assertEqual(_expr("5U"), "5")
self.assertEqual(_expr("3.14F"), "3.14")
def test_boolean_ops(self):
self.assertIn("and", _expr("a && b"))
self.assertIn("or", _expr("a || b"))
self.assertIn("!=", _expr("a <> b"))
def test_pack16(self):
result = _expr("{ a, b }")
self.assertIn("_pack", result)
def test_type_cast_strip(self):
self.assertEqual(_expr("64'U(x)"), "(x)")
self.assertEqual(_expr("32'I(y)"), "(y)")
class TestExecContext(unittest.TestCase):
def test_float_add(self):
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
def test_float_mul(self):
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
ctx.run("D0.f32 = S0.f32 * S1.f32")
self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
def test_scc_comparison(self):
ctx = ExecContext(s0=42, s1=42)
ctx.run("SCC = S0.u32 == S1.u32")
self.assertEqual(ctx.SCC._val, 1)
def test_scc_comparison_false(self):
ctx = ExecContext(s0=42, s1=43)
ctx.run("SCC = S0.u32 == S1.u32")
self.assertEqual(ctx.SCC._val, 0)
def test_ternary(self):
code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
ctx = ExecContext(s0=5, s1=3)
ctx.run(code)
self.assertEqual(ctx.D0._val, 1)
def test_pack(self):
code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
ctx = ExecContext(s0=0x1234, s1=0x5678)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0x56781234)
def test_tmp_with_typed_access(self):
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
self.assertEqual(ctx.D0._val, 300)
def test_s_add_u32_pattern(self):
# Real pseudocode pattern from S_ADD_U32
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
# Test overflow case
ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0) # Wraps to 0
self.assertEqual(ctx.SCC._val, 1) # Carry set
def test_s_add_u32_no_overflow(self):
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
self.assertEqual(ctx.D0._val, 300)
self.assertEqual(ctx.SCC._val, 0) # No carry
def test_vcc_lane_read(self):
ctx = ExecContext(vcc=0b1010, lane=1)
# Lane 1 is set
self.assertEqual(ctx.VCC.u64[1], 1)
self.assertEqual(ctx.VCC.u64[2], 0)
def test_vcc_lane_write(self):
ctx = ExecContext(vcc=0, lane=0)
ctx.VCC.u64[3] = 1
ctx.VCC.u64[1] = 1
self.assertEqual(ctx.VCC._val, 0b1010)
def test_for_loop(self):
# CTZ pattern - find first set bit
code = _compile_pseudocode("""tmp = -1
for i in 0 : 31 do
if S0.u32[i] == 1 then
tmp = i
endif
endfor
D0.i32 = tmp""")
ctx = ExecContext(s0=0b1000) # Bit 3 is set
ctx.run(code)
self.assertEqual(ctx.D0._val & MASK32, 3)
def test_result_dict(self):
ctx = ExecContext(s0=5, s1=3)
ctx.D0.u32 = 42
ctx.SCC._val = 1
result = ctx.result()
self.assertEqual(result['d0'], 42)
self.assertEqual(result['scc'], 1)
class TestPseudocodeRegressions(unittest.TestCase):
"""Regression tests for pseudocode instruction emulation bugs."""
def test_v_div_scale_f32_vcc_always_returned(self):
"""V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
This caused division to produce wrong results for multiple lanes."""
# Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
s0 = 0x3f800000 # 1.0
s1 = 0x40400000 # 3.0
s2 = 0x3f800000 # 1.0 (numerator)
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
# Must always have VCC in result
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
def test_v_cmp_class_f32_detects_quiet_nan(self):
"""V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
quiet_nan = 0x7fc00000 # quiet NaN: exponent=255, bit22=1
signal_nan = 0x7f800001 # signaling NaN: exponent=255, bit22=0
# Test quiet NaN detection (bit 1 in mask)
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
# Test signaling NaN detection (bit 0 in mask)
s1_signal = 0b0000000001 # bit 0 = signaling NaN
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
# Test that quiet NaN doesn't match signaling NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
# Test that signaling NaN doesn't match quiet NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
def testisNAN_with_typed_view(self):
"""isNAN must work with TypedView objects, not just Python floats.
Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
nan_reg = Reg(0x7fc00000) # quiet NaN
normal_reg = Reg(0x3f800000) # 1.0
inf_reg = Reg(0x7f800000) # +inf
self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
class TestBF16(unittest.TestCase):
"""Tests for BF16 (bfloat16) support."""
def test_bf16_conversion(self):
"""Test bf16 <-> f32 conversion."""
# bf16 is just the top 16 bits of f32
# 1.0f = 0x3f800000, bf16 = 0x3f80
self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
self.assertEqual(_ibf16(1.0), 0x3f80)
# 2.0f = 0x40000000, bf16 = 0x4000
self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
self.assertEqual(_ibf16(2.0), 0x4000)
# -1.0f = 0xbf800000, bf16 = 0xbf80
self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
self.assertEqual(_ibf16(-1.0), 0xbf80)
def test_bf16_special_values(self):
"""Test bf16 special values (inf, nan)."""
import math
# +inf: f32 = 0x7f800000, bf16 = 0x7f80
self.assertTrue(math.isinf(_bf16(0x7f80)))
self.assertEqual(_ibf16(float('inf')), 0x7f80)
# -inf: f32 = 0xff800000, bf16 = 0xff80
self.assertTrue(math.isinf(_bf16(0xff80)))
self.assertEqual(_ibf16(float('-inf')), 0xff80)
# NaN: quiet NaN bf16 = 0x7fc0
self.assertTrue(math.isnan(_bf16(0x7fc0)))
self.assertEqual(_ibf16(float('nan')), 0x7fc0)
def test_bf16_register_property(self):
"""Test Reg.bf16 property."""
r = Reg(0)
r.bf16 = 3.0 # 3.0f = 0x40400000, bf16 = 0x4040
self.assertEqual(r._val & 0xffff, 0x4040)
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
def test_bf16_slice_property(self):
"""Test TypedView.bf16 property."""
r = Reg(0x40404040) # Two bf16 3.0 values
self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
class TestBytePermute(unittest.TestCase):
"""Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
def test_byte_select_0_to_7(self):
"""Test selecting bytes 0-7 from 64-bit data."""
# data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
# Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
data = 0x0706050403020100
for i in range(8):
self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
def test_sign_extend_bytes(self):
"""Test sign extension selectors 8-11."""
# sel 8: sign of byte 1 (bits 15:8)
# sel 9: sign of byte 3 (bits 31:24)
# sel 10: sign of byte 5 (bits 47:40)
# sel 11: sign of byte 7 (bits 63:56)
data = 0x8000800080008000 # All relevant bytes have sign bit set
self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
data = 0x7f007f007f007f00 # No sign bits set
self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
def test_constant_zero(self):
"""Test selector 12 returns 0x00."""
self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
def test_constant_ff(self):
"""Test selectors >= 13 return 0xFF."""
for sel in [13, 14, 15, 255]:
self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
class TestSADHelpers(unittest.TestCase):
"""Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
def test_v_sad_u8_basic(self):
"""Test v_sad_u8 with simple values."""
# s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
result = v_sad_u8(0x04030201, 0x04030201, 0)
self.assertEqual(result, 0)
# s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
result = v_sad_u8(0x05040302, 0x04030201, 0)
self.assertEqual(result, 4)
def test_v_sad_u8_with_accumulator(self):
"""Test v_sad_u8 with non-zero accumulator."""
# s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
result = v_sad_u8(0x05040302, 0x04030201, 100)
self.assertEqual(result, 104)
def test_v_sad_u8_large_diff(self):
"""Test v_sad_u8 with maximum byte differences."""
# s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
result = v_sad_u8(0xffffffff, 0x00000000, 0)
self.assertEqual(result, 1020)
def test_v_msad_u8_basic(self):
"""Test v_msad_u8 masks when reference byte is 0."""
# s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
result = v_msad_u8(0x10101010, 0x00000000, 0)
self.assertEqual(result, 0)
# s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
result = v_msad_u8(0x10101010, 0x01010101, 0)
self.assertEqual(result, 60)
def test_v_msad_u8_partial_mask(self):
"""Test v_msad_u8 with partial masking."""
# s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
# diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
result = v_msad_u8(0x10101010, 0x00010001, 0)
self.assertEqual(result, 30)
def test_v_msad_u8_with_accumulator(self):
"""Test v_msad_u8 with non-zero accumulator."""
result = v_msad_u8(0x10101010, 0x01010101, 50)
self.assertEqual(result, 110) # 60 + 50
if __name__ == '__main__':
unittest.main()
+2 -2
View File
@@ -1,7 +1,7 @@
#!/usr/bin/env python3
"""Test PDF pseudocode extraction from amdxml.py."""
"""Test PDF pseudocode extraction from generate.py."""
import unittest
from extra.assembly.amd.amdxml import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
from extra.assembly.amd.generate import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
EXPECTED_PAGES = {"rdna3": 655, "rdna4": 711, "cdna": 610}
+98
View File
@@ -0,0 +1,98 @@
import unittest, ctypes
from extra.assembly.amd.autogen.rdna4 import ins as ir4
from extra.assembly.amd.dsl import v, s
from extra.assembly.amd.emu import WaveState, decode_program
from tinygrad.device import Buffer, BufferSpec
from tinygrad.dtype import dtypes
class TestRDNA4Emu(unittest.TestCase):
def _run(self, insts: list, sgprs: dict[int, int] = None, vgprs: dict[tuple[int, int], int] = None) -> WaveState:
"""Run instructions and return final WaveState."""
# Add S_ENDPGM if not present
if not any(isinstance(i, ir4.SOPP) and i.op == ir4.SOPPOp.S_ENDPGM for i in insts):
insts = list(insts) + [ir4.SOPP(ir4.SOPPOp.S_ENDPGM, simm=0)]
# Assemble and decode
code = b''.join(i.to_bytes() for i in insts)
code_buf = (ctypes.c_uint8 * len(code)).from_buffer_copy(code)
code_addr = ctypes.addressof(code_buf)
program_raw = decode_program(code, "rdna4")
program = {code_addr + offset: val for offset, val in program_raw.items()}
# Setup wave state
st = WaveState(n_lanes=1)
st.pc = code_addr
if sgprs:
for idx, val in sgprs.items(): st._write_sgpr(idx, val)
if vgprs:
for (reg, lane), val in vgprs.items(): st._write_vgpr(reg, lane, val)
# Setup vmem buffer with external_ptr=0 (maps to address 0, allows any pointer access)
vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
# Execute
c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(0), ctypes.c_uint64(0)]
for _ in range(100):
if (pc := st.pc) == 0xFFFFFFFFFFFFFFFF or pc not in program: break
_, fxn, globals_list, _ = program[pc]
fxn(*[c_bufs[g] for g in globals_list])
return st
def test_vopd_dual_mov(self):
"""Test VOPD with two V_DUAL_MOV_B32 operations: v[1]=s[1], v[2]=s[2]."""
insts = [ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0])]
st = self._run(insts, sgprs={1: 0x40e00000, 2: 0x41100000}) # 7.0f, 9.0f
self.assertEqual(st._read_vgpr(1, 0), 0x40e00000) # v[1] = 7.0
self.assertEqual(st._read_vgpr(2, 0), 0x41100000) # v[2] = 9.0
def test_vopd_dual_mov_after_other_vopd(self):
"""Test VOPD reuse: first VOPD(v[3]=0, v[0]=?), then VOPD(v[1]=s[1], v[2]=s[2])."""
# This matches the BEAM kernel sequence that fails
insts = [
ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]), # v[3]=0, v[0]=s[0]
ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]), # v[1]=s[1], v[2]=s[2]
]
st = self._run(insts, sgprs={0: 0x40a00000, 1: 0x40e00000, 2: 0x41100000}) # 5.0f, 7.0f, 9.0f
self.assertEqual(st._read_vgpr(1, 0), 0x40e00000) # v[1] = 7.0
self.assertEqual(st._read_vgpr(2, 0), 0x41100000) # v[2] = 9.0
def test_vopd_with_s_add_f32_sequence(self):
"""Test full BEAM kernel sequence: s_add_f32 then VOPD."""
# This is the exact sequence from the failing BEAM kernel
insts = [
ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[0], ssrc0=s[0], ssrc1=s[8]), # s[0] = s[0] + s[8]
ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[1], ssrc0=s[1], ssrc1=s[9]), # s[1] = s[1] + s[9]
ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[2], ssrc0=s[2], ssrc1=s[10]), # s[2] = s[2] + s[10]
ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),
ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
]
# Input: s[0:2] = [1,2,3], s[8:10] = [4,5,6]
# After s_add_f32: s[0:2] = [5,7,9]
st = self._run(insts, sgprs={0: 0x3f800000, 1: 0x40000000, 2: 0x40400000, # 1.0, 2.0, 3.0
8: 0x40800000, 9: 0x40a00000, 10: 0x40c00000}) # 4.0, 5.0, 6.0
self.assertEqual(st._read_vgpr(1, 0), 0x40e00000) # v[1] = 7.0
self.assertEqual(st._read_vgpr(2, 0), 0x41100000) # v[2] = 9.0
def test_s_mov_b32_then_vopd(self):
"""Test s_mov_b32 followed by VOPD - simulates BEAM kernel sequence."""
# Use s_mov_b32 with SGPR source (copy from pre-initialized SGPRs)
# s[10:12] will have values set by test harness, copy to s[0:2], then VOPD to VGPRs
insts = [
ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[0], ssrc0=s[10]), # s[0] = s[10]
ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[1], ssrc0=s[11]), # s[1] = s[11]
ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[2], ssrc0=s[12]), # s[2] = s[12]
ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
]
st = self._run(insts, sgprs={10: 0x40a00000, 11: 0x40e00000, 12: 0x41100000}) # 5.0, 7.0, 9.0
self.assertEqual(st._read_vgpr(1, 0), 0x40e00000) # v[1] = 7.0
self.assertEqual(st._read_vgpr(2, 0), 0x41100000) # v[2] = 9.0
if __name__ == '__main__':
unittest.main()
+9 -6
View File
@@ -2,7 +2,7 @@
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re, subprocess, os
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump, get_target, get_mattr
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
@@ -83,17 +83,21 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
arch = self.arch
from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
from tinygrad.helpers import AMD_LLVM
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
compiler = HIPCompiler(get_target(arch))
# rendered source can be C or llvmir
compiler = (AMDLLVMCompiler if AMD_LLVM else HIPCompiler)(get_target(arch))
# First pass: decode all instructions and collect info
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
for ki, kernel in enumerate(kernels):
offset = 0
while offset < len(kernel.code):
remaining = kernel.code[offset:]
code = next((s.content for s in elf_loader(compiler.compile(kernel.src))[1] if s.name == ".text"))
while offset < len(code):
remaining = code[offset:]
fmt = detect_format(remaining, arch)
if fmt is None:
decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
@@ -234,7 +238,6 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
@unittest.skip("RDNA4 decode roundtrip not yet supported")
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
@unittest.skip("CDNA decode roundtrip not yet supported")
+34 -19
View File
@@ -5,11 +5,11 @@ from pathlib import Path
from tinygrad.helpers import DEBUG
from tinygrad.runtime.autogen import rocprof
from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES, InstOp, print_packets)
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_L4, WAVEEND, INST, INST_L4, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES_L3, PACKET_TYPES_L4, InstOp, InstOpL4, print_packets)
from extra.assembly.amd.test.helpers import TARGET_TO_ARCH
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
@@ -19,6 +19,7 @@ OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LD
InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
InstOp.OTHER_GLOBAL_STORE_VADDR_128}
OTHER_SIMD_OPS_L4 = {InstOpL4.OTHER_VMEM, InstOpL4.UNK_60}
# ═══════════════════════════════════════════════════════════════════════════════
# ROCPROF DECODER
@@ -91,11 +92,12 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
return occupancy_records, wave_insts
class TestSQTTExamples(unittest.TestCase):
target = "gfx1100"
class SQTTExamplesTestBase(unittest.TestCase):
target: str
@classmethod
def setUpClass(cls):
if cls is SQTTExamplesTestBase: raise unittest.SkipTest("base class")
cls.examples = {}
for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
with open(pkl_path, "rb") as f:
@@ -118,18 +120,20 @@ class TestSQTTExamples(unittest.TestCase):
self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
def test_packet_types_valid(self):
all_classes = set(PACKET_TYPES_L3.values()) | set(PACKET_TYPES_L4.values())
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
for pkt in decode(event.blob):
self.assertIn(type(pkt), PACKET_TYPES, f"unknown packet type {type(pkt)} in {name}")
# Use isinstance to handle layout-specific subclasses (e.g., WAVESTART_L4)
self.assertTrue(any(isinstance(pkt, cls) for cls in all_classes), 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, (WAVESTART, WAVESTART_L4))]), 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):
@@ -144,17 +148,11 @@ class TestSQTTExamples(unittest.TestCase):
if "gemm" not in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, INST)]), 0, f"no INST packets in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_L4))]), 0, f"no INST packets in {name}")
expected = {
"profile_empty_run_0": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_empty_run_1": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_gemm_run_0": [2531, 1844, 1864, 1915, 1942, 1848, 3074, 1919, 1939, 1990, 2017, 1923, 19026, 1919, 1939, 1990, 2017, 1929],
"profile_gemm_run_1": [2554, 1844, 1864, 1915, 1942, 1848, 3084, 1919, 1939, 1990, 2017, 1923, 19010, 1919, 1939, 1990, 2017, 1923],
"profile_plus_run_0": [1900, 1908, 1928, 1979, 2006, 1912],
"profile_plus_run_1": [1856, 1908, 1928, 1979, 2006, 1912],
}
expected: dict[str, list[int]] = {} # override in subclasses
def test_packet_counts(self):
if not self.expected: self.skipTest("no expected packet counts for this target")
for name, (events, *_) in self.examples.items():
with self.subTest(example=name):
if not self.expected.get(name): continue
@@ -178,7 +176,7 @@ class TestSQTTExamples(unittest.TestCase):
for event in events:
wave_starts: dict[tuple[int, int, int], int] = {}
for p in decode(event.blob):
if isinstance(p, WAVESTART): wave_starts[(p.wave, p.simd, p.cu)] = p._time
if isinstance(p, (WAVESTART, WAVESTART_L4)): 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}")
@@ -195,15 +193,32 @@ class TestSQTTExamples(unittest.TestCase):
for event in events:
for p in decode(event.blob):
if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
elif isinstance(p, INST_L4) and p.op not in OTHER_SIMD_OPS_L4: our_insts.append(p._time)
elif isinstance(p, VALUINST): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE_MASK):
for _ in range(bin(p.mask).count('1')): our_insts.append(p._time)
self.assertEqual(sorted(our_insts), sorted(roc_insts), f"instruction times mismatch in {name}")
#class TestSQTTExamplesRDNA4(TestSQTTExamples): target = "gfx1200"
class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
target = "gfx1100"
expected = {
"profile_empty_run_0": [1844, 1885, 1905, 1956, 1983, 1889],
"profile_empty_run_1": [1780, 1885, 1905, 1956, 1983, 1889],
"profile_gemm_run_0": [2656, 2025, 2045, 2096, 2123, 2029, 3183, 2019, 2039, 2090, 2117, 2023, 19119, 2013, 2033, 2084, 2111, 2017],
"profile_gemm_run_1": [2662, 2025, 2045, 2096, 2123, 2029, 3179, 2019, 2039, 2090, 2117, 2023, 19113, 2071, 2091, 2142, 2169, 2075],
"profile_plus_run_0": [1886, 2013, 2033, 2084, 2111, 2017],
"profile_plus_run_1": [1988, 2071, 2091, 2142, 2169, 2075],
}
#class TestSQTTExamplesCDNA(TestSQTTExamples): target = "gfx950"
class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
# CDNA/MI300 (gfx950) uses a completely different 16-bit header packet format, not the nibble-based format.
# See decode_tt_header_stream in ghidra/librocprof-trace-decoder.c - it reads 16-bit headers and uses
# pkt_fmt = header & 0xf to look up packet_class (0x10=2bytes, 0x20=4bytes, 0x30=6bytes, 0x40=8bytes).
# This is NOT implemented yet - the nibble decoder produces garbage for CDNA data.
@unittest.skip("CDNA/MI300 uses 16-bit header format, not nibble-based - decoder not implemented")
class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
target = "gfx950"
if __name__ == "__main__":
unittest.main()
+170
View File
@@ -0,0 +1,170 @@
"""Tests comparing sqtt.py PACKET_TYPES_L3/L4 against AMD's rocprof-trace-decoder binary."""
import unittest, struct, ctypes, pickle
from pathlib import Path
ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
def _find_segment(perms: str):
"""Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
with open('/proc/self/maps', 'r') as f:
for line in f:
if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
parts = line.split()
return int(parts[0].split('-')[0], 16), int(parts[2], 16)
return None, None
def _read_array(file_offset: int, count: int):
"""Read an array of uint8 at file_offset from the loaded library."""
base, seg_offset = _find_segment('rw-p')
if base is None: return None
return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
def _load_lib():
if not ROCPROF_LIB.exists(): return False
ctypes.CDLL(str(ROCPROF_LIB))
return True
# ═══════════════════════════════════════════════════════════════════════════════
# RDNA EXTRACTION (nibble-based format)
# ═══════════════════════════════════════════════════════════════════════════════
def extract_bit_tables():
"""Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
if not _load_lib(): return None
return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
def extract_delta_fields():
"""Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
if not _load_lib(): return None
ro_base, ro_offset = _find_segment('r--p')
if ro_base is None: return None
def read_table(file_offset, num_entries):
addr = ro_base + (file_offset - ro_offset)
data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
return {type_id: (lo, hi) for j in range(0, len(data), 12)
for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
def extract_packet_encodings():
"""Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
if not _load_lib(): return None
rw_base, rw_offset = _find_segment('rw-p')
if rw_base is None: return None
# Read base encodings from registration vector at 0x2d340
vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
base = {}
if vec_start and vec_end:
for i in range((vec_end - vec_start) // 32):
addr = vec_start + i * 32
type_id = ctypes.c_uint8.from_address(addr).value
pat_start = ctypes.c_void_p.from_address(addr + 8).value
pat_end = ctypes.c_void_p.from_address(addr + 16).value
if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
pat = list((ctypes.c_uint8 * n).from_address(pat_start))
base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base} # L2 has overrides
# ═══════════════════════════════════════════════════════════════════════════════
# CDNA EXTRACTION (16-bit header format)
# ═══════════════════════════════════════════════════════════════════════════════
def extract_cdna_packet_sizes():
"""Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
from extra.assembly.amd.test.test_sqtt_examples import run_rocprof_decoder
if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
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 not sqtt_events or not prg: return None
# Run decoder to trigger hash table initialization
run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
# Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
rw_base, rw_offset = _find_segment('rw-p')
if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
pkt_sizes, node, seen = {}, head, set()
while node and node not in seen and len(pkt_sizes) < 20:
seen.add(node)
key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
node = ctypes.c_void_p.from_address(node).value
return pkt_sizes if len(pkt_sizes) == 16 else None
# ═══════════════════════════════════════════════════════════════════════════════
# TESTS
# ═══════════════════════════════════════════════════════════════════════════════
class TestSQTTMatchesBinary(unittest.TestCase):
def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
def test_encodings_match_layout3(self): self._test_encodings(3)
def test_encodings_match_layout4(self): self._test_encodings(4)
def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
def test_cdna_packet_sizes(self):
"""Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
pkt_sizes = extract_cdna_packet_sizes()
self.assertIsNotNone(pkt_sizes, "failed to extract CDNA packet sizes")
from extra.assembly.amd.sqtt_cdna import CDNA_PKT_SIZES
for pkt_fmt, size in CDNA_PKT_SIZES.items():
with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
def _test_bit_counts(self, layout: int):
if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
with self.subTest(packet=pkt_cls.__name__):
self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])
def _test_encodings(self, layout: int):
if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
with self.subTest(packet=pkt_cls.__name__):
self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
def _test_delta_fields(self, layout: int):
if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
if type_id not in deltas[layout - 2]: continue
delta = getattr(pkt_cls, 'delta', None)
actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
if __name__ == "__main__":
tables = extract_bit_tables()
encodings = extract_packet_encodings()
deltas = extract_delta_fields()
TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
if encodings and tables:
print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12} {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
print("-" * 140)
for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12} {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
cdna = extract_cdna_packet_sizes()
if cdna: print(f"\nCDNA packet sizes: {cdna}")
unittest.main()
+1 -2
View File
@@ -67,12 +67,11 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
if not wasm:
thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
for name,cl in bufs_to_save.items():
weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
return '\n'.join(headers + cprog)
else:
if bufs_to_save:
+1 -1
View File
@@ -28,7 +28,7 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
_, input_uop, weight_uop = kernel.src
_, input_uop, weight_uop = kernel.src[1:]
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)
+22 -21
View File
@@ -168,12 +168,14 @@ PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR,
# =============================================================================
class Kernel:
def __init__(self, arch='gfx1100'):
self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
def label(self, name): self.labels[name] = self.pos
def 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 emit(self, inst, target=None):
self.instructions.append(inst)
inst._target, inst._pos = target, self.pos
self.pos += inst.size()
return inst
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."""
@@ -182,16 +184,15 @@ class Kernel:
self.emit(s_waitcnt(simm16=waitcnt))
def to_asm(self):
import re
# Instruction stream with labels
label_at = {pos: name for name, pos in self.labels.items()}
body = []
for i, inst in enumerate(self.instructions):
if i in label_at: body.append(f'.{label_at[i]}:')
asm = inst.disasm()
if i in self.branch_targets:
asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
body.append('\t' + asm)
# Patch branch offsets: simm16 = (target_pos - branch_end_pos) / 4
for inst in self.instructions:
if inst._target is None: continue
offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
inst.simm16 = offset_dwords
# TODO: replace this with direct ELF
body = ['\t' + inst.disasm() for inst in self.instructions]
# limit wave occupancy by using more LDS
lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
@@ -315,7 +316,7 @@ def build_kernel(arch='gfx1100'):
k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
# S_LOOP_CTR is already 0 from prologue initialization
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_ENTRY')
k.emit(s_branch(), target='LOOP_ENTRY')
# ===========================================================================
# MAIN GEMM LOOP
@@ -326,12 +327,12 @@ def build_kernel(arch='gfx1100'):
k.label('LOOP_INC')
k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
k.emit(s_cbranch_scc1(simm16=0)); k.branch_to('EPILOGUE')
k.emit(s_cbranch_scc1(), target='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
k.emit(s_cbranch_scc0(), target='SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
if not NO_GLOBAL:
# Advance prefetch pointers (VGPR)
@@ -402,7 +403,7 @@ def build_kernel(arch='gfx1100'):
offset = i * 64
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_INC')
k.emit(s_branch(), target='LOOP_INC')
# ===========================================================================
# EPILOGUE: Permute and store results
@@ -470,7 +471,7 @@ THREADS = 128
def test_matmul():
dev = Device[Device.DEFAULT]
print(f"Device arch: {dev.arch}")
print(f"Device arch: {dev.renderer.arch}")
if getenv("STOCK", 0):
# Load the stock kernel from amd_seb/kernel8_batched_gmem.s
@@ -478,7 +479,7 @@ def test_matmul():
asm = stock_path.read_text()
print(f"Loaded stock kernel from {stock_path}")
else:
asm = build_kernel(dev.arch)
asm = build_kernel(dev.renderer.arch)
binary = dev.compiler.compile(asm)
print(f"Compiled! Binary size: {len(binary)} bytes")
File diff suppressed because it is too large Load Diff
+100
View File
@@ -0,0 +1,100 @@
import atexit, functools
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad import Tensor, Device, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.renderer import Estimates
from tinygrad.helpers import getenv, all_same, dedup
from extra.gemm.asm.cdna.asm import build_kernel, GEMM_ARGS
# ** CDNA4 assembly gemm
WORKGROUP_SIZE = 256
def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
batch, M, K = A.shape
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2
lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
gidx = UOp.special(wg, "gidx0")
k = build_kernel(batch, M, N, K, A.dtype.base)
sink = UOp.sink(C.base, A.base, B.base, lidx, gidx,
arg=KernelInfo(name=k.name, estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
binary = HIPCompiler(arch).compile(k.to_asm())
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=k.to_text()), UOp(Ops.BINARY, arg=binary)))
counters = {"used":0, "todos":[]}
def todo(msg:str) -> bool: counters["todos"].append(msg); return False
atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
N = b.shape[1]
# only sharding on the batch or K is tested, others might work too
if isinstance(a.device, tuple):
if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
dname = a.device[0]
else: dname = a.device
arch = getattr(Device[dname].renderer, "arch", "")
if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
if (key:=(M, N, K)) not in GEMM_ARGS and arch == "gfx950": return todo(f"GEMM shape not supported {key} on {arch}")
return True
# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
# note: this can be removed after we have GEMM on mixins
def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
M, K = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2
m = UOp.range(M, 1, AxisType.LOOP)
n = UOp.range(N, 2, AxisType.LOOP)
k = UOp.range(K, 0, AxisType.REDUCE)
mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
# ** backward gemm, might use the asm gemm
def custom_gemm_bw(gradient:UOp, kernel:UOp):
out, a, b = kernel.src[1:]
assert all_same([gradient.device, a.device, b.device, out.device])
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
# TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
g_t = g_t[:a.shape[0]]
grad_a = (g_t @ b_t.T).uop
grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
return (None, grad_a, grad_b)
# ** main gemm function
def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
counters["used"] += 1
squeeze = a.ndim == 2
if squeeze: a = a.unsqueeze(0)
batch, M, K = a.shape
N = b.shape[1]
is_multi = isinstance(a.device, tuple)
if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
if is_multi:
out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
else:
out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
dname = a.device[0] if is_multi else a.device
arch = getattr(Device[dname].renderer, "arch", "")
if arch.startswith("gfx950") and getenv("USE_ASM", 1):
numWG = GEMM_ARGS[(M, N, K)][0]
out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
else:
out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
if k_sharded: out = out.sum(0)
return out.squeeze(0) if squeeze else out
File diff suppressed because it is too large Load Diff
-78
View File
@@ -1,78 +0,0 @@
.text
.section .text.
.global gemm
.p2align 8
.type gemm,@function
gemm:
INSTRUCTIONS
.section .rodata,"a",@progbits
.p2align 6, 0x0
.amdhsa_kernel gemm
# basic memory requirements
.amdhsa_group_segment_fixed_size 133120
.amdhsa_private_segment_fixed_size 0
.amdhsa_kernarg_size 28
# register usage (RSRC1)
.amdhsa_next_free_vgpr 504
.amdhsa_next_free_sgpr 96
# workgroup / workitem IDs (RSRC2)
.amdhsa_system_sgpr_workgroup_id_x 1
.amdhsa_system_sgpr_workgroup_id_y 1
.amdhsa_system_sgpr_workgroup_id_z 1
# user SGPRs, we only specify the kernel args ptr in s[0:1]
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_user_sgpr_count 2
.amdhsa_user_sgpr_kernarg_preload_length 0
.amdhsa_user_sgpr_kernarg_preload_offset 0
# gfx90a / gfx940 specifics (RSRC3)
.amdhsa_accum_offset 248
.amdhsa_uses_dynamic_stack 0
.amdhsa_tg_split 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.kernels:
- .name: gemm
.symbol: gemm.kd
.args:
- .name: C
.address_space: global
.offset: 0
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: B
.address_space: global
.offset: 8
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: A
.address_space: global
.offset: 16
.size: 8
.value_kind: global_buffer
.value_type: bf16
- .name: sz
.offset: 24
.size: 4
.value_kind: by_value
.value_type: u32
.group_segment_fixed_size: 133120
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.kernarg_segment_size: 28
.max_flat_workgroup_size: 256
.sgpr_count: 88
.sgpr_spill_count: 0
.vgpr_count: 248
.vgpr_spill_count: 0
.wavefront_size: 64
amdhsa.version:
- 1
- 0
...
.end_amdgpu_metadata
-72
View File
@@ -1,72 +0,0 @@
# Run assembly on the AMD runtime and check correctness
# VIZ=2 to profile
import pathlib
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.helpers import getenv
fp = pathlib.Path(__file__).parent/"gemm.s"
N = getenv("N", 8192)
THREADS_PER_WG = 256
NUM_WG = N//THREADS_PER_WG * N//THREADS_PER_WG
assert N % THREADS_PER_WG == 0, "N must be divisible by THREADS_PER_WG"
# ** generate inputs on CPU
scale = 10.0
import torch
torch.manual_seed(0)
A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
Bt = B.t().contiguous() # transpose B for the asm gemm
C_torch = A@B
# ** copy buffers to AMD
# input creation and validation run on the copy engine for simpler tracing
def from_torch(t:torch.Tensor) -> Tensor:
return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
C_tiny = from_torch(A) @ from_torch(B)
C_asm = Tensor.empty_like(C_tiny)
# ** assembly custom kernel
def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
lidx = UOp.special(THREADS_PER_WG, "lidx0")
gidx = UOp.special(NUM_WG, "gidx0")
src = (pathlib.Path(__file__).parent/"template.s").read_text().replace("INSTRUCTIONS", fp.read_text())
sz = UOp.variable("SZ", 256, 8192)
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
# ** run gemms
sched = Tensor.schedule(C_tiny, C_asm)
eis = [si.lower() for si in sched]
with Context(DEBUG=2):
for ei in eis:
et = ei.run({"SZ":N}, wait=True)
print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
# ** correctness
import ctypes
def torch_bf16(t:Tensor) -> torch.tensor:
asm_out = t.to("cpu").realize().uop.buffer._buf
buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)
+1 -1
View File
@@ -37,7 +37,7 @@ b.copyin(row.data)
c.copyin(mat.data)
ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
print(ret)
out = np.frombuffer(a.as_buffer(), np.float32)
out = np.frombuffer(a.as_memoryview(), np.float32)
real = row.astype(np.float32)@mat.T.astype(np.float32)
print("out:", out)
print("real", real)
+2 -2
View File
@@ -98,10 +98,10 @@ if __name__ == "__main__":
# check correctness
if getenv("VERIFY"):
from tinygrad.engine.realize import run_schedule
triton_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
print(triton_buf)
run_schedule(sched)
tinygrad_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
print(tinygrad_buf)
np.testing.assert_allclose(triton_buf, tinygrad_buf)
print("correct!")
+6 -14
View File
@@ -1,14 +1,15 @@
#!/usr/bin/env python3
import argparse, glob, os, time, subprocess, sys
from tinygrad.helpers import temp
def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
devs = []
for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
dev_id = dev[8:-5]
if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
for dev in glob.glob(temp(f'{prefix}_*.lock')):
dev_id = dev.split('/')[-1][len(prefix)+1:-5]
if dev_id.startswith(target_dev): devs.append(dev_id)
return devs
def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
@@ -53,16 +54,7 @@ def cmd_show_pids(args):
for dev in devs:
try:
pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
print(f"{dev}: {pid}")
except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
def cmd_kill_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
for dev in devs:
try:
pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
print(f"{dev}: {pid}")
except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
@@ -74,7 +66,7 @@ def cmd_kill_pids(args):
if i > 0: time.sleep(0.2)
try:
try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
except subprocess.CalledProcessError: break
print(f"Killing process {pid} (which uses {dev})")
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
import numpy as np
def to_np(buf): return np.frombuffer(buf.as_buffer().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
big = to_np(buf0)
print(big)
@@ -8,14 +8,14 @@ from tinygrad.helpers import _ensure_downloads_dir
DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
from tinygrad.helpers import tqdm
def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, local_dir: str|Path|None=None,
tries: int=2, **kwargs) -> Path:
for attempt in range(tries):
try:
return Path(snapshot_download(
repo_id=repo_id,
allow_patterns=allow_patterns,
cache_dir=str(cache_dir) if cache_dir is not None else None,
local_dir=str(local_dir) if local_dir is not None else None,
**kwargs
))
except Exception as e:
@@ -144,14 +144,14 @@ class HuggingFaceONNXManager:
root_path = snapshot_download_with_retry(
repo_id=model_id,
allow_patterns=allow_patterns,
cache_dir=str(self.models_dir)
local_dir=str(self.models_dir / model_id)
)
# Download config files (usually small)
snapshot_download_with_retry(
repo_id=model_id,
allow_patterns=["*config.json"],
cache_dir=str(self.models_dir)
local_dir=str(self.models_dir / model_id)
)
model_data["download_path"] = str(root_path)
+4 -4
View File
@@ -88,8 +88,8 @@ if __name__ == "__main__":
# repo id
# validates all onnx models inside repo
repo_id = "/".join(path)
root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], local_dir=DOWNLOADS_DIR / repo_id)
snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
config = get_config(root_path)
for onnx_model in root_path.rglob("*.onnx"):
rtol, atol = get_tolerances(onnx_model.name)
@@ -101,8 +101,8 @@ if __name__ == "__main__":
onnx_model = path[-1]
assert path[-1].endswith(".onnx")
repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], local_dir=DOWNLOADS_DIR / repo_id)
snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
config = get_config(root_path)
rtol, atol = get_tolerances(onnx_model)
print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")
+70 -62
View File
@@ -5,7 +5,7 @@ os.environ["AMD_AQL"] = "1"
from tinygrad.device import Device
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.runtime.ops_amd import AMDProgram
from extra.assembly.amd.dsl import Reg, Inst, s, v
NUM_WORKGROUPS = 96
WAVE_SIZE = 32
@@ -17,84 +17,92 @@ DIRECTIVE = ".amdhsa_wavefront_size32 1"
assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> bytes:
preamble = s_mov_b32(counter_sreg, n).to_bytes()
insts_bytes = b"".join([inst.to_bytes() for inst in insts])
sub_inst, cmp_inst = s_sub_u32(counter_sreg, counter_sreg, 1), s_cmp_lg_i32(counter_sreg, 0)
loop_sz = len(insts_bytes) + sub_inst.size() + cmp_inst.size()
branch_inst = s_cbranch_scc1(simm16=-((loop_sz // 4) + 1) & 0xFFFF)
return preamble + insts_bytes + sub_inst.to_bytes() + cmp_inst.to_bytes() + branch_inst.to_bytes() + s_endpgm().to_bytes()
def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs):
if accum:
instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
vgprIndices[1], vgprIndices[2],
vgprIndices[1], vgprIndices[2], extra)
inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1, acc_cd=1, **kwargs)
elif dense:
instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
vgprIndices[1], vgprIndices[2],
vgprIndices[1], vgprIndices[2])
inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1)
else:
instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
vgprIndices[1], vgprIndices[2],
vgprIndices[3], vgprIndices[4],
vgprIndices[5])
src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[3]:vgprIndices[4]], v[vgprIndices[5]])
vgprs:set = set()
for n,_ in inst._fields:
if isinstance(val:=getattr(inst, n), Reg) and val.offset >= v.offset: vgprs |= {val.offset+i for i in range(val.sz)}
inst_bytes = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
inst_hex = "\n".join(" .byte " + ",".join(f"0x{b:02x}" for b in inst_bytes[i:i+16]) for i in range(0, len(inst_bytes), 16)) + "\n"
src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", inst_hex).replace("VGPR_COUNT", str(len(vgprs)))
src = src.replace("DIRECTIVE", DIRECTIVE)
lib = COMPILER.compile(src)
fxn = AMDProgram(DEV, "matmul", lib)
fxn = DEV.runtime("matmul", lib)
elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
if __name__=="__main__":
DEVICENUM = os.getenv("DEVICENUM", "0")
try:
DEV = Device['AMD:' + DEVICENUM]
except:
raise RuntimeError("Error while initiating AMD device")
DEV = Device[Device.DEFAULT]
arch = DEV.renderer.arch
COMPILER = HIPCompiler(DEV.arch)
if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 32
launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
elif DEV.arch == 'gfx1201':
COMPILER = HIPCompiler(arch)
if arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
from extra.assembly.amd.autogen.rdna3.ins import *
if arch == 'gfx1103': NUM_WORKGROUPS = 8
if arch == 'gfx1151': NUM_WORKGROUPS = 32
launchBenchmark(v_wmma_bf16_16x16x16_bf16, (7,8,15))
launchBenchmark(v_wmma_f16_16x16x16_f16, (7,8,15))
launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,15))
launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,15))
launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,9))
launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,11))
elif arch in {'gfx1200', 'gfx1201'}:
from extra.assembly.amd.autogen.rdna4.ins import *
# this instruction does not exist in the rdna4 isa, use the co version
s_sub_u32 = s_sub_co_u32
NUM_WORKGROUPS = 64
launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,11))
launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,11))
launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,8))
launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_fp8_fp8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_fp8_bf8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_bf8_fp8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_bf8_bf8", (7,8,9))
launchBenchmark(v_wmma_bf16_16x16x16_bf16, (3,4,7))
launchBenchmark(v_wmma_f16_16x16x16_f16, (3,4,7))
launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,11))
launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,11))
launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,8))
launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_fp8_fp8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_fp8_bf8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_bf8_fp8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_bf8_bf8, (7,8,9))
FLOPS_PER_MATMUL = 16*16*32*2
launchBenchmark("v_wmma_i32_16X16X32_iu4", (7,8,9))
launchBenchmark("v_swmmac_f32_16x16x32_f16", (7,8,11,12,19,20), False)
launchBenchmark("v_swmmac_f32_16x16x32_bf16", (7,8,11,12,19,20), False)
launchBenchmark("v_swmmac_f16_16x16x32_f16", (3,4,7,8,15,16), False)
launchBenchmark("v_swmmac_bf16_16x16x32_bf16", (3,4,7,8,15,16), False)
launchBenchmark("v_swmmac_i32_16x16x32_iu8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_i32_16x16x32_iu4", (7,8,8,9,10,11), False)
launchBenchmark("v_swmmac_f32_16x16x32_fp8_fp8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_f32_16x16x32_fp8_bf8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
launchBenchmark(v_wmma_i32_16x16x32_iu4, (7,8,9))
launchBenchmark(v_swmmac_f32_16x16x32_f16, (7,8,11,12,19,20), False)
launchBenchmark(v_swmmac_f32_16x16x32_bf16, (7,8,11,12,19,20), False)
launchBenchmark(v_swmmac_f16_16x16x32_f16, (3,4,7,8,15,16), False)
launchBenchmark(v_swmmac_bf16_16x16x32_bf16, (3,4,7,8,15,16), False)
launchBenchmark(v_swmmac_i32_16x16x32_iu8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_i32_16x16x32_iu4, (7,8,8,9,10,11), False)
launchBenchmark(v_swmmac_f32_16x16x32_fp8_fp8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_f32_16x16x32_fp8_bf8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_f32_16x16x32_bf8_fp8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_f32_16x16x32_bf8_bf8, (7,8,9,10,13,14), False)
FLOPS_PER_MATMUL = 16*16*64*2
launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
elif DEV.arch == 'gfx950':
launchBenchmark(v_swmmac_i32_16x16x64_iu4, (7,8,9,10,13,14), False)
elif arch == 'gfx950':
from extra.assembly.amd.autogen.cdna.ins import *
DIRECTIVE = ".amdhsa_accum_offset 4"
NUM_WORKGROUPS = 256
WAVE_SIZE = 64
NUM_WAVES = 4
launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
launchBenchmark(v_mfma_f32_16x16x16_f16, (3,0,1), accum=True)
launchBenchmark(v_mfma_f32_16x16x16_bf16, (3,0,1), accum=True)
FLOPS_PER_MATMUL = 16*16*32*2
launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
launchBenchmark(v_mfma_f32_16x16x32_f16, (3,0,3), accum=True)
launchBenchmark(v_mfma_f32_16x16x32_bf16, (3,0,3), accum=True)
FLOPS_PER_MATMUL = 16*16*128*2
launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,7), accum=True) # fp8
launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,5), accum=True, cbsz=2, blgp=2) # fp6
launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,3), accum=True, cbsz=4, blgp=4) # fp4
else:
raise RuntimeError(f"arch {DEV.arch} not supported.")
raise RuntimeError(f"arch {arch} not supported.")
+3 -10
View File
@@ -3,20 +3,13 @@
.p2align 8
.type matmul,@function
matmul:
s_mov_b32 s1, INTERNAL_LOOP
s_mov_b32 s2, 0
inner_loop:
INSTRUCTION
s_sub_u32 s1, s1, 1
s_cmp_lg_i32 s1, s2
s_cbranch_scc1 inner_loop
s_endpgm
.rodata
.p2align 6
.amdhsa_kernel matmul
.amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
.amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
.amdhsa_next_free_vgpr VGPR_COUNT
.amdhsa_next_free_sgpr 3
DIRECTIVE
.end_amdhsa_kernel
@@ -37,4 +30,4 @@ amdhsa.kernels:
.vgpr_count: 32
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
.end_amdgpu_metadata
+2 -2
View File
@@ -5,7 +5,6 @@ from tinygrad.nn import Linear, LayerNorm, Embedding, Conv2d
from typing import List, Optional, Union, Tuple, Dict
from abc import ABC, abstractmethod
from functools import lru_cache
from PIL import Image
import numpy as np
import re, gzip
@@ -444,7 +443,8 @@ class OpenClipEncoder:
# TODO:
# Should be doable in pure tinygrad, would just require some work and verification.
# This is very desirable since it would allow for full generation->evaluation in a single JIT call.
def prepare_image(self, image:Image.Image) -> Tensor:
def prepare_image(self, image) -> Tensor:
from PIL import Image
SIZE = 224
w, h = image.size
scale = min(SIZE / h, SIZE / w)
+17 -9
View File
@@ -55,12 +55,15 @@ class Attention:
xqkv = x @ self.wqkv.T
xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
else:
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
if self.q_norm is not None and self.k_norm is not None:
xq = self.q_norm(xq)
xk = self.k_norm(xk)
# cast_float_to_bf16 is expensive in reduction loops, break it out
if x.dtype == dtypes.bfloat16: xq, xk = xq.contiguous_backward(), xk.contiguous_backward()
xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
@@ -86,20 +89,23 @@ class Attention:
assert start_pos == 0
keys, values = xk, xv
keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
if self.max_context:
keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
else:
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
if getenv("STUB_ATTENTION"):
# TODO: do we need mask?
from tinygrad.uop.ops import UOp, KernelInfo
def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
grad_q = Tensor.empty_like(q:=Tensor(kernel.src[2]))
grad_k = Tensor.empty_like(k:=Tensor(kernel.src[3]))
grad_v = Tensor.empty_like(v:=Tensor(kernel.src[4]))
ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
return (None, ck[0].uop, ck[1].uop, ck[2].uop)
attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
@@ -197,7 +203,9 @@ class Transformer:
h = self.tok_embeddings(tokens)
freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]
mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
if self.max_context != 0 and seqlen > 1:
mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
else: mask = None
for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
logits = self.output(self.norm(h))
if math.isnan(temperature): return logits
+112 -24
View File
@@ -1,6 +1,7 @@
# type: ignore
import ctypes, ctypes.util, struct, platform, pathlib, re, time, os, signal
from tinygrad.helpers import from_mv, to_mv, getenv, init_c_struct_t
from tinygrad.helpers import from_mv, to_mv, getenv
from tinygrad.runtime.support.c import init_c_struct_t
from hexdump import hexdump
start = time.perf_counter()
@@ -10,18 +11,21 @@ processor = platform.processor()
IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]
MMAP_SYSCALL = {"aarch64": 0xde, "x86_64":0x09}[processor]
IOCTL_PRINT = getenv("IOCTL_PRINT", getenv("IOCTL", 0))
GRAB_PMA = getenv("GRAB_PMA", 0)
def get_struct(argp, stype):
return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents
def dump_struct(st):
if getenv("IOCTL", 0) == 0: return
if IOCTL_PRINT == 0: return
print("\t", st.__class__.__name__, end=" { ")
for v in type(st)._fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
for v in type(st)._real_fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
print("}")
def format_struct(s):
sdats = []
for field in s._fields_:
for field in s._real_fields_:
dat = getattr(s, field[0])
if isinstance(dat, int): sdats.append(f"{field[0]}:0x{dat:X}")
else: sdats.append(f"{field[0]}:{dat}")
@@ -58,6 +62,29 @@ def install_hook(c_function, python_function):
return orig_func
# *** ioctl lib end ***
# PMA buffer tracking for raw PC sampling data (only when GRAB_PMA is enabled)
pma_mem_handle = 0 # hMemPmaBuffer from ALLOC_PMA_STREAM
pma_buffer_size = 0
pma_buffer_va = 0 # actual mapped VA (found via /proc/self/maps)
pma_get_offset = 0 # current read offset in ring buffer
pma_pending_map = False # flag to check for new mapping on next ioctl
pma_maps_before = set() # mappings before MAP_MEMORY
pma_raw_dumps: list[bytes] = [] # raw PMA buffer dumps
def get_pma_raw_dumps() -> list[bytes]: return pma_raw_dumps
def clear_pma_raw_dumps(): pma_raw_dumps.clear()
def get_proc_maps():
"""Read current process memory mappings as set of (start, end) tuples."""
result = set()
with open("/proc/self/maps", "r") as f:
for line in f:
addr_range = line.split()[0]
start, end = addr_range.split("-")
result.add((int(start, 16), int(end, 16)))
return result
from tinygrad.runtime.autogen import nv_570 as nv_gpu
nvescs = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("NV_ESC")}
nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(nv_gpu, "struct_"+x.replace("_CMD_", "_")+"_PARAMS", None))) for x in dir(nv_gpu) if \
@@ -69,6 +96,7 @@ def get_classes():
"NV20_SUBDEVICE_0"}
for nm,val in nv_gpu.__dict__.items():
if not isinstance(val, int): continue
if nm.endswith("PARAMETERS_MESSAGE_ID"): continue
if 0x3000 < val < 0xffff: res[val] = nm
if nm in known_classes: res[val] = nm
return res
@@ -80,37 +108,92 @@ global_ioctl_id = 0
gpus_user_modes = []
gpus_mmio = []
gpus_fifo = []
offset_load = 0
@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
def ioctl(fd, request, argp):
global global_ioctl_id, gpus_user_modes, gpus_mmio
global pma_mem_handle, pma_buffer_size, pma_buffer_va, pma_get_offset, pma_pending_map, pma_maps_before
global_ioctl_id += 1
# Check for new PMA buffer mapping from previous MAP_MEMORY call (only when GRAB_PMA is enabled)
if GRAB_PMA and pma_pending_map:
pma_pending_map = False
new_maps = get_proc_maps()
for start, end in new_maps - pma_maps_before:
if end - start == pma_buffer_size:
pma_buffer_va = start
if IOCTL_PRINT >= 1: print(f"\t PMA buffer mapped at CPU VA=0x{pma_buffer_va:x}")
break
st = time.perf_counter()
ret = libc.syscall(IOCTL_SYSCALL, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
et = time.perf_counter()-st
fn = os.readlink(f"/proc/self/fd/{fd}")
#print(f"ioctl {request:8x} {fn:20s}")
idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
if getenv("IOCTL", 0) >= 1: print(f"#{global_ioctl_id}: ", end="")
if IOCTL_PRINT >= 1: print(f"#{global_ioctl_id}: ", end="")
if itype == ord(nv_gpu.NV_IOCTL_MAGIC):
if nr == nv_gpu.NV_ESC_RM_CONTROL:
s = get_struct(argp, nv_gpu.NVOS54_PARAMETERS)
if s.cmd in nvcmds:
name, struc = nvcmds[s.cmd]
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
print(f"NV_ESC_RM_CONTROL cmd={name:30s} hClient={s.hClient}, hObject={s.hObject}, flags={s.flags}, params={s.params}, paramsSize={s.paramsSize}, status={s.status}")
if struc is not None: dump_struct(get_struct(s.params, struc))
elif hasattr(nv_gpu, name+"_PARAMS"): dump_struct(get_struct(argp, getattr(nv_gpu, name+"_PARAMS")))
elif name == "NVA06C_CTRL_CMD_GPFIFO_SCHEDULE": dump_struct(get_struct(argp, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS))
elif name == "NV83DE_CTRL_CMD_GET_MAPPINGS": dump_struct(get_struct(s.params, nv_gpu.NV83DE_CTRL_DEBUG_GET_MAPPINGS_PARAMETERS))
elif name == "NVB0CC_CTRL_CMD_SET_HS_CREDITS":
hs_params = get_struct(s.params, nv_gpu.NVB0CC_CTRL_SET_HS_CREDITS_PARAMS)
dump_struct(hs_params)
if IOCTL_PRINT >= 2:
for i in range(hs_params.numEntries):
print(f"\t\t", end="")
dump_struct(hs_params.creditInfo[i])
# PMA buffer tracking (only when GRAB_PMA is enabled)
if GRAB_PMA and name == "NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM":
pma_params = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS)
pma_mem_handle = pma_params.hMemPmaBuffer
pma_buffer_size = pma_params.pmaBufferSize
pma_get_offset = 0 # Reset read offset for new stream
if IOCTL_PRINT >= 1: print(f"\t PMA buffer: hMem=0x{pma_mem_handle:x} size={pma_buffer_size}")
if GRAB_PMA and name == "NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT":
pma_update = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS)
if pma_update.bytesAvailable > 0 and pma_buffer_va and pma_buffer_size > 0:
avail = pma_update.bytesAvailable
read_offset = pma_get_offset
# Handle ring buffer wrap-around
if pma_get_offset + avail <= pma_buffer_size:
pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, avail))
else:
# Wrap around: read to end, then from start
first_part = pma_buffer_size - pma_get_offset
second_part = avail - first_part
pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, first_part)) + bytes(to_mv(pma_buffer_va, second_part))
pma_raw_dumps.append(pma_data)
pma_get_offset = (pma_get_offset + avail) % pma_buffer_size
if IOCTL_PRINT >= 2:
print(f"\t PMA data: {avail} bytes from offset=0x{read_offset:x}, new offset=0x{pma_get_offset:x}")
hexdump(pma_data)
# Dump regOps for EXEC_REG_OPS when IOCTL >= 3
if name == "NVB0CC_CTRL_CMD_EXEC_REG_OPS" and struc is not None and IOCTL_PRINT >= 3:
reg_params = get_struct(s.params, struc)
for i in range(reg_params.regOpCount):
print(f"\t\t", end="")
dump_struct(reg_params.regOps[i])
# val = (op.regValueHi << 32) | op.regValueLo
# print(f"\t regOps[{i:3d}]: op={op.regOp} type={op.regType} status={op.regStatus} offset=0x{op.regOffset:08x} value=0x{val:016x}")
else:
if getenv("IOCTL", 0) >= 1: print("unhandled cmd", hex(s.cmd))
if IOCTL_PRINT >= 1: print("unhandled cmd", hex(s.cmd))
# format_struct(s)
# print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
elif nr == nv_gpu.NV_ESC_RM_ALLOC:
s = get_struct(argp, nv_gpu.NVOS21_PARAMETERS)
if getenv("IOCTL", 0) >= 1: print(f"NV_ESC_RM_ALLOC hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
if IOCTL_PRINT >= 1: print(f"NV_ESC_RM_ALLOC hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
if s.pAllocParms is not None:
if s.hClass == nv_gpu.NV01_DEVICE_0: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV0080_ALLOC_PARAMETERS))
if s.hClass == nv_gpu.FERMI_VASPACE_A: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS))
@@ -118,7 +201,8 @@ def ioctl(fd, request, argp):
if s.hClass == nv_gpu.NV1_MEMORY_USER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
if s.hClass == nv_gpu.NV1_MEMORY_SYSTEM: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
if s.hClass == nv_gpu.GT200_DEBUGGER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV83DE_ALLOC_PARAMETERS))
if s.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
if s.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE: dump_struct(get_struct(s.pAllocParms, nv_gpu.NVB2CC_ALLOC_PARAMETERS))
if s.hClass in {nv_gpu.AMPERE_CHANNEL_GPFIFO_A, nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A}:
sx = get_struct(s.pAllocParms, nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS)
dump_struct(sx)
gpus_fifo.append((sx.gpFifoOffset, sx.gpFifoEntries))
@@ -126,31 +210,35 @@ def ioctl(fd, request, argp):
if s.hClass == nv_gpu.TURING_USERMODE_A: gpus_user_modes.append(s.hObjectNew)
elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
# nv_ioctl_nvos33_parameters_with_fd
if getenv("IOCTL", 0) >= 1:
s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
if IOCTL_PRINT >= 1:
print(f"NV_ESC_RM_MAP_MEMORY hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, length={s.length} flags={s.flags} pLinearAddress={s.pLinearAddress}")
# Track PMA buffer mapping - save maps now, check for new mapping on next ioctl (after mmap happens)
if GRAB_PMA and pma_mem_handle and s.hMemory == pma_mem_handle:
pma_maps_before = get_proc_maps()
pma_pending_map = True
elif nr == nv_gpu.NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO:
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
s = get_struct(argp, nv_gpu.NVOS56_PARAMETERS)
print(f"NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, pOldCpuAddress={s.pOldCpuAddress} pNewCpuAddress={s.pNewCpuAddress} status={s.status}")
elif nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
s = get_struct(argp, nv_gpu.nv_ioctl_nvos02_parameters_with_fd)
print(f"NV_ESC_RM_ALLOC_MEMORY fd={s.fd}, hRoot={s.params.hRoot}, hObjectParent={s.params.hObjectParent}, hObjectNew={s.params.hObjectNew}, hClass={s.params.hClass}, flags={s.params.flags}, pMemory={s.params.pMemory}, limit={s.params.limit}, status={s.params.status}")
elif nr == nv_gpu.NV_ESC_ALLOC_OS_EVENT:
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
s = get_struct(argp, nv_gpu.nv_ioctl_alloc_os_event_t)
print(f"NV_ESC_ALLOC_OS_EVENT hClient={s.hClient} hDevice={s.hDevice} fd={s.fd} Status={s.Status}")
elif nr == nv_gpu.NV_ESC_REGISTER_FD:
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
s = get_struct(argp, nv_gpu.nv_ioctl_register_fd_t)
print(f"NV_ESC_REGISTER_FD fd={s.ctl_fd}")
elif nr in nvescs:
if getenv("IOCTL", 0) >= 1: print(nvescs[nr])
if IOCTL_PRINT >= 1: print(nvescs[nr])
else:
if getenv("IOCTL", 0) >= 1: print("unhandled NR", nr)
if IOCTL_PRINT >= 1: print("unhandled NR", nr)
elif fn.endswith("nvidia-uvm"):
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
print(f"{nvuvms.get(request, f'UVM UNKNOWN {request=}')}")
if nvuvms.get(request) is not None: dump_struct(get_struct(argp, getattr(nv_gpu, nvuvms.get(request)+"_PARAMS")))
if nvuvms.get(request) == "UVM_MAP_EXTERNAL_ALLOCATION":
@@ -159,7 +247,7 @@ def ioctl(fd, request, argp):
print("perGpuAttributes[{i}] = ", end="")
dump_struct(st.perGpuAttributes[i])
if getenv("IOCTL") >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
if IOCTL_PRINT >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
return ret
@ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long)
@@ -172,14 +260,14 @@ def _mmap(addr, length, prot, flags, fd, offset):
return ret
install_hook(libc.ioctl, ioctl)
if getenv("IOCTL") >= 3: orig_mmap_mv = install_hook(libc.mmap, _mmap)
if getenv("IOCTL") >= 4: orig_mmap_mv = install_hook(libc.mmap, _mmap)
import collections
old_gpputs = collections.defaultdict(int)
def _dump_gpfifo(mark):
launches = []
# print("_dump_gpfifo:", mark)
print("_dump_gpfifo:", mark)
for start, size in gpus_fifo:
gpfifo_controls = nv_gpu.AmpereAControlGPFifo.from_address(start+size*8)
gpfifo = to_mv(start, size * 8).cast("Q")
@@ -205,7 +293,7 @@ def make_qmd_struct_type():
fields.append((name.replace("NVC6C0_QMDV03_00_", "").lower(), ctypes.c_uint32, data[0]-data[1]+1))
if len(fields) >= 2 and fields[-2][0].endswith('_lower') and fields[-1][0].endswith('_upper') and fields[-1][0][:-6] == fields[-2][0][:-6]:
fields = fields[:-2] + [(fields[-1][0][:-6], ctypes.c_uint64, fields[-1][2] + fields[-2][2])]
return init_c_struct_t(tuple(fields))
return init_c_struct_t(0x40 * 4, tuple(fields))
qmd_struct_t = make_qmd_struct_type()
assert ctypes.sizeof(qmd_struct_t) == 0x40 * 4
@@ -222,7 +310,7 @@ def _dump_qmd(address, packets):
subc = (dat>>13) & 7
mthd = (dat<<2) & 0x7FFF
method_name = nvqcmds.get(mthd, f"unknown method #{mthd}")
if getenv("IOCTL", 0) >= 1:
if IOCTL_PRINT >= 1:
print(f"\t\t{method_name}, {typ=} {size=} {subc=} {mthd=}")
for j in range(size): print(f"\t\t\t{j}: {gpfifo[i+j+1]} | 0x{gpfifo[i+j+1]:x}")
if mthd == 792:
+1
View File
@@ -0,0 +1 @@
examples/
+135
View File
@@ -0,0 +1,135 @@
import pickle, os, sys, functools, numpy as np
from pathlib import Path
os.environ["DEV"] = "CUDA"
os.environ["PROFILE"] = os.environ.get("PROFILE", "2")
from extra.nv_pma.cupti import cu_prof_ext
cu_prof_ext.enable_auto()
from tinygrad import Tensor, Device
if not os.environ.get("IOCTL") or not os.environ.get("GRAB_PMA"):
print("Usage: GRAB_PMA=1 IOCTL=1 IOCTL_PRINT=0 python3 extra/nv_pma/collect.py")
sys.exit(1)
assert Device.DEFAULT == "CUDA", "only works with CUDA"
EXAMPLES_DIR = Path(__file__).parent / "examples"
_collectors: list[tuple[str, callable]] = []
def pcsampling_test(name: str):
def decorator(fn):
@functools.wraps(fn)
def wrapper():
cu_prof_ext.clear_pma_raw_dumps()
cu_prof_ext.clear_cupti_pc_samples()
fn()
Device["CUDA"].synchronize()
dumps = cu_prof_ext.get_pma_raw_dumps()
# from hexdump import hexdump
# hexdump(dumps[0][:0x40])
return {"test_name": name, "pma_raw_dumps": list(cu_prof_ext.get_pma_raw_dumps()), "cupti_pc_samples": list(cu_prof_ext.get_cupti_pc_samples())}
_collectors.append((name, wrapper))
return wrapper
return decorator
# Refs
@pcsampling_test("test_plus")
def test_plus():
a = Tensor([1, 2, 3, 4])
b = Tensor([5, 6, 7, 8])
(a + b).realize()
@pcsampling_test("test_matmul")
def test_matmul():
a = Tensor(np.random.rand(12, 12).astype(np.float32))
b = Tensor(np.random.rand(12, 12).astype(np.float32))
(a @ b).realize()
@pcsampling_test("test_reduce_sum")
def test_reduce_sum():
a = Tensor(np.random.rand(1024).astype(np.float32))
a.sum().realize()
@pcsampling_test("test_reduce_max")
def test_reduce_max():
a = Tensor(np.random.rand(1024).astype(np.float32))
a.max().realize()
@pcsampling_test("test_exp")
def test_exp():
a = Tensor(np.random.rand(256).astype(np.float32))
a.exp().realize()
@pcsampling_test("test_softmax")
def test_softmax():
a = Tensor(np.random.rand(64, 64).astype(np.float32))
a.softmax().realize()
@pcsampling_test("test_conv2d")
def test_conv2d():
x = Tensor(np.random.rand(1, 3, 32, 32).astype(np.float32))
w = Tensor(np.random.rand(8, 3, 3, 3).astype(np.float32))
x.conv2d(w).realize()
@pcsampling_test("test_large_matmul")
def test_large_matmul():
a = Tensor(np.random.rand(128, 128).astype(np.float32))
b = Tensor(np.random.rand(128, 128).astype(np.float32))
(a @ b).realize()
@pcsampling_test("test_elementwise_chain")
def test_elementwise_chain():
a = Tensor(np.random.rand(512).astype(np.float32))
((a + 1) * 2 - 0.5).relu().realize()
@pcsampling_test("test_broadcast")
def test_broadcast():
a = Tensor(np.random.rand(64, 1).astype(np.float32))
b = Tensor(np.random.rand(1, 64).astype(np.float32))
(a + b).realize()
@pcsampling_test("test_plus_big")
def test_plus_big():
a = Tensor(np.random.rand(64, 32).astype(np.float32))
b = Tensor(np.random.rand(64, 32).astype(np.float32))
(a + b).realize()
def save_example(name: str, data: dict):
pma_bytes = sum(len(d) for d in data['pma_raw_dumps'])
cupti_samples = sum(r['samples'] for r in data['cupti_pc_samples'])
print(f" PMA: {len(data['pma_raw_dumps'])} buffers, {pma_bytes} bytes")
print(f" CUPTI: {len(data['cupti_pc_samples'])} records, {cupti_samples} samples")
outfile = EXAMPLES_DIR / f"{name}.pkl"
with open(outfile, "wb") as f:
pickle.dump(data, f)
print(f" Saved to {outfile}")
if __name__ == "__main__":
EXAMPLES_DIR.mkdir(exist_ok=True)
# Run specific tests if provided as arguments, otherwise run all
if len(sys.argv) > 1:
test_names = sys.argv[1:]
collectors = [(name, fn) for name, fn in _collectors if name in test_names]
if not collectors:
print(f"Unknown tests: {test_names}")
print(f"Available: {[name for name, _ in _collectors]}")
sys.exit(1)
else:
collectors = _collectors
for name, collect_fn in collectors:
print(f"\nCollecting {name}...")
try:
data = collect_fn()
save_example(name, data)
except Exception as e:
print(f" ERROR: {e}")
import traceback
traceback.print_exc()
+26
View File
@@ -0,0 +1,26 @@
# CUPTI autogen loader for nv_pma
# To regenerate: REGEN=1 python -c "import extra.nv_pma.cupti"
import importlib, pathlib
from tinygrad.helpers import getenv
root = pathlib.Path(__file__).parents[3]
here = pathlib.Path(__file__).parent
def load(name, dll, files, **kwargs):
if not (f:=here/f"{name}.py").exists() or getenv('REGEN'):
kwargs['args'] = kwargs.get('args', [])
f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(name, dll, files, **kwargs))
return importlib.import_module(f"extra.nv_pma.cupti.{name}")
def __getattr__(nm):
match nm:
case "cupti":
return load("cupti", "'/usr/local/cuda/targets/x86_64-linux/lib/libcupti.so'", [
"/usr/local/cuda/include/cupti_result.h", "/usr/local/cuda/include/cupti_activity.h",
"/usr/local/cuda/include/cupti_callbacks.h", "/usr/local/cuda/include/cupti_events.h",
"/usr/local/cuda/include/cupti_metrics.h", "/usr/local/cuda/include/cupti_driver_cbid.h",
"/usr/local/cuda/include/cupti_runtime_cbid.h", "/usr/local/cuda/include/cupti_profiler_target.h",
"/usr/local/cuda/include/cupti_profiler_host.h", "/usr/local/cuda/include/cupti_pmsampling.h",
"/usr/local/cuda/include/generated_cuda_meta.h", "/usr/local/cuda/include/generated_cuda_runtime_api_meta.h"
], args=["-D__CUDA_API_VERSION_INTERNAL", "-I/usr/local/cuda/include"], parse_macros=False)
case _: raise AttributeError(f"no such autogen: {nm}")
+164
View File
@@ -0,0 +1,164 @@
from __future__ import annotations
import ctypes
from tinygrad.helpers import DEBUG, getenv
from extra.nv_pma.cupti import cupti
def stall_reason_name(reason: int) -> str:
name = cupti.CUpti_ActivityPCSamplingStallReason.get(reason)
return name.replace("CUPTI_ACTIVITY_PC_SAMPLING_STALL_", "").lower() if name else str(reason)
class CUPTIProfiler:
def __init__(self):
self.initialized = False
self.pc_sampling_enabled = False
self.buffers: list[ctypes.Array] = []
self.kernel_stalls: dict[int, dict[int, int]] = {}
self.raw_buffers: list[bytes] = []
self.pc_samples: list[dict] = []
def _check_cupti(self, status, soft=False):
if status != cupti.CUPTI_SUCCESS:
if soft: return False
raise RuntimeError(f"CUPTI Error {status}")
return True
def init(self, ctx, device_id: int = 0, profile_level: int = 2):
if self.initialized: return
# Initialize profiler API
init_params = cupti.CUpti_Profiler_Initialize_Params()
init_params.structSize = 16
cupti.cuptiProfilerInitialize(ctypes.byref(init_params))
# Register buffer callbacks for Activity API
self._buf_req_cb = cupti.CUpti_BuffersCallbackRequestFunc(self._buffer_requested)
self._buf_comp_cb = cupti.CUpti_BuffersCallbackCompleteFunc(self._buffer_completed)
self._check_cupti(cupti.cuptiActivityRegisterCallbacks(self._buf_req_cb, self._buf_comp_cb))
# PROFILE=1: kernel timing, PROFILE=2: PC sampling with stall reasons
if profile_level >= 2:
# PC sampling for stall analysis (requires elevated privileges)
if DEBUG >= 1: print(" CUPTI: PC sampling mode (before)")
pc_status = cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING)
if pc_status == cupti.CUPTI_SUCCESS:
config = cupti.CUpti_ActivityPCSamplingConfig()
config.size, config.samplingPeriod = 16, cupti.CUPTI_ACTIVITY_PC_SAMPLING_PERIOD_MIN
cfg_status = cupti.dll.cuptiActivityConfigurePCSampling(ctx, ctypes.byref(config))
if cfg_status == cupti.CUPTI_SUCCESS:
if DEBUG >= 1: print(" CUPTI: PC sampling mode (before stall analysis)")
cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO)
self.pc_sampling_enabled = True
if DEBUG >= 1: print(" CUPTI: PC sampling mode (stall analysis)")
elif cfg_status == 35:
if DEBUG >= 1: print(" CUPTI: PC sampling needs: echo 'options nvidia NVreg_RestrictProfilingToAdminUsers=0'|sudo tee /etc/modprobe.d/nvidia.conf && sudo reboot")
# Fall back to kernel timing if PC sampling setup failed
if not self.pc_sampling_enabled:
self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
else:
# Kernel activity tracing for timing
self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
self.initialized = True
def _buffer_requested(self, buffer, size, max_num_records):
buf = (ctypes.c_uint8 * 1024 * 1024)() # 1MB buffer
self.buffers.append(buf)
buffer[0] = ctypes.cast(buf, ctypes.POINTER(ctypes.c_uint8))
size[0] = ctypes.sizeof(buf)
max_num_records[0] = 0
def _buffer_completed(self, ctx, stream_id, buffer, size, valid_size):
if valid_size > 0:
record = ctypes.POINTER(cupti.CUpti_Activity)()
while cupti.cuptiActivityGetNextRecord(buffer, valid_size, ctypes.byref(record)) == cupti.CUPTI_SUCCESS:
kind = record.contents.kind
if kind == cupti.CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL:
kernel = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityKernel9)).contents
name = ctypes.string_at(kernel.name).decode() if kernel.name else "unknown"
duration_us = (kernel.end - kernel.start) / 1000.0
grid, block = (kernel.gridX, kernel.gridY, kernel.gridZ), (kernel.blockX, kernel.blockY, kernel.blockZ)
print(f" CUPTI: {name[:40]:40s} | {duration_us:10.2f} us | grid={grid} block={block} | regs={kernel.registersPerThread:3d} smem={kernel.staticSharedMemory + kernel.dynamicSharedMemory:6d}B")
elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING:
pc = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSampling3)).contents
cid = pc.correlationId
if cid not in self.kernel_stalls: self.kernel_stalls[cid] = {}
self.kernel_stalls[cid][pc.stallReason] = self.kernel_stalls[cid].get(pc.stallReason, 0) + pc.samples
self.pc_samples.append({
'correlationId': pc.correlationId, 'pcOffset': pc.pcOffset, 'stallReason': pc.stallReason,
'samples': pc.samples, 'latencySamples': pc.latencySamples, 'functionId': pc.functionId, 'sourceLocatorId': pc.sourceLocatorId
})
if DEBUG >= 3:
print(f" PC {pc.pcOffset:#x} stall={stall_reason_name(pc.stallReason)} samples={pc.samples} latency={pc.latencySamples} func={pc.functionId} src={pc.sourceLocatorId}")
elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO:
info = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSamplingRecordInfo)).contents
cid = info.correlationId
if cid in self.kernel_stalls:
stalls = self.kernel_stalls[cid]
total = sum(stalls.values())
if total > 0:
top = sorted(stalls.items(), key=lambda x: -x[1])[:5]
stall_str = " ".join(f"{stall_reason_name(r)}:{100*c//total}%" for r,c in top if c > 0)
print(f" CUPTI stalls (corr={cid}): {total} samples | {stall_str}")
del self.kernel_stalls[cid]
else: print(f" CUPTI: Unhandled activity kind {kind}")
def flush(self):
if not self.initialized: return
self._check_cupti(cupti.cuptiActivityFlushAll(0))
# Module-level profiler instance
_profiler: CUPTIProfiler | None = None
def get_profiler() -> CUPTIProfiler | None:
return _profiler
def get_cupti_raw_buffers() -> list[bytes]:
return _profiler.raw_buffers if _profiler else []
def clear_cupti_raw_buffers():
if _profiler: _profiler.raw_buffers.clear()
def get_cupti_pc_samples() -> list[dict]:
return _profiler.pc_samples if _profiler else []
def clear_cupti_pc_samples():
if _profiler: _profiler.pc_samples.clear()
# Raw PMA buffer access (from ioctl interception)
def get_pma_raw_dumps() -> list[bytes]:
try:
from extra.nv_gpu_driver.nv_ioctl import get_pma_raw_dumps as _get
return _get()
except ImportError: return []
def clear_pma_raw_dumps():
try:
from extra.nv_gpu_driver.nv_ioctl import clear_pma_raw_dumps as _clear
_clear()
except ImportError: pass
def enable(profile_level:int=2):
global _profiler
if _profiler is not None: return
_profiler = CUPTIProfiler()
# Patch CUDADevice to initialize CUPTI profiler
from tinygrad.runtime.ops_cuda import CUDADevice
_orig_init = CUDADevice.__init__
_orig_sync = CUDADevice.synchronize
def _patched_init(self, device: str):
_orig_init(self, device)
device_id = int(device.split(":")[1]) if ":" in device else 0
_profiler.init(self.context, device_id, profile_level)
def _patched_sync(self):
_orig_sync(self)
if _profiler: _profiler.flush()
CUDADevice.__init__ = _patched_init
CUDADevice.synchronize = _patched_sync
def enable_auto():
if (profile_level:=getenv("PROFILE", 0)) > 0: enable(profile_level)
File diff suppressed because it is too large Load Diff
+193
View File
@@ -0,0 +1,193 @@
#!/usr/bin/env python3
from __future__ import annotations
import enum, collections
from typing import Iterator
from tinygrad.helpers import colored
from extra.assembly.amd.sqtt import PacketType, bits
# ═══════════════════════════════════════════════════════════════════════════════
# STALL REASONS
# ═══════════════════════════════════════════════════════════════════════════════
class StallReason(enum.IntEnum):
# Based on CUpti_ActivityPCSamplingStallReason
INVALID = 0
NONE = 1 # selected, selected_not_issued
INST_FETCH = 2 # branch_resolving, no_instructions
EXEC_DEPENDENCY = 3 # short_scoreboard, wait
MEMORY_DEPENDENCY = 4 # long_scoreboard
TEXTURE = 5 # tex_throttle
SYNC = 6 # barrier, membar
CONSTANT_MEMORY = 7 # imc_miss
PIPE_BUSY = 8 # mio_throttle, math_pipe_throttle
MEMORY_THROTTLE = 9 # drain, lg_throttle
NOT_SELECTED = 10 # not_selected
OTHER = 11 # misc, dispatch_stall
SLEEPING = 12 # sleeping
STALL_KEY_MAP_AMPERE: dict[int, StallReason] = {
1: StallReason.MEMORY_THROTTLE, 15: StallReason.MEMORY_THROTTLE,
2: StallReason.CONSTANT_MEMORY,
3: StallReason.SYNC,
6: StallReason.INST_FETCH, 11: StallReason.INST_FETCH,
7: StallReason.EXEC_DEPENDENCY, 10: StallReason.EXEC_DEPENDENCY,
9: StallReason.MEMORY_DEPENDENCY,
12: StallReason.PIPE_BUSY,
17: StallReason.OTHER, 20: StallReason.OTHER,
18: StallReason.NONE,
}
STALL_KEY_MAP_BLACKWELL: dict[int, StallReason] = {
0x01: StallReason.MEMORY_THROTTLE, 0x0e: StallReason.MEMORY_THROTTLE,
0x02: StallReason.SYNC,
0x05: StallReason.INST_FETCH, 0x0a: StallReason.INST_FETCH,
0x06: StallReason.EXEC_DEPENDENCY, 0x09: StallReason.EXEC_DEPENDENCY,
0x08: StallReason.MEMORY_DEPENDENCY,
0x0b: StallReason.PIPE_BUSY, 0x0f: StallReason.PIPE_BUSY,
0x10: StallReason.OTHER, 0x13: StallReason.OTHER,
0x11: StallReason.NONE,
}
# Lookup table for extracting sample bytes from 32-byte packet (bytes 0-3, 8-31, skipping header at 4-7)
LOOKUP_28B = [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET HEADER
# ═══════════════════════════════════════════════════════════════════════════════
class PMAHeader(PacketType):
num_bytes = bits[4:0] # number of sample bytes in this packet
tpc_id_lo = bits[15:8] # TPC identifier low 8 bits
tpc_id_hi = bits[27:25] # TPC identifier high 3 bits
dropped = bits[28:28] # dropped flag (resets byte accumulator)
@property
def tpc_id(self) -> int: return self.tpc_id_lo | (self.tpc_id_hi << 8)
# ═══════════════════════════════════════════════════════════════════════════════
# 8-BYTE SAMPLE FORMAT (Ampere/Ada/Hopper)
# ═══════════════════════════════════════════════════════════════════════════════
class PMASampleAmpere8B(PacketType):
pc_raw = bits[44:0] # raw PC value (pc_offset = pc_raw << 4)
stall_key = bits[49:45] # stall reason key
wave_id = bits[55:50] # warp/wave identifier
active = bits[62:62] # 1 if warp was executing, 0 if scheduled but not issued
@property
def pc_offset(self) -> int: return self.pc_raw << 4
@property
def stall_reason(self) -> StallReason: return STALL_KEY_MAP_AMPERE.get(self.stall_key, StallReason.OTHER)
# ═══════════════════════════════════════════════════════════════════════════════
# 9-BYTE SAMPLE FORMAT (Blackwell+)
# ═══════════════════════════════════════════════════════════════════════════════
class PMASampleBlackwell9B(PacketType):
stall_key = bits[5:0] # stall reason key
pc_raw = bits[60:8] # raw PC value (pc_offset = pc_raw << 4)
wave_hi = bits[7:6] # wave_id high 2 bits
wave_lo = bits[71:68] # wave_id low 4 bits
active = bits[67:67] # 1 if warp was executing, 0 if scheduled but not issued
@property
def pc_offset(self) -> int: return self.pc_raw << 4
@property
def stall_reason(self) -> StallReason: return STALL_KEY_MAP_BLACKWELL.get(self.stall_key, StallReason.OTHER)
@property
def wave_id(self) -> int: return (self.wave_hi << 4) | self.wave_lo
PMASample = PMASampleAmpere8B|PMASampleBlackwell9B
def decode(data: bytes, sm_version: int = 0x800) -> Iterator[tuple[PMASample, int]]:
use_9byte = sm_version >= 0xa04
record_size = 9 if use_9byte else 8
sample_cls = PMASampleBlackwell9B if use_9byte else PMASampleAmpere8B
tpc_state: dict[int, list[int]] = collections.defaultdict(list)
for pkt_idx in range(len(data) // 32):
pkt = data[pkt_idx * 32:(pkt_idx + 1) * 32]
hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
if hdr.dropped: tpc_state[hdr.tpc_id].clear()
for i in range(hdr.num_bytes):
tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[i]])
while len(tpc_state[hdr.tpc_id]) >= record_size:
yield sample_cls.from_raw(int.from_bytes(bytes(tpc_state[hdr.tpc_id][:record_size]), 'little')), hdr.tpc_id
del tpc_state[hdr.tpc_id][:record_size]
# ═══════════════════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════════════════
STALL_COLORS = {
StallReason.NONE: "green", StallReason.INST_FETCH: "yellow", StallReason.EXEC_DEPENDENCY: "cyan",
StallReason.MEMORY_DEPENDENCY: "red", StallReason.SYNC: "magenta", StallReason.CONSTANT_MEMORY: "blue",
StallReason.PIPE_BUSY: "yellow", StallReason.MEMORY_THROTTLE: "RED", StallReason.OTHER: "white",
}
def decode_tpc_id(tpc_id:int) -> tuple[int, int, int]:
# NOTE: valid only for ops_nv, cuda encoding is different
return (tpc_id >> 5, (tpc_id >> 1) & 0xf, tpc_id & 1)
def print_samples(samples:list[tuple[PMASample, int]]) -> None:
if not samples: return
base_pc = min(s.pc_offset for s, _ in samples)
for s, tpc_id in samples:
gpc, tpc, sm = decode_tpc_id(tpc_id)
stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
print(f"pc=0x{s.pc_offset - base_pc:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
def print_packets(data:bytes, sm_version:int=0x800) -> None:
record_size = 9 if sm_version >= 0x890 else 8
tpc_state: dict[int, list[int]] = collections.defaultdict(list)
for i in range(len(data) // 32):
pkt = data[i * 32:(i + 1) * 32]
hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
if hdr.dropped: tpc_state[hdr.tpc_id].clear()
for j in range(hdr.num_bytes): tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[j]])
# Show complete records extracted from this packet
records = []
while len(tpc_state[hdr.tpc_id]) >= record_size:
records.append(bytes(tpc_state[hdr.tpc_id][:record_size]).hex())
del tpc_state[hdr.tpc_id][:record_size]
leftover = len(tpc_state[hdr.tpc_id])
print(f"Pkt {i:3d}: tpc={hdr.tpc_id:4d} n={hdr.num_bytes:2d} drop={hdr.dropped} left={leftover} | {' '.join(records)}")
def print_aggregated(samples:list[tuple[PMASample, int]]) -> None:
if not samples: return
base_pc = min(s.pc_offset for s, _ in samples)
counter: collections.Counter[tuple[int, StallReason]] = collections.Counter((s.pc_offset, s.stall_reason) for s, _ in samples)
print(f"\nAggregated samples (base_pc=0x{base_pc:x}):")
for (pc, reason), cnt in sorted(counter.items()):
stall_str = colored(f"{reason.name:17}", STALL_COLORS.get(reason, "white"))
print(f" pc=0x{pc - base_pc:06x} {stall_str} samples={cnt:4d}")
if __name__ == "__main__":
import sys, pickle
if len(sys.argv) < 2:
print("Usage: python decode.py <pkl_file> [--raw] [--sm=0xNNN]")
sys.exit(1)
with open(sys.argv[1], "rb") as f:
data = pickle.load(f)
if isinstance(data, dict):
sm_version = 0x800 # default to Ampere
for arg in sys.argv:
if arg.startswith("--sm="): sm_version = int(arg[5:], 0)
dumps = [(i, x, sm_version) for i, x in enumerate(data["pma_raw_dumps"])]
else:
devs = {e.device: e for e in data if type(e).__name__ == "ProfileDeviceEvent"}
dumps = []
for i, e in enumerate(e for e in data if type(e).__name__ == "ProfilePMAEvent"):
dumps.append((i, e.blob, devs[e.device].props.get('sm_version', 0x800)))
for dump_idx, raw, sm_ver in dumps:
print(f"\n{'='*60}\nDump {dump_idx} ({len(raw)} bytes, {len(raw)//32} packets)\n{'='*60}")
if "--raw" in sys.argv: print_packets(raw, sm_ver)
else:
samples = list(decode(raw, sm_ver))
print(f"\nDecoded {len(samples)} samples:")
print_samples(samples)
print_aggregated(samples)
+76
View File
@@ -0,0 +1,76 @@
import pickle, unittest
from collections import Counter
from pathlib import Path
from extra.nv_pma.decode import decode
from tinygrad.helpers import DEBUG
EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
EXAMPLES_5090_DIR = Path(__file__).parent.parent / "examples_5090"
def decode_and_aggregate(raw_dumps: list[bytes], sm_version: int = 0x800) -> Counter[tuple[int, int]]:
"""Decode all PMA buffers and aggregate by (relative_pc, stall_reason). Each dump is normalized separately."""
result: Counter[tuple[int, int]] = Counter()
for raw in raw_dumps:
samples = [s for s, _ in decode(raw, sm_version)]
if not samples: continue
base_pc = min(s.pc_offset for s in samples)
result += Counter((s.pc_offset - base_pc, int(s.stall_reason)) for s in samples)
return result
def cupti_to_counter(cupti_records: list[dict]) -> Counter[tuple[int, int]]:
"""Convert CUPTI records to Counter[(pcOffset, stallReason)]."""
counter: Counter[tuple[int, int]] = Counter()
for r in cupti_records:
counter[(r['pcOffset'], r['stallReason'])] += r['samples']
return counter
class TestNVProf(unittest.TestCase):
def _test_example(self, name: str, sm_version: int = 0x800, examples_dir: Path = EXAMPLES_DIR):
pkl_file = examples_dir / f"{name}.pkl"
if not pkl_file.exists():
self.skipTest(f"Example data not found: {pkl_file}. Run collect.py first.")
with open(pkl_file, "rb") as f:
data = pickle.load(f)
self.assertEqual(data["test_name"], name)
pma_agg = decode_and_aggregate(data["pma_raw_dumps"], sm_version)
cupti_agg = cupti_to_counter(data["cupti_pc_samples"])
if DEBUG >= 2:
total = sum(cupti_agg.values())
mismatched = sum(abs(pma_agg.get(k, 0) - v) for k, v in cupti_agg.items())
mismatched += sum(v for k, v in pma_agg.items() if k not in cupti_agg)
mismatched //= 2
print(f"\n=== Test: {name} ===")
print(f"Total samples: {total}, Mismatched: {mismatched} ({mismatched/total*100 if total else 0:.1f}%)")
self.assertEqual(pma_agg, cupti_agg, f"PMA: {dict(pma_agg)}\nCUPTI: {dict(cupti_agg)}")
# Ampere tests (8-byte format)
def test_decode_test_plus(self): self._test_example("test_plus")
def test_decode_test_reduce_sum(self): self._test_example("test_reduce_sum")
def test_decode_test_broadcast(self): self._test_example("test_broadcast")
def test_decode_test_matmul(self): self._test_example("test_matmul")
def test_decode_test_plus_big(self): self._test_example("test_plus_big")
def test_decode_test_elementwise_chain(self): self._test_example("test_elementwise_chain")
def test_decode_test_conv2d(self): self._test_example("test_conv2d")
def test_decode_test_large_matmul(self): self._test_example("test_large_matmul")
# Blackwell/5090 tests (9-byte format)
def test_5090_test_plus(self): self._test_example("test_plus", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_plus_big(self): self._test_example("test_plus_big", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_broadcast(self): self._test_example("test_broadcast", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_matmul(self): self._test_example("test_matmul", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_large_matmul(self): self._test_example("test_large_matmul", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_reduce_sum(self): self._test_example("test_reduce_sum", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_reduce_max(self): self._test_example("test_reduce_max", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_elementwise_chain(self): self._test_example("test_elementwise_chain", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_conv2d(self): self._test_example("test_conv2d", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_exp(self): self._test_example("test_exp", 0xa04, EXAMPLES_5090_DIR)
def test_5090_test_softmax(self): self._test_example("test_softmax", 0xa04, EXAMPLES_5090_DIR)
if __name__ == "__main__":
unittest.main()

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