Compare commits

...
Author SHA1 Message Date
geohot 65aa41a116 hwtest fixes for rdna3 dsl 2025-12-28 20:23:10 -05:00
George HotzandGitHub d9603c1bee improve asm dsl syntax (#13864)
* improve asm dsl syntax

* improve asm dsl syntax
2025-12-28 20:04:59 -05:00
chenyuandGitHub f5090192c8 reorder AMD tensor core benchmark test (#13860)
* reorder AMD tensor core benchmark test

* disable that
2025-12-28 12:29:51 -05:00
qazalandGitHub 066d96c397 print tflops in asm gemm test (#13859)
* print tflops in asm gemm test

* change order
2025-12-29 02:26:40 +09:00
chenyuandGitHub a03cd43e78 fix typing in compute_gradient (#13852) 2025-12-28 11:52:14 -05:00
chenyuandGitHub cba05acadf re-enable TYPED=1 import test (#13858) 2025-12-28 11:49:06 -05:00
qazalandGitHub 2cfbabdc34 mi350x 1tflop bf16 gemm in extra (#13702) 2025-12-28 21:45:42 +09:00
qazalandGitHub 2180eee5e4 use the asm dsl in remu hwtest.py (#13856)
* remu hw test with the asm dsl

* simpler

* nthreads and exec mask

* cmp/cmpx

* assembler error in s_mov_b32

* vopd in dsl?
2025-12-28 11:32:41 +09:00
chenyuandGitHub 784b919f7f Revert "optim empty shard #13513 (#13598)" (#13855)
* Revert "optim empty shard #13513 (#13598)"

This reverts commit 76d465dbc3.

* test_arange_shrink

* update test
2025-12-27 21:10:23 -05:00
anuandGitHub 9b4de8abc7 fix beam in python 3.14+ (#13836)
* fix beam search on python 3.14

* add PickleableCount class to helpers

* change name, add test, add step

* tidy count init
2025-12-27 16:24:22 -05:00
chenyuandGitHub 0f74909ae9 clean up rearrange (#13851) 2025-12-27 11:06:10 -05:00
qazalandGitHub f6c660f7fa simplify sqtt decoder infra (#13849)
* more work

* simpler
2025-12-28 00:31:16 +09:00
Clément VerrierandGitHub ae013beab8 handle empty VECTORIZE in UOp.render() (#13847)
`UOp.render()` crashed with `IndexError: tuple index out of range` when
the UOp graph contained a `VECTORIZE` with empty `src=()`. This occurs
when reshaping to scalar shape `()`, e.g., `Tensor.ones(4).sum()`.

The bug was in the renderer's VECTORIZE pattern: `all_same(())` returns
`True` (vacuous truth), causing the code to access `x.src[0]` on an
empty tuple.

- Fix `IndexError` when calling `UOp.render()` on graphs containing
  empty `VECTORIZE` nodes.
- Add test for empty `VECTORIZE` rendering.
2025-12-27 10:09:39 -05:00
qazalandGitHub a2da61d096 use new style amd compiler in viz (#13848)
* working version, handcode gfx1100 arch

* get target from device properties

* lib in cfg test program spec
2025-12-27 23:59:30 +09:00
JINO ROHITandGitHub 1ee92003ea minor typo (#13846) 2025-12-27 09:34:57 -05:00
nimlgenandGitHub 276159cb87 system: add base_class to pci_scan_bus (#13845)
* system: add base_class to pci_scan_bus

* fix
2025-12-27 13:22:21 +03:00
Francis LataandGitHub fac137779e remove flux1 seed image (#13843) 2025-12-27 00:45:11 -05:00
qazalandGitHub f6de9095a0 switch asm tests to dsl (#13840)
* switch asm tests to dsl

* labeled basic blocks also work

* indenting for basic blocks

* allow define from star import
2025-12-27 02:15:16 +09:00
chenyuandGitHub ba922094f2 remove redudant check in disk_supports_fast_copyout (#13838) 2025-12-26 11:30:55 -05:00
George HotzandGitHub e9f2aaba2a simplify rdna3 asm (#13835)
* simplify rdna3 asm

* cleanups

* fix names

* fix tests

* fixes

* more test fixes

* type fixes

* tests pass + mypy passes

* 3.11 syntax
2025-12-26 11:21:03 -05:00
nimlgenandGitHub c44b4f9ae0 am: fix sdma warm boot (#13837) 2025-12-26 12:38:06 +03:00
George HotzandGitHub c6937fa744 more work on RDNA3 asm (#13833)
* more llvm asm tests

* roundtrip test

* work

* more handwritten

* more handwritten

* work

* tests pass

* dual mov

* all tests pass

* all tests pass fast
2025-12-25 23:28:14 -05:00
George HotzandGitHub f1111ac7de move amd compilers to new style (#13831)
* move amd compilers to new style

* simplest diff

* AMDHIPrenderer
2025-12-25 13:42:24 -05:00
George HotzandGitHub 9d94b8c6b2 python asm dsl in extra + python REMU (#13436)
* having fun with python asm dsl

* rdna3

* meh

* all in rdna3

* work

* more work

* work

* integration

* tests

* simpler

* simpler

* asm

* better

* simpler

* progress

* emu

* simpler

* emu

* tests

* types

* vopd

* cleaups

* work

* memory ranges

* add tracing

* refactors

* run_asm exit

* more readable

* compare to remu

* test gemm

* bug + stale

* more tests

* refactor

* tests fix

* more ins

* more instructions

* refactor

* faster

* match case

* match case

* simpler

* work

* tests

* run_asm

* work

* bug fixes

* more emu

* alu/emu

* refactor

* no pipeline emu yet

* alu direct

* fix

* bugfixes + new test

* fix exceptions in emulators

* update gen.py

* pylint

* no pdf

* improve bench_emu

* speedups

* cleanups

* more tests
2025-12-25 13:04:14 -05:00
nimlgenandGitHub b5f3a5ad79 am: cleanup comment (#13828) 2025-12-25 18:00:28 +03:00
chenyuandGitHub 8985a4a023 one less branch in Buffer.view [pr] (#13829) 2025-12-25 09:34:15 -05:00
chenyuandGitHub 094753b4e0 renderer arch version cleanup [pr] (#13830) 2025-12-25 09:32:56 -05:00
chenyuandGitHub 54af29dbdb trange can just be a function (#13827) 2025-12-24 23:57:10 -05:00
qazalandGitHub a1c1684b91 set .amdhsa_kernarg_size in asm test (#13826) 2025-12-25 13:08:14 +09:00
chenyuandGitHub da1cb6a9ec update llama dataloader (#13825)
separate creating dataset from itererating over the dataset to not create eval data for each eval
2025-12-24 17:42:08 -05:00
chenyuandGitHub a7fc0c288b clean up BufferCopy init [pr] (#13824) 2025-12-24 10:40:15 -05:00
chenyuandGitHub 903753c60c llama wandb logging (#13822) 2025-12-24 10:24:59 -05:00
qazalandGitHub e3a646dce3 viz: skip plaintext disassemble for cfg (#13821) 2025-12-24 23:16:59 +09:00
chenyuandGitHub cb07c5d0e8 fewer import annotations (#13819) 2025-12-23 18:45:50 -05:00
62 changed files with 24486 additions and 516 deletions
+8 -6
View File
@@ -429,13 +429,15 @@ jobs:
# LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
- name: Test speed vs theoretical
run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
- name: Test tensor cores
run: |
AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
- name: Test tensor cores AMD_LLVM=0
run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
# TODO: this is flaky
# - name: Test tensor cores AMD_LLVM=1
# run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
- name: Run Tensor Core GEMM (AMD)
run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
run: |
AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
- name: Test AMD=1
run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
#- name: Test HIP=1
+3 -4
View File
@@ -241,9 +241,8 @@ jobs:
run: |
python -m mypy --strict-equality --lineprecision-report .
cat lineprecision.txt
# broken because of UPatAny
#- name: Run TYPED=1
# run: TYPED=1 python -c "import tinygrad"
- name: Run TYPED=1
run: TYPED=1 python -c "import tinygrad"
unittest:
name: Unit Tests
@@ -642,7 +641,7 @@ jobs:
if: matrix.backend=='amdllvm'
run: python test/device/test_amd_llvm.py
- name: Run pytest (amd)
run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
- name: Run pytest (amd)
run: python -m pytest test/external/external_test_am.py --durations=20
- name: Run TRANSCENDENTAL math
Binary file not shown.

Before

Width:  |  Height:  |  Size: 286 KiB

+16 -38
View File
@@ -763,48 +763,26 @@ class BlendedGPTDataset:
return dataset_idx, dataset_sample_idx
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
if small:
if val:
return BlendedGPTDataset(
[base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
return BlendedGPTDataset(
[base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
if val:
dataset = BlendedGPTDataset([
base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
], [
1.0
], samples, seqlen, seed, False)
else:
dataset = BlendedGPTDataset([
base_dir / "c4-train.en_6_text_document",
base_dir / "c4-train.en_7_text_document",
], [
1.0, 1.0
], samples, seqlen, seed, True)
return BlendedGPTDataset(
[base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
return BlendedGPTDataset(
[base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
for b in range(math.ceil(samples / bs)):
batch = []
for i in range(bs):
tokens = dataset.get(b * bs + i)
batch.append(tokens)
def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
for b in range(math.ceil(dataset.samples / bs)):
batch = [dataset.get(b * bs + i) for i in range(bs)]
yield Tensor.stack(batch, dim=0)
def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
if val:
dataset = BlendedGPTDataset([
base_dir / "c4-validation-91205-samples.en_text_document",
], [
1.0
], samples, seqlen, seed, False)
else:
dataset = BlendedGPTDataset([
base_dir / "c4-train.en_6_text_document",
], [
1.0
], samples, seqlen, seed, True)
for b in range(math.ceil(samples / bs)):
batch = []
for i in range(bs):
tokens = dataset.get(b * bs + i)
batch.append(tokens)
yield Tensor.stack(batch, dim=0)
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
if __name__ == "__main__":
def load_unet3d(val):
+3 -6
View File
@@ -234,12 +234,9 @@ def eval_llama3():
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float()
if SMALL:
from examples.mlperf.dataloader import batch_load_llama3_small
iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
else:
from examples.mlperf.dataloader import batch_load_llama3
iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
iter = iterate_llama3_dataset(eval_dataset, BS)
losses = []
for tokens in tqdm(iter, total=5760//BS):
+28 -17
View File
@@ -1314,6 +1314,13 @@ def train_llama3():
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
opt_end_learning_rate = getenv("END_LR", 8e-7)
# ** init wandb **
WANDB = getenv("WANDB")
if WANDB:
import wandb
wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")
model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
# vocab_size from the mixtral tokenizer
if not SMALL: model_params |= {"vocab_size": 32000}
@@ -1417,23 +1424,20 @@ def train_llama3():
if getenv("FAKEDATA", 0):
return fake_data(BS, SAMPLES)
else:
if SMALL:
from examples.mlperf.dataloader import batch_load_llama3_small
return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
if getenv("FAKEDATA", 0):
eval_dataset = None
else:
from examples.mlperf.dataloader import get_llama3_dataset
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
def get_eval_iter():
if getenv("FAKEDATA", 0):
if eval_dataset is None:
return fake_data(EVAL_BS, 5760)
else:
if SMALL:
from examples.mlperf.dataloader import batch_load_llama3_small
return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
from examples.mlperf.dataloader import iterate_llama3_dataset
return iterate_llama3_dataset(eval_dataset, EVAL_BS)
iter = get_train_iter()
i, sequences_seen = resume_ckpt, 0
@@ -1449,13 +1453,17 @@ def train_llama3():
sequences_seen += tokens.shape[0]
sec = time.perf_counter()-t
mem_gb = GlobalCounters.mem_used / 1e9
gflops = GlobalCounters.global_ops / 1e9 / sec
tqdm.write(
f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, "
f"{GlobalCounters.global_ops * 1e-9 / sec:9.2f} GFLOPS")
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.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
if WANDB:
wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
tqdm.write("saving checkpoint")
@@ -1481,6 +1489,9 @@ def train_llama3():
tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
if WANDB:
wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
if log_perplexity < EVAL_TARGET:
tqdm.write(f"target achieved after {sequences_seen} sequences")
if getenv("CKPT"):
+254
View File
@@ -0,0 +1,254 @@
# Pure combinational ALU functions for RDNA3 emulation
from __future__ import annotations
import struct, math
from typing import Callable
from extra.assembly.rdna3.autogen import SOP1Op, SOP2Op, SOPCOp, SOPKOp, VOP1Op, VOP2Op, VOP3Op
# Format base offsets for unified opcode space
SOP2_BASE, SOP1_BASE, SOPC_BASE, SOPK_BASE = 0x000, 0x100, 0x200, 0x300
VOP2_BASE, VOP1_BASE = 0x100, 0x180
# Float conversion helpers
_I, _f, _H, _e = struct.Struct('<I'), struct.Struct('<f'), struct.Struct('<H'), struct.Struct('<e')
def f32(i: int) -> float: return _f.unpack(_I.pack(i & 0xffffffff))[0]
def i32(f: float) -> int:
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
try: return _I.unpack(_f.pack(f))[0]
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
def f16(i: int) -> float: return _e.unpack(_H.pack(i & 0xffff))[0]
def i16(f: float) -> int:
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
try: return _H.unpack(_e.pack(f))[0]
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
def sext(v: int, b: int) -> int: return v - (1 << b) if v & (1 << (b-1)) else v
def clz(x: int) -> int: return 32 - x.bit_length() if x else 32
def cls(x: int) -> int: x &= 0xffffffff; return 31 if x in (0, 0xffffffff) else clz(~x & 0xffffffff if x >> 31 else x) - 1
def _cvt_i32_f32(v): return (0x7fffffff if v > 0 else 0x80000000) if math.isinf(v) else (0 if math.isnan(v) else max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff)
def _cvt_u32_f32(v): return (0xffffffff if v > 0 else 0) if math.isinf(v) else (0 if math.isnan(v) or v < 0 else min(0xffffffff, int(v)))
# SALU: op -> fn(s0, s1, scc_in) -> (result, scc_out)
SALU: dict[int, Callable] = {
# SOP2
SOP2_BASE + SOP2Op.S_ADD_U32: lambda a, b, scc: ((a + b) & 0xffffffff, int((a + b) >= 0x100000000)),
SOP2_BASE + SOP2Op.S_SUB_U32: lambda a, b, scc: ((a - b) & 0xffffffff, int(b > a)),
SOP2_BASE + SOP2Op.S_ADDC_U32: lambda a, b, scc: ((r := a + b + scc) & 0xffffffff, int(r >= 0x100000000)),
SOP2_BASE + SOP2Op.S_SUBB_U32: lambda a, b, scc: ((a - b - scc) & 0xffffffff, int((b + scc) > a)),
SOP2_BASE + SOP2Op.S_ADD_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 32)) & 0xffffffff, int(((a >> 31) == (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
SOP2_BASE + SOP2Op.S_SUB_I32: lambda a, b, scc: ((r := sext(a, 32) - sext(b, 32)) & 0xffffffff, int(((a >> 31) != (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
SOP2_BASE + SOP2Op.S_AND_B32: lambda a, b, scc: ((r := a & b), int(r != 0)),
SOP2_BASE + SOP2Op.S_OR_B32: lambda a, b, scc: ((r := a | b), int(r != 0)),
SOP2_BASE + SOP2Op.S_XOR_B32: lambda a, b, scc: ((r := a ^ b), int(r != 0)),
SOP2_BASE + SOP2Op.S_AND_NOT1_B32: lambda a, b, scc: ((r := a & (~b & 0xffffffff)), int(r != 0)),
SOP2_BASE + SOP2Op.S_OR_NOT1_B32: lambda a, b, scc: ((r := a | (~b & 0xffffffff)), int(r != 0)),
SOP2_BASE + SOP2Op.S_LSHL_B32: lambda a, b, scc: ((r := (a << (b & 0x1f)) & 0xffffffff), int(r != 0)),
SOP2_BASE + SOP2Op.S_LSHR_B32: lambda a, b, scc: ((r := a >> (b & 0x1f)), int(r != 0)),
SOP2_BASE + SOP2Op.S_ASHR_I32: lambda a, b, scc: ((r := sext(a, 32) >> (b & 0x1f)) & 0xffffffff, int(r != 0)),
SOP2_BASE + SOP2Op.S_MUL_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 32)) & 0xffffffff, scc),
SOP2_BASE + SOP2Op.S_MUL_HI_U32: lambda a, b, scc: (((a * b) >> 32) & 0xffffffff, scc),
SOP2_BASE + SOP2Op.S_MUL_HI_I32: lambda a, b, scc: (((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff, scc),
SOP2_BASE + SOP2Op.S_MIN_I32: lambda a, b, scc: (a, 1) if sext(a, 32) < sext(b, 32) else (b, 0),
SOP2_BASE + SOP2Op.S_MIN_U32: lambda a, b, scc: (a, 1) if a < b else (b, 0),
SOP2_BASE + SOP2Op.S_MAX_I32: lambda a, b, scc: (a, 1) if sext(a, 32) > sext(b, 32) else (b, 0),
SOP2_BASE + SOP2Op.S_MAX_U32: lambda a, b, scc: (a, 1) if a > b else (b, 0),
SOP2_BASE + SOP2Op.S_CSELECT_B32: lambda a, b, scc: (a if scc else b, scc),
SOP2_BASE + SOP2Op.S_BFE_U32: lambda a, b, scc: ((r := ((a >> (b & 0x1f)) & ((1 << ((b >> 16) & 0x7f)) - 1)) if (b >> 16) & 0x7f else 0), int(r != 0)),
SOP2_BASE + SOP2Op.S_BFE_I32: lambda a, b, scc: ((r := sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := (b >> 16) & 0x7f) else 0), int(r != 0)),
SOP2_BASE + SOP2Op.S_PACK_LL_B32_B16: lambda a, b, scc: ((a & 0xffff) | ((b & 0xffff) << 16), scc),
SOP2_BASE + SOP2Op.S_PACK_LH_B32_B16: lambda a, b, scc: ((a & 0xffff) | (b & 0xffff0000), scc),
SOP2_BASE + SOP2Op.S_PACK_HH_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | (b & 0xffff0000), scc),
SOP2_BASE + SOP2Op.S_PACK_HL_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | ((b & 0xffff) << 16), scc),
SOP2_BASE + SOP2Op.S_ADD_F32: lambda a, b, scc: (i32(f32(a) + f32(b)), scc),
SOP2_BASE + SOP2Op.S_SUB_F32: lambda a, b, scc: (i32(f32(a) - f32(b)), scc),
SOP2_BASE + SOP2Op.S_MUL_F32: lambda a, b, scc: (i32(f32(a) * f32(b)), scc),
# SOP1
SOP1_BASE + SOP1Op.S_MOV_B32: lambda a, b, scc: (a, scc),
SOP1_BASE + SOP1Op.S_NOT_B32: lambda a, b, scc: ((r := (~a) & 0xffffffff), int(r != 0)),
SOP1_BASE + SOP1Op.S_BREV_B32: lambda a, b, scc: (int(f'{a & 0xffffffff:032b}'[::-1], 2), scc),
SOP1_BASE + SOP1Op.S_CLZ_I32_U32: lambda a, b, scc: (clz(a), scc),
SOP1_BASE + SOP1Op.S_CLS_I32: lambda a, b, scc: (cls(a), scc),
SOP1_BASE + SOP1Op.S_SEXT_I32_I8: lambda a, b, scc: (sext(a & 0xff, 8) & 0xffffffff, scc),
SOP1_BASE + SOP1Op.S_SEXT_I32_I16: lambda a, b, scc: (sext(a & 0xffff, 16) & 0xffffffff, scc),
SOP1_BASE + SOP1Op.S_ABS_I32: lambda a, b, scc: ((r := abs(sext(a, 32)) & 0xffffffff), int(r != 0)),
SOP1_BASE + SOP1Op.S_CVT_F32_I32: lambda a, b, scc: (i32(float(sext(a, 32))), scc),
SOP1_BASE + SOP1Op.S_CVT_F32_U32: lambda a, b, scc: (i32(float(a)), scc),
SOP1_BASE + SOP1Op.S_CVT_I32_F32: lambda a, b, scc: (_cvt_i32_f32(f32(a)), scc),
SOP1_BASE + SOP1Op.S_CVT_U32_F32: lambda a, b, scc: (_cvt_u32_f32(f32(a)), scc),
SOP1_BASE + SOP1Op.S_CEIL_F32: lambda a, b, scc: (i32(math.ceil(f32(a))), scc),
SOP1_BASE + SOP1Op.S_FLOOR_F32: lambda a, b, scc: (i32(math.floor(f32(a))), scc),
SOP1_BASE + SOP1Op.S_TRUNC_F32: lambda a, b, scc: (i32(math.trunc(f32(a))), scc),
SOP1_BASE + SOP1Op.S_RNDNE_F32: lambda a, b, scc: (i32(round(f32(a))), scc),
SOP1_BASE + SOP1Op.S_CVT_F16_F32: lambda a, b, scc: (i16(f32(a)), scc),
SOP1_BASE + SOP1Op.S_CVT_F32_F16: lambda a, b, scc: (i32(f16(a)), scc),
# SOPC
SOPC_BASE + SOPCOp.S_CMP_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 32))),
SOPC_BASE + SOPCOp.S_CMP_EQ_U32: lambda a, b, scc: (0, int(a == b)),
SOPC_BASE + SOPCOp.S_CMP_LG_U32: lambda a, b, scc: (0, int(a != b)),
SOPC_BASE + SOPCOp.S_CMP_GT_U32: lambda a, b, scc: (0, int(a > b)),
SOPC_BASE + SOPCOp.S_CMP_GE_U32: lambda a, b, scc: (0, int(a >= b)),
SOPC_BASE + SOPCOp.S_CMP_LT_U32: lambda a, b, scc: (0, int(a < b)),
SOPC_BASE + SOPCOp.S_CMP_LE_U32: lambda a, b, scc: (0, int(a <= b)),
SOPC_BASE + SOPCOp.S_BITCMP0_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) == 0)),
SOPC_BASE + SOPCOp.S_BITCMP1_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) != 0)),
# SOPK
SOPK_BASE + SOPKOp.S_MOVK_I32: lambda a, b, scc: (sext(b, 16) & 0xffffffff, scc),
SOPK_BASE + SOPKOp.S_CMOVK_I32: lambda a, b, scc: ((sext(b, 16) & 0xffffffff) if scc else a, scc),
SOPK_BASE + SOPKOp.S_ADDK_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 16)) & 0xffffffff, int(((a >> 31) == ((b >> 15) & 1)) and ((a >> 31) != ((r >> 31) & 1)))),
SOPK_BASE + SOPKOp.S_MULK_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 16)) & 0xffffffff, scc),
SOPK_BASE + SOPKOp.S_CMPK_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 16))),
SOPK_BASE + SOPKOp.S_CMPK_EQ_U32: lambda a, b, scc: (0, int(a == (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_LG_U32: lambda a, b, scc: (0, int(a != (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_GT_U32: lambda a, b, scc: (0, int(a > (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_GE_U32: lambda a, b, scc: (0, int(a >= (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_LT_U32: lambda a, b, scc: (0, int(a < (b & 0xffff))),
SOPK_BASE + SOPKOp.S_CMPK_LE_U32: lambda a, b, scc: (0, int(a <= (b & 0xffff))),
}
# VALU: op -> fn(s0, s1, s2) -> result
VALU: dict[int, Callable] = {
# VOP2
VOP2_BASE + VOP2Op.V_ADD_F32: lambda a, b, c: i32(f32(a) + f32(b)),
VOP2_BASE + VOP2Op.V_SUB_F32: lambda a, b, c: i32(f32(a) - f32(b)),
VOP2_BASE + VOP2Op.V_SUBREV_F32: lambda a, b, c: i32(f32(b) - f32(a)),
VOP2_BASE + VOP2Op.V_MUL_F32: lambda a, b, c: i32(f32(a) * f32(b)),
VOP2_BASE + VOP2Op.V_MIN_F32: lambda a, b, c: i32(min(f32(a), f32(b))),
VOP2_BASE + VOP2Op.V_MAX_F32: lambda a, b, c: i32(max(f32(a), f32(b))),
VOP2_BASE + VOP2Op.V_ADD_NC_U32: lambda a, b, c: (a + b) & 0xffffffff,
VOP2_BASE + VOP2Op.V_SUB_NC_U32: lambda a, b, c: (a - b) & 0xffffffff,
VOP2_BASE + VOP2Op.V_SUBREV_NC_U32: lambda a, b, c: (b - a) & 0xffffffff,
VOP2_BASE + VOP2Op.V_AND_B32: lambda a, b, c: a & b,
VOP2_BASE + VOP2Op.V_OR_B32: lambda a, b, c: a | b,
VOP2_BASE + VOP2Op.V_XOR_B32: lambda a, b, c: a ^ b,
VOP2_BASE + VOP2Op.V_XNOR_B32: lambda a, b, c: (~(a ^ b)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_LSHLREV_B32: lambda a, b, c: (b << (a & 0x1f)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_LSHRREV_B32: lambda a, b, c: b >> (a & 0x1f),
VOP2_BASE + VOP2Op.V_ASHRREV_I32: lambda a, b, c: (sext(b, 32) >> (a & 0x1f)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MIN_I32: lambda a, b, c: a if sext(a, 32) < sext(b, 32) else b,
VOP2_BASE + VOP2Op.V_MAX_I32: lambda a, b, c: a if sext(a, 32) > sext(b, 32) else b,
VOP2_BASE + VOP2Op.V_MIN_U32: lambda a, b, c: min(a, b),
VOP2_BASE + VOP2Op.V_MAX_U32: lambda a, b, c: max(a, b),
VOP2_BASE + VOP2Op.V_MUL_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MUL_HI_I32_I24: lambda a, b, c: ((sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) >> 32) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MUL_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff)) & 0xffffffff,
VOP2_BASE + VOP2Op.V_MUL_HI_U32_U24: lambda a, b, c: (((a & 0xffffff) * (b & 0xffffff)) >> 32) & 0xffffffff,
VOP2_BASE + VOP2Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
VOP2_BASE + VOP2Op.V_LDEXP_F16: lambda a, b, c: i16(math.ldexp(f16(a), sext(b, 32))),
VOP2_BASE + VOP2Op.V_ADD_F16: lambda a, b, c: i16(f16(a) + f16(b)),
VOP2_BASE + VOP2Op.V_SUB_F16: lambda a, b, c: i16(f16(a) - f16(b)),
VOP2_BASE + VOP2Op.V_MUL_F16: lambda a, b, c: i16(f16(a) * f16(b)),
VOP2_BASE + VOP2Op.V_MIN_F16: lambda a, b, c: i16(min(f16(a), f16(b))),
VOP2_BASE + VOP2Op.V_MAX_F16: lambda a, b, c: i16(max(f16(a), f16(b))),
# VOP1
VOP1_BASE + VOP1Op.V_MOV_B32: lambda a, b, c: a,
VOP1_BASE + VOP1Op.V_NOT_B32: lambda a, b, c: (~a) & 0xffffffff,
VOP1_BASE + VOP1Op.V_BFREV_B32: lambda a, b, c: int(f'{a & 0xffffffff:032b}'[::-1], 2),
VOP1_BASE + VOP1Op.V_CLZ_I32_U32: lambda a, b, c: clz(a),
VOP1_BASE + VOP1Op.V_CLS_I32: lambda a, b, c: cls(a),
VOP1_BASE + VOP1Op.V_CVT_F32_I32: lambda a, b, c: i32(float(sext(a, 32))),
VOP1_BASE + VOP1Op.V_CVT_F32_U32: lambda a, b, c: i32(float(a)),
VOP1_BASE + VOP1Op.V_CVT_I32_F32: lambda a, b, c: _cvt_i32_f32(f32(a)),
VOP1_BASE + VOP1Op.V_CVT_U32_F32: lambda a, b, c: _cvt_u32_f32(f32(a)),
VOP1_BASE + VOP1Op.V_CVT_F16_F32: lambda a, b, c: i16(f32(a)),
VOP1_BASE + VOP1Op.V_CVT_F32_F16: lambda a, b, c: i32(f16(a)),
VOP1_BASE + VOP1Op.V_RCP_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
VOP1_BASE + VOP1Op.V_RCP_IFLAG_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
VOP1_BASE + VOP1Op.V_RSQ_F32: lambda a, b, c: i32(1.0 / math.sqrt(f32(a)) if f32(a) > 0 else (float('nan') if f32(a) < 0 else float('inf'))),
VOP1_BASE + VOP1Op.V_SQRT_F32: lambda a, b, c: i32(math.sqrt(f32(a)) if f32(a) >= 0 else float('nan')),
VOP1_BASE + VOP1Op.V_LOG_F32: lambda a, b, c: i32(math.log2(f32(a)) if f32(a) > 0 else (float('-inf') if f32(a) == 0 else float('nan'))),
VOP1_BASE + VOP1Op.V_EXP_F32: lambda a, b, c: i32(float('inf') if f32(a) > 128 else (0.0 if f32(a) < -150 else math.pow(2.0, f32(a)))),
VOP1_BASE + VOP1Op.V_SIN_F32: lambda a, b, c: i32(math.sin(f32(a) * 2 * math.pi)),
VOP1_BASE + VOP1Op.V_COS_F32: lambda a, b, c: i32(math.cos(f32(a) * 2 * math.pi)),
VOP1_BASE + VOP1Op.V_FLOOR_F32: lambda a, b, c: i32(math.floor(f32(a))),
VOP1_BASE + VOP1Op.V_CEIL_F32: lambda a, b, c: i32(math.ceil(f32(a))),
VOP1_BASE + VOP1Op.V_TRUNC_F32: lambda a, b, c: i32(math.trunc(f32(a))),
VOP1_BASE + VOP1Op.V_RNDNE_F32: lambda a, b, c: i32(round(f32(a))),
VOP1_BASE + VOP1Op.V_FRACT_F32: lambda a, b, c: i32((v := f32(a)) - math.floor(v)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE0: lambda a, b, c: i32(float(a & 0xff)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE1: lambda a, b, c: i32(float((a >> 8) & 0xff)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE2: lambda a, b, c: i32(float((a >> 16) & 0xff)),
VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE3: lambda a, b, c: i32(float((a >> 24) & 0xff)),
VOP1_BASE + VOP1Op.V_FREXP_MANT_F32: lambda a, b, c: i32(math.frexp(v)[0] if (v := f32(a)) != 0 else 0.0),
VOP1_BASE + VOP1Op.V_FREXP_EXP_I32_F32: lambda a, b, c: (math.frexp(v)[1] if (v := f32(a)) != 0 else 0) & 0xffffffff,
# VOP3
VOP3Op.V_FMA_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
VOP3Op.V_DIV_FMAS_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
VOP3Op.V_ADD3_U32: lambda a, b, c: (a + b + c) & 0xffffffff,
VOP3Op.V_LSHL_ADD_U32: lambda a, b, c: ((a << (b & 0x1f)) + c) & 0xffffffff,
VOP3Op.V_ADD_LSHL_U32: lambda a, b, c: ((a + b) << (c & 0x1f)) & 0xffffffff,
VOP3Op.V_XOR3_B32: lambda a, b, c: a ^ b ^ c,
VOP3Op.V_OR3_B32: lambda a, b, c: a | b | c,
VOP3Op.V_AND_OR_B32: lambda a, b, c: (a & b) | c,
VOP3Op.V_LSHL_OR_B32: lambda a, b, c: ((a << (b & 0x1f)) | c) & 0xffffffff,
VOP3Op.V_XAD_U32: lambda a, b, c: ((a ^ b) + c) & 0xffffffff,
VOP3Op.V_MAD_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff) + c) & 0xffffffff,
VOP3Op.V_MAD_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24) + sext(c, 32)) & 0xffffffff,
VOP3Op.V_BFE_U32: lambda a, b, c: (a >> (b & 0x1f)) & ((1 << (c & 0x1f)) - 1) if c & 0x1f else 0,
VOP3Op.V_BFE_I32: lambda a, b, c: sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := c & 0x1f) else 0,
VOP3Op.V_ALIGNBIT_B32: lambda a, b, c: (((a << 32) | b) >> (c & 0x1f)) & 0xffffffff,
VOP3Op.V_MUL_LO_U32: lambda a, b, c: (a * b) & 0xffffffff,
VOP3Op.V_MUL_HI_U32: lambda a, b, c: ((a * b) >> 32) & 0xffffffff,
VOP3Op.V_MUL_HI_I32: lambda a, b, c: ((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff,
VOP3Op.V_LDEXP_F32: lambda a, b, c: i32(math.ldexp(f32(a), sext(b, 32))),
VOP3Op.V_DIV_FIXUP_F32: lambda a, b, c: i32(math.copysign(float('inf'), f32(c)) if f32(b) == 0.0 else f32(c) / f32(b)),
VOP3Op.V_PACK_B32_F16: lambda a, b, c: (a & 0xffff) | ((b & 0xffff) << 16),
VOP3Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
VOP3Op.V_LSHLREV_B16: lambda a, b, c: ((b & 0xffff) << (a & 0xf)) & 0xffff,
VOP3Op.V_LSHRREV_B16: lambda a, b, c: (b & 0xffff) >> (a & 0xf),
VOP3Op.V_ASHRREV_I16: lambda a, b, c: (sext(b & 0xffff, 16) >> (a & 0xf)) & 0xffff,
VOP3Op.V_ADD_NC_U16: lambda a, b, c: ((a & 0xffff) + (b & 0xffff)) & 0xffff,
VOP3Op.V_SUB_NC_U16: lambda a, b, c: ((a & 0xffff) - (b & 0xffff)) & 0xffff,
VOP3Op.V_MUL_LO_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff)) & 0xffff,
VOP3Op.V_MIN_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff),
VOP3Op.V_MAX_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff),
VOP3Op.V_MIN_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) < sext(b & 0xffff, 16) else (b & 0xffff),
VOP3Op.V_MAX_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) > sext(b & 0xffff, 16) else (b & 0xffff),
VOP3Op.V_MAD_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff) + (c & 0xffff)) & 0xffff,
VOP3Op.V_MAD_I16: lambda a, b, c: (sext(a & 0xffff, 16) * sext(b & 0xffff, 16) + sext(c & 0xffff, 16)) & 0xffff,
VOP3Op.V_FMA_F16: lambda a, b, c: i16(f16(a) * f16(b) + f16(c)),
VOP3Op.V_MIN3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[0] & 0xffffffff,
VOP3Op.V_MAX3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[2] & 0xffffffff,
VOP3Op.V_MED3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[1] & 0xffffffff,
VOP3Op.V_MIN3_F16: lambda a, b, c: i16(min(f16(a), f16(b), f16(c))),
VOP3Op.V_MAX3_F16: lambda a, b, c: i16(max(f16(a), f16(b), f16(c))),
VOP3Op.V_MED3_F16: lambda a, b, c: i16(sorted([f16(a), f16(b), f16(c)])[1]),
VOP3Op.V_MIN3_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff, c & 0xffff),
VOP3Op.V_MAX3_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff, c & 0xffff),
VOP3Op.V_MED3_U16: lambda a, b, c: sorted([a & 0xffff, b & 0xffff, c & 0xffff])[1],
VOP3Op.V_MIN3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[0] & 0xffff,
VOP3Op.V_MAX3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[2] & 0xffff,
VOP3Op.V_MED3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[1] & 0xffff,
}
def _cmp8(a, b): return [False, a < b, a == b, a <= b, a > b, a != b, a >= b, True]
def _cmp6(a, b): return [a < b, a == b, a <= b, a > b, a != b, a >= b]
def vopc(op: int, s0: int, s1: int, s0_hi: int = 0, s1_hi: int = 0) -> int:
base = op & 0x7f
if 16 <= base <= 31: # F32
f0, f1, cmp, nan = f32(s0), f32(s1), base - 16, math.isnan(f32(s0)) or math.isnan(f32(s1))
return int([False, f0<f1, f0==f1, f0<=f1, f0>f1, f0!=f1, f0>=f1, not nan, nan, f0<f1 or nan, f0==f1 or nan, f0<=f1 or nan, f0>f1 or nan, f0!=f1 or nan, f0>=f1 or nan, True][cmp])
if 49 <= base <= 54: return int(_cmp6(sext(s0 & 0xffff, 16), sext(s1 & 0xffff, 16))[base - 49]) # I16
if 57 <= base <= 62: return int(_cmp6(s0 & 0xffff, s1 & 0xffff)[base - 57]) # U16
if 64 <= base <= 79: # I32/U32
cmp = (base - 64) % 8
return int(_cmp8(sext(s0, 32), sext(s1, 32))[cmp] if base < 72 else _cmp8(s0, s1)[cmp])
if 80 <= base <= 95: # I64/U64
s0_64, s1_64 = s0 | (s0_hi << 32), s1 | (s1_hi << 32)
return int(_cmp8(sext(s0_64, 64), sext(s1_64, 64))[(base - 80) % 8] if base < 88 else _cmp8(s0_64, s1_64)[(base - 80) % 8])
if base == 126: # CLASS_F32
f, mask = f32(s0), s1
if math.isnan(f): return int(bool(mask & 0x3))
if math.isinf(f): return int(bool(mask & (0x4 if f < 0 else 0x200)))
if f == 0.0: return int(bool(mask & (0x20 if (s0 >> 31) & 1 else 0x40)))
exp, sign = (s0 >> 23) & 0xff, (s0 >> 31) & 1
return int(bool(mask & ((0x10 if sign else 0x80) if exp == 0 else (0x8 if sign else 0x100))))
raise NotImplementedError(f"VOPC op {op} (base {base})")
+602
View File
@@ -0,0 +1,602 @@
# RDNA3 assembler and disassembler
from __future__ import annotations
import re
from extra.assembly.rdna3.lib import Inst, RawImm, Reg, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
# Decoding helpers
SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
SPECIAL_PAIRS = {106: "vcc", 126: "exec"} # Special register pairs (for 64-bit ops)
# GFX11 hwreg names (IDs 16-17 are TBA - not supported, IDs 18-19 are PERF_SNAPSHOT)
HWREG_NAMES = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI',
22: 'HW_REG_XNACK_MASK', 23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
HWREG_IDS = {v.lower(): k for k, v in HWREG_NAMES.items()} # Reverse map for assembler
MSG_NAMES = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
_16BIT_TYPES = ('f16', 'i16', 'u16', 'b16')
def _is_16bit(s: str) -> bool: return any(s.endswith(x) for x in _16BIT_TYPES)
def decode_src(val: int) -> str:
if val <= 105: return f"s{val}"
if val in SPECIAL_DEC: return SPECIAL_DEC[val]
if 108 <= val <= 123: return f"ttmp{val - 108}"
if 128 <= val <= 192: return str(val - 128)
if 193 <= val <= 208: return str(-(val - 192))
if 256 <= val <= 511: return f"v{val - 256}"
return "lit" if val == 255 else f"?{val}"
def _reg(prefix: str, base: int, cnt: int = 1) -> str: return f"{prefix}{base}" if cnt == 1 else f"{prefix}[{base}:{base+cnt-1}]"
def _sreg(base: int, cnt: int = 1) -> str: return _reg("s", base, cnt)
def _vreg(base: int, cnt: int = 1) -> str: return _reg("v", base, cnt)
def _fmt_sdst(v: int, cnt: int = 1) -> str:
"""Format SGPR destination with special register names."""
if v == 124: return "null"
if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
if cnt > 1 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if cnt > 1: return _sreg(v, cnt)
return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
def _fmt_ssrc(v: int, cnt: int = 1) -> str:
"""Format SGPR source with special register names and pairs."""
if cnt == 2:
if v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if v <= 105: return _sreg(v, 2)
if 108 <= v <= 123: return _reg("ttmp", v - 108, 2)
return decode_src(v)
def _fmt_src_n(v: int, cnt: int) -> str:
"""Format source with given register count (1, 2, or 4)."""
if cnt == 1: return decode_src(v)
if v >= 256: return _vreg(v - 256, cnt)
if v <= 105: return _sreg(v, cnt)
if cnt == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
return decode_src(v)
def _fmt_src64(v: int) -> str:
"""Format 64-bit source (VGPR pair, SGPR pair, or special pair)."""
return _fmt_src_n(v, 2)
def _parse_sop_sizes(op_name: str) -> tuple[int, ...]:
"""Parse dst and src sizes from SOP instruction name. Returns (dst_cnt, src0_cnt) or (dst_cnt, src0_cnt, src1_cnt)."""
if op_name in ('s_bitset0_b64', 's_bitset1_b64'): return (2, 1)
if op_name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return (2, 2, 1)
if op_name in ('s_bfm_b64',): return (2, 1, 1)
# SOPC: s_bitcmp0_b64, s_bitcmp1_b64 - 64-bit src0, 32-bit src1 (bit index)
if op_name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return (1, 2, 1)
if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', op_name):
return (2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1)
if m := re.search(r'_(b|i|u)(32|64)$', op_name):
sz = 2 if m.group(2) == '64' else 1
return (sz, sz)
return (1, 1)
# Waitcnt helpers (RDNA3 format: bits 15:10=vmcnt, bits 9:4=lgkmcnt, bits 3:0=expcnt)
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
def decode_waitcnt(val: int) -> tuple[int, int, int]:
return (val >> 10) & 0x3f, val & 0xf, (val >> 4) & 0x3f # vmcnt, expcnt, lgkmcnt
# VOP3SD opcodes (shared encoding with VOP3 but different field layout)
# Note: opcodes 0-255 are VOPC promoted to VOP3 - never treat as VOP3SD
VOP3SD_OPCODES = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
# Disassembler
def disasm(inst: Inst) -> str:
op_val = unwrap(inst._values.get('op', 0))
cls_name = inst.__class__.__name__
# VOP3 and VOP3SD share encoding - check opcode to determine which
is_vop3sd = cls_name == 'VOP3' and op_val in VOP3SD_OPCODES
try:
from extra.assembly.rdna3 import autogen
if is_vop3sd:
op_name = autogen.VOP3SDOp(op_val).name.lower()
else:
op_name = getattr(autogen, f"{cls_name}Op")(op_val).name.lower() if hasattr(autogen, f"{cls_name}Op") else f"op_{op_val}"
except (ValueError, KeyError): op_name = f"op_{op_val}"
def fmt_src(v): return f"0x{inst._literal:x}" if v == 255 and getattr(inst, '_literal', None) else decode_src(v)
# VOP1
if cls_name == 'VOP1':
vdst, src0 = unwrap(inst._values['vdst']), unwrap(inst._values['src0'])
if op_name == 'v_nop': return 'v_nop'
if op_name == 'v_pipeflush': return 'v_pipeflush'
parts = op_name.split('_')
is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name
_F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
is_f64_src = op_name in _F64_OPS or op_name in ('v_cvt_f32_f64', 'v_cvt_i32_f64', 'v_cvt_u32_f64', 'v_frexp_exp_i32_f64')
if op_name == 'v_readfirstlane_b32':
return f"v_readfirstlane_b32 {decode_src(vdst)}, v{src0 - 256 if src0 >= 256 else src0}"
dst_str = _vreg(vdst, 2) if is_f64_dst else f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}" if is_16bit_dst else f"v{vdst}"
src_str = _fmt_src64(src0) if is_f64_src else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit_src and src0 >= 256 else fmt_src(src0)
return f"{op_name}_e32 {dst_str}, {src_str}"
# VOP2
if cls_name == 'VOP2':
vdst, src0_raw, vsrc1 = unwrap(inst._values['vdst']), unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
suffix = "" if op_name == "v_dot2acc_f32_f16" else "_e32"
is_16bit_op = ('_f16' in op_name or '_i16' in op_name or '_u16' in op_name) and '_f32' not in op_name and '_i32' not in op_name and 'pk_' not in op_name
if is_16bit_op:
dst_str = f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}"
src0_str = f"v{(src0_raw - 256) & 0x7f}.{'h' if src0_raw >= 384 else 'l'}" if src0_raw >= 256 else fmt_src(src0_raw)
vsrc1_str = f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}"
else:
dst_str, src0_str, vsrc1_str = f"v{vdst}", fmt_src(src0_raw), f"v{vsrc1}"
return f"{op_name}{suffix} {dst_str}, {src0_str}, {vsrc1_str}" + (", vcc_lo" if op_name == "v_cndmask_b32" else "")
# VOPC
if cls_name == 'VOPC':
src0, vsrc1 = unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
is_64bit = any(x in op_name for x in ('f64', 'i64', 'u64'))
is_64bit_vsrc1 = is_64bit and 'class' not in op_name
is_16bit = any(x in op_name for x in ('_f16', '_i16', '_u16')) and 'f32' not in op_name
is_cmpx = op_name.startswith('v_cmpx') # VOPCX writes to exec, no vcc destination
src0_str = _fmt_src64(src0) if is_64bit else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit and src0 >= 256 else fmt_src(src0)
vsrc1_str = _vreg(vsrc1, 2) if is_64bit_vsrc1 else f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}" if is_16bit else f"v{vsrc1}"
return f"{op_name}_e32 {src0_str}, {vsrc1_str}" if is_cmpx else f"{op_name}_e32 vcc_lo, {src0_str}, {vsrc1_str}"
# SOPP
if cls_name == 'SOPP':
simm16 = unwrap(inst._values.get('simm16', 0))
# No-operand instructions (simm16 is ignored)
no_imm_ops = ('s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_ttracedata_imm',
's_wait_idle', 's_endpgm_saved', 's_code_end', 's_endpgm_ordered_ps_done')
if op_name in no_imm_ops: return op_name
if op_name == 's_waitcnt':
vmcnt, expcnt, lgkmcnt = decode_waitcnt(simm16)
parts = []
if vmcnt != 0x3f: parts.append(f"vmcnt({vmcnt})")
if expcnt != 0x7: parts.append(f"expcnt({expcnt})")
if lgkmcnt != 0x3f: parts.append(f"lgkmcnt({lgkmcnt})")
return f"s_waitcnt {' '.join(parts)}" if parts else "s_waitcnt 0"
if op_name == 's_delay_alu':
dep_names = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
skip_names = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = simm16 & 0xf, (simm16 >> 4) & 0x7, (simm16 >> 7) & 0xf
def dep_name(v): return dep_names[v-1] if 0 < v <= len(dep_names) else str(v)
parts = [f"instid0({dep_name(id0)})"] if id0 else []
if skip: parts.append(f"instskip({skip_names[skip]})")
if id1: parts.append(f"instid1({dep_name(id1)})")
return f"s_delay_alu {' | '.join(p for p in parts if p)}" if parts else "s_delay_alu 0"
if op_name.startswith('s_cbranch') or op_name.startswith('s_branch'):
return f"{op_name} {simm16}"
# Most SOPP ops require immediate (s_nop, s_setkill, s_sethalt, s_sleep, s_setprio, s_sendmsg*, etc.)
return f"{op_name} 0x{simm16:x}"
# SMEM
if cls_name == 'SMEM':
if op_name in ('s_gl1_inv', 's_dcache_inv'): return op_name
sdata, sbase, soffset, offset = unwrap(inst._values['sdata']), unwrap(inst._values['sbase']), unwrap(inst._values['soffset']), unwrap(inst._values.get('offset', 0))
glc, dlc = unwrap(inst._values.get('glc', 0)), unwrap(inst._values.get('dlc', 0))
# Format offset: "soffset offset:X" if both, "0x{offset:x}" if only imm, or decode_src(soffset)
off_str = f"{decode_src(soffset)} offset:0x{offset:x}" if offset and soffset != 124 else f"0x{offset:x}" if offset else decode_src(soffset)
sbase_idx, sbase_cnt = sbase * 2, 4 if (8 <= op_val <= 12 or op_name == 's_atc_probe_buffer') else 2
sbase_str = _fmt_ssrc(sbase_idx, sbase_cnt) if sbase_cnt == 2 else _sreg(sbase_idx, sbase_cnt) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_cnt)
if op_name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{op_name} {sdata}, {sbase_str}, {off_str}"
width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val, 1)
mods = [m for m in ["glc" if glc else "", "dlc" if dlc else ""] if m]
return f"{op_name} {_fmt_sdst(sdata, width)}, {sbase_str}, {off_str}" + (" " + " ".join(mods) if mods else "")
# FLAT
if cls_name == 'FLAT':
vdst, addr, data, saddr, offset, seg = [unwrap(inst._values.get(f, 0)) for f in ['vdst', 'addr', 'data', 'saddr', 'offset', 'seg']]
instr = f"{['flat', 'scratch', 'global'][seg] if seg < 3 else 'flat'}_{op_name.split('_', 1)[1] if '_' in op_name else op_name}"
width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'u8':1, 'i8':1, 'u16':1, 'i16':1}.get(op_name.split('_')[-1], 1)
addr_str = _vreg(addr, 2) if saddr == 0x7F else _vreg(addr)
saddr_str = "" if saddr == 0x7F else f", {_sreg(saddr, 2)}" if saddr < 106 else ", off" if saddr == 124 else f", {decode_src(saddr)}"
off_str = f" offset:{offset}" if offset else ""
vdata_str = _vreg(data if 'store' in op_name else vdst, width)
return f"{instr} {addr_str}, {vdata_str}{saddr_str}{off_str}" if 'store' in op_name else f"{instr} {vdata_str}, {addr_str}{saddr_str}{off_str}"
# VOP3: vector ops with modifiers (can be 1, 2, or 3 sources depending on opcode range)
if cls_name == 'VOP3':
# Handle VOP3SD opcodes (same encoding, different field layout)
if is_vop3sd:
vdst = unwrap(inst._values.get('vdst', 0))
# VOP3SD: sdst is at bits [14:8], but VOP3 decodes opsel at [14:11], abs at [10:8], clmp at [15]
# We need to reconstruct sdst from these fields
opsel_raw = unwrap(inst._values.get('opsel', 0))
abs_raw = unwrap(inst._values.get('abs', 0))
clmp_raw = unwrap(inst._values.get('clmp', 0))
sdst = (clmp_raw << 7) | (opsel_raw << 3) | abs_raw
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg = unwrap(inst._values.get('neg', 0))
omod = unwrap(inst._values.get('omod', 0))
omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
is_f64 = 'f64' in op_name
# v_mad_i64_i32/v_mad_u64_u32: 64-bit dst and src2, 32-bit src0/src1
is_mad64 = 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
def fmt_sd_src(v, neg_bit, is_64bit=False):
s = _fmt_src64(v) if (is_64bit or is_f64) else fmt_src(v)
return f"-{s}" if neg_bit else s
src0_str, src1_str = fmt_sd_src(src0, neg & 1), fmt_sd_src(src1, neg & 2)
src2_str = fmt_sd_src(src2, neg & 4, is_mad64)
dst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}"
sdst_str = _fmt_sdst(sdst, 1)
# v_add_co_u32, v_sub_co_u32, v_subrev_co_u32, v_add_co_ci_u32, etc. only use 2 sources
if op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32', 'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'):
return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}"
# v_div_scale uses 3 sources
return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}, {src2_str}" + omod_str
vdst = unwrap(inst._values.get('vdst', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, abs_, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('abs', 0)), unwrap(inst._values.get('clmp', 0))
opsel = unwrap(inst._values.get('opsel', 0))
# Check if 64-bit op (needs register pairs)
is_f64 = 'f64' in op_name or 'i64' in op_name or 'u64' in op_name or 'b64' in op_name
# v_cmp_class_* has 64-bit src0 but 32-bit src1 (class mask)
is_class = 'class' in op_name
# Shift ops: v_*rev_*64 have 32-bit shift amount (src0), 64-bit value (src1)
is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_')
# v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent)
is_ldexp64 = op_name == 'v_ldexp_f64'
# v_trig_preop_f64: 64-bit dst/src0, 32-bit src1 (exponent/scale)
is_trig_preop = op_name == 'v_trig_preop_f64'
# v_readlane_b32: destination is SGPR (despite vdst field)
is_readlane = op_name == 'v_readlane_b32'
# SAD/QSAD/MQSAD instructions have mixed sizes
# v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1
# v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1
is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk'))
is_mqsad_u32 = 'mqsad_u32' in op_name
# Detect 16-bit and 64-bit operand sizes for various instruction patterns
if 'cvt_pk' in op_name:
is_f16_dst, is_f16_src, is_f16_src2 = False, op_name.endswith('16'), False
elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', op_name):
dst_type, src_type = m.group(1), m.group(2)
is_f16_dst, is_f16_src, is_f16_src2 = _is_16bit(dst_type), _is_16bit(src_type), _is_16bit(src_type)
is_f64_dst, is_f64_src, is_f64 = '64' in dst_type, '64' in src_type, False
elif re.match(r'v_mad_[iu]32_[iu]16', op_name):
is_f16_dst, is_f16_src, is_f16_src2 = False, True, False # 32-bit dst, 16-bit src0/src1, 32-bit src2
elif 'pack_b32' in op_name:
is_f16_dst, is_f16_src, is_f16_src2 = False, True, True # 32-bit dst, 16-bit sources
else:
is_16bit_op = any(x in op_name for x in _16BIT_TYPES) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad'))
is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op
# Check if any opsel bit is set (any operand uses .h) - if so, we need explicit .l for low-half
any_hi = opsel != 0
def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False):
s = _fmt_src_n(v, reg_cnt) if reg_cnt > 1 else f"v{v - 256}.h" if is_16 and v >= 256 and hi_bit else f"v{v - 256}.l" if is_16 and v >= 256 and any_hi else fmt_src(v)
if abs_bit: s = f"|{s}|"
return f"-{s}" if neg_bit else s
# Determine register count for each source (check for cvt-specific 64-bit flags first)
is_src0_64 = locals().get('is_f64_src', is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32
is_src1_64 = is_f64 and not is_class and not is_ldexp64 and not is_trig_preop
src0_cnt = 2 if is_src0_64 else 1
src1_cnt = 2 if is_src1_64 else 1
src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1
src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src)
src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src)
src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2)
# Format destination - for 16-bit ops, use .h/.l suffix; readlane uses SGPR dest
is_dst_64 = locals().get('is_f64_dst', is_f64) or is_sad64
dst_cnt = 4 if is_mqsad_u32 else 2 if is_dst_64 else 1
if is_readlane:
dst_str = _fmt_sdst(vdst, 1)
elif dst_cnt > 1:
dst_str = _vreg(vdst, dst_cnt)
elif is_f16_dst:
dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l" if any_hi else f"v{vdst}"
else:
dst_str = f"v{vdst}"
clamp_str = " clamp" if clmp else ""
omod = unwrap(inst._values.get('omod', 0))
omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
# op_sel for non-VGPR sources (when opsel bits are set but source is not a VGPR)
# For 16-bit ops with VGPR sources, opsel is encoded in .h/.l suffix
# For non-VGPR sources or non-16-bit ops, we need explicit op_sel
has_nonvgpr_opsel = (src0 < 256 and (opsel & 1)) or (src1 < 256 and (opsel & 2)) or (src2 < 256 and (opsel & 4))
need_opsel = has_nonvgpr_opsel or (opsel and not is_f16_src)
# Helper to format opsel string based on source count
def fmt_opsel(num_src):
if not need_opsel: return ""
# When dst is .h (for 16-bit ops) and non-VGPR sources have opsel, use all 1s
if is_f16_dst and (opsel & 8): # dst is .h
return f" op_sel:[1,1,1{',1' if num_src == 3 else ''}]"
# Otherwise output actual opsel values
if num_src == 3:
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
# Determine number of sources based on opcode range:
# 0-255: VOPC promoted (comparison, 2 src, sdst)
# 256-383: VOP2 promoted (2 src)
# 384-511: VOP1 promoted (1 src)
# 512+: Native VOP3 (2 or 3 src depending on instruction)
if op_val < 256: # VOPC promoted
# VOPCX (v_cmpx_*) writes to exec, no explicit destination
if op_name.startswith('v_cmpx'):
return f"{op_name}_e64 {src0_str}, {src1_str}"
return f"{op_name}_e64 {_fmt_sdst(vdst, 1)}, {src0_str}, {src1_str}"
elif op_val < 384: # VOP2 promoted
# v_cndmask_b32 in VOP3 format has 3 sources (src2 is mask selector)
if 'cndmask' in op_name:
return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
elif op_val < 512: # VOP1 promoted
if op_name in ('v_nop', 'v_pipeflush'): return f"{op_name}_e64"
return f"{op_name}_e64 {dst_str}, {src0_str}" + fmt_opsel(1) + clamp_str + omod_str
else: # Native VOP3 - determine 2 vs 3 sources based on instruction name
# 3-source ops: fma, mad, min3, max3, med3, div_fixup, div_fmas, sad, msad, qsad, mqsad, lerp, alignbit/byte, cubeid/sc/tc/ma, bfe, bfi, perm_b32, permlane, cndmask
# Note: v_writelane_b32 is 2-src (src0, src1 with vdst as 3rd operand - read-modify-write)
is_3src = any(x in op_name for x in ('fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube',
'bfe', 'bfi', 'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add',
'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit'))
if is_3src:
return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
return f"{op_name} {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
# VOP3SD: 3-source with scalar destination (v_div_scale_*, v_add_co_u32, v_mad_*64_*32, etc.)
if cls_name == 'VOP3SD':
vdst, sdst = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('sdst', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, omod, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('omod', 0)), unwrap(inst._values.get('clmp', 0))
is_f64, is_mad64 = 'f64' in op_name, 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
def fmt_neg(v, neg_bit, is_64=False): return f"-{_fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)}" if neg_bit else _fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)
srcs = [fmt_neg(src0, neg & 1), fmt_neg(src1, neg & 2), fmt_neg(src2, neg & 4, is_mad64)]
dst_str, sdst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}", _fmt_sdst(sdst, 1)
clamp_str, omod_str = " clamp" if clmp else "", {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
is_2src = op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32')
suffix = "_e64" if op_name.startswith('v_') and 'co_' in op_name else ""
return f"{op_name}{suffix} {dst_str}, {sdst_str}, {', '.join(srcs[:2] if is_2src else srcs)}" + clamp_str + omod_str
# VOPD: dual-issue instructions
if cls_name == 'VOPD':
from extra.assembly.rdna3 import autogen
opx, opy, vdstx, vdsty_enc = [unwrap(inst._values.get(f, 0)) for f in ('opx', 'opy', 'vdstx', 'vdsty')]
srcx0, vsrcx1, srcy0, vsrcy1 = [unwrap(inst._values.get(f, 0)) for f in ('srcx0', 'vsrcx1', 'srcy0', 'vsrcy1')]
vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1) # Decode vdsty
def fmt_vopd(op, vdst, src0, vsrc1):
try: name = autogen.VOPDOp(op).name.lower()
except (ValueError, KeyError): name = f"op_{op}"
return f"{name} v{vdst}, {fmt_src(src0)}" if 'mov' in name else f"{name} v{vdst}, {fmt_src(src0)}, v{vsrc1}"
return f"{fmt_vopd(opx, vdstx, srcx0, vsrcx1)} :: {fmt_vopd(opy, vdsty, srcy0, vsrcy1)}"
# VOP3P: packed vector ops
if cls_name == 'VOP3P':
vdst, clmp = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('clmp', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, neg_hi = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('neg_hi', 0))
opsel, opsel_hi, opsel_hi2 = unwrap(inst._values.get('opsel', 0)), unwrap(inst._values.get('opsel_hi', 0)), unwrap(inst._values.get('opsel_hi2', 0))
is_wmma, is_3src = 'wmma' in op_name, any(x in op_name for x in ('fma', 'mad', 'dot', 'wmma'))
def fmt_bits(name, val, n): return f"{name}:[{','.join(str((val >> i) & 1) for i in range(n))}]"
# WMMA: f16/bf16 use 8-reg sources, iu8 uses 4-reg, iu4 uses 2-reg; all have 8-reg dst
if is_wmma:
src_cnt = 2 if 'iu4' in op_name else 4 if 'iu8' in op_name else 8
src0_str, src1_str, src2_str = _fmt_src_n(src0, src_cnt), _fmt_src_n(src1, src_cnt), _fmt_src_n(src2, 8)
dst_str = _vreg(vdst, 8)
else:
src0_str, src1_str, src2_str = _fmt_src_n(src0, 1), _fmt_src_n(src1, 1), _fmt_src_n(src2, 1)
dst_str = f"v{vdst}"
n = 3 if is_3src else 2
full_opsel_hi = opsel_hi | (opsel_hi2 << 2)
mods = [fmt_bits("op_sel", opsel, n)] if opsel else []
if full_opsel_hi != (0b111 if is_3src else 0b11): mods.append(fmt_bits("op_sel_hi", full_opsel_hi, n))
if neg: mods.append(fmt_bits("neg_lo", neg, n))
if neg_hi: mods.append(fmt_bits("neg_hi", neg_hi, n))
if clmp: mods.append("clamp")
mod_str = " " + " ".join(mods) if mods else ""
return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}{mod_str}" if is_3src else f"{op_name} {dst_str}, {src0_str}, {src1_str}{mod_str}"
# VINTERP: interpolation instructions
if cls_name == 'VINTERP':
vdst = unwrap(inst._values.get('vdst', 0))
src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
neg, waitexp, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('waitexp', 0)), unwrap(inst._values.get('clmp', 0))
def fmt_neg_vi(v, neg_bit): return f"-{v}" if neg_bit else v
srcs = [fmt_neg_vi(f"v{s - 256}" if s >= 256 else fmt_src(s), neg & (1 << i)) for i, s in enumerate([src0, src1, src2])]
mods = [m for m in [f"wait_exp:{waitexp}" if waitexp else "", "clamp" if clmp else ""] if m]
return f"{op_name} v{vdst}, {', '.join(srcs)}" + (" " + " ".join(mods) if mods else "")
# MUBUF/MTBUF helpers
def _buf_vaddr(vaddr, offen, idxen): return _vreg(vaddr, 2) if offen and idxen else f"v{vaddr}" if offen or idxen else "off"
def _buf_srsrc(srsrc): srsrc_base = srsrc * 4; return _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
# MUBUF: buffer load/store
if cls_name == 'MUBUF':
vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
offset, offen, idxen = unwrap(inst._values.get('offset', 0)), unwrap(inst._values.get('offen', 0)), unwrap(inst._values.get('idxen', 0))
glc, dlc, slc, tfe = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc', 'tfe')]
if op_name in ('buffer_gl0_inv', 'buffer_gl1_inv'): return op_name
# Determine data width from op name
if 'd16' in op_name: width = 2 if any(x in op_name for x in ('xyz', 'xyzw')) else 1
elif 'atomic' in op_name:
base_width = 2 if any(x in op_name for x in ('b64', 'u64', 'i64')) else 1
width = base_width * 2 if 'cmpswap' in op_name else base_width
else: width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'b16':1, 'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
if tfe: width += 1
mods = [m for m in ["offen" if offen else "", "idxen" if idxen else "", f"offset:{offset}" if offset else "",
"glc" if glc else "", "dlc" if dlc else "", "slc" if slc else "", "tfe" if tfe else ""] if m]
return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)}" + (" " + " ".join(mods) if mods else "")
# MTBUF: typed buffer load/store
if cls_name == 'MTBUF':
vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
offset, tbuf_fmt, offen, idxen = [unwrap(inst._values.get(f, 0)) for f in ('offset', 'format', 'offen', 'idxen')]
glc, dlc, slc = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc')]
mods = [f"format:{tbuf_fmt}"] + [m for m in ["idxen" if idxen else "", "offen" if offen else "", f"offset:{offset}" if offset else "",
"glc" if glc else "", "dlc" if dlc else "", "slc" if slc else ""] if m]
width = 2 if 'd16' in op_name and any(x in op_name for x in ('xyz', 'xyzw')) else 1 if 'd16' in op_name else {'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)} {' '.join(mods)}"
# SOP1/SOP2/SOPC/SOPK
if cls_name in ('SOP1', 'SOP2', 'SOPC', 'SOPK'):
sizes = _parse_sop_sizes(op_name)
dst_cnt, src0_cnt = sizes[0], sizes[1]
src1_cnt = sizes[2] if len(sizes) > 2 else src0_cnt
if cls_name == 'SOP1':
sdst, ssrc0 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('ssrc0', 0))
if op_name == 's_getpc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}"
if op_name in ('s_setpc_b64', 's_rfe_b64'): return f"{op_name} {_fmt_ssrc(ssrc0, 2)}"
if op_name == 's_swappc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}, {_fmt_ssrc(ssrc0, 2)}"
if op_name in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'):
return f"{op_name} {_fmt_sdst(sdst, 2 if 'b64' in op_name else 1)}, sendmsg({MSG_NAMES.get(ssrc0, str(ssrc0))})"
ssrc0_str = fmt_src(ssrc0) if src0_cnt == 1 else _fmt_ssrc(ssrc0, src0_cnt)
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {ssrc0_str}"
if cls_name == 'SOP2':
sdst, ssrc0, ssrc1 = [unwrap(inst._values.get(f, 0)) for f in ('sdst', 'ssrc0', 'ssrc1')]
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {_fmt_ssrc(ssrc0, src0_cnt)}, {_fmt_ssrc(ssrc1, src1_cnt)}"
if cls_name == 'SOPC':
return f"{op_name} {_fmt_ssrc(unwrap(inst._values.get('ssrc0', 0)), src0_cnt)}, {_fmt_ssrc(unwrap(inst._values.get('ssrc1', 0)), src1_cnt)}"
if cls_name == 'SOPK':
sdst, simm16 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('simm16', 0))
if op_name == 's_version': return f"{op_name} 0x{simm16:x}"
if op_name in ('s_setreg_b32', 's_getreg_b32'):
hwreg_id, hwreg_offset, hwreg_size = simm16 & 0x3f, (simm16 >> 6) & 0x1f, ((simm16 >> 11) & 0x1f) + 1
hwreg_str = f"0x{simm16:x}" if hwreg_id in (16, 17) else f"hwreg({HWREG_NAMES.get(hwreg_id, str(hwreg_id))}, {hwreg_offset}, {hwreg_size})"
return f"{op_name} {hwreg_str}, {_fmt_sdst(sdst, 1)}" if op_name == 's_setreg_b32' else f"{op_name} {_fmt_sdst(sdst, 1)}, {hwreg_str}"
return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, 0x{simm16:x}"
# Generic fallback
def fmt_field(n, v):
v = unwrap(v)
if n in SRC_FIELDS: return fmt_src(v) if v != 255 else "0xff"
if n in ('sdst', 'vdst'): return f"{'s' if n == 'sdst' else 'v'}{v}"
return f"v{v}" if n == 'vsrc1' else f"0x{v:x}" if n == 'simm16' else str(v)
ops = [fmt_field(n, inst._values.get(n, 0)) for n in inst._fields if n not in ('encoding', 'op')]
return f"{op_name} {', '.join(ops)}" if ops else op_name
# Assembler
SPECIAL_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125), 'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'scc': RawImm(253)}
FLOAT_CONSTS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
def parse_operand(op: str) -> tuple:
op = op.strip().lower()
neg = op.startswith('-') and not op[1:2].isdigit(); op = op[1:] if neg else op
abs_ = op.startswith('|') and op.endswith('|') or op.startswith('abs(') and op.endswith(')')
op = op[1:-1] if op.startswith('|') else op[4:-1] if op.startswith('abs(') else op
hi_half = op.endswith('.h')
op = re.sub(r'\.[lh]$', '', op)
if op in FLOAT_CONSTS: return (FLOAT_CONSTS[op], neg, abs_, hi_half)
if re.match(r'^-?\d+$', op): return (int(op), neg, abs_, hi_half)
if m := re.match(r'^-?0x([0-9a-f]+)$', op):
v = -int(m.group(1), 16) if op.startswith('-') else int(m.group(1), 16)
return (v, neg, abs_, hi_half)
if op in SPECIAL_REGS: return (SPECIAL_REGS[op], neg, abs_, hi_half)
if op == 'lit': return (RawImm(255), neg, abs_, hi_half) # literal marker (actual value comes from literal word)
if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', op): return (REG_MAP[m.group(1)][int(m.group(2)):int(m.group(3))], neg, abs_, hi_half)
if m := re.match(r'^([svt](?:tmp)?)(\d+)$', op):
reg = REG_MAP[m.group(1)][int(m.group(2))]
reg.hi = hi_half
return (reg, neg, abs_, hi_half)
# hwreg(name, offset, size) or hwreg(name) -> simm16 encoding
if m := re.match(r'^hwreg\((\w+)(?:,\s*(\d+),\s*(\d+))?\)$', op):
name_str = m.group(1).lower()
hwreg_id = HWREG_IDS.get(name_str, int(name_str) if name_str.isdigit() else None)
if hwreg_id is None: raise ValueError(f"unknown hwreg name: {name_str}")
offset, size = int(m.group(2)) if m.group(2) else 0, int(m.group(3)) if m.group(3) else 32
return (((size - 1) << 11) | (offset << 6) | hwreg_id, neg, abs_, hi_half)
raise ValueError(f"cannot parse operand: {op}")
SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
SOP1_SRC_ONLY = {'s_setpc_b64', 's_rfe_b64'}
SOP1_MSG_IMM = {'s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'}
SOPK_IMM_ONLY = {'s_version'}
SOPK_IMM_FIRST = {'s_setreg_b32'}
SOPK_UNSUPPORTED = {'s_setreg_imm32_b32'}
def asm(text: str) -> Inst:
from extra.assembly.rdna3 import autogen
text = text.strip()
clamp = 'clamp' in text.lower()
if clamp: text = re.sub(r'\s+clamp\s*$', '', text, flags=re.I)
modifiers = {}
if m := re.search(r'\s+wait_exp:(\d+)', text, re.I): modifiers['waitexp'] = int(m.group(1)); text = text[:m.start()] + text[m.end():]
parts = text.replace(',', ' ').split()
if not parts: raise ValueError("empty instruction")
mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip()
# Handle s_waitcnt specially before operand parsing
if mnemonic == 's_waitcnt':
vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f
for part in op_str.replace(',', ' ').split():
if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1))
elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1))
elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1))
elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0))
return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt))
# Handle VOPD dual-issue instructions: opx dst, src :: opy dst, src
if '::' in text:
x_part, y_part = text.split('::')
x_parts, y_parts = x_part.strip().replace(',', ' ').split(), y_part.strip().replace(',', ' ').split()
opx_name, opy_name = x_parts[0].upper(), y_parts[0].upper()
opx, opy = autogen.VOPDOp[opx_name], autogen.VOPDOp[opy_name]
x_ops, y_ops = [parse_operand(p)[0] for p in x_parts[1:]], [parse_operand(p)[0] for p in y_parts[1:]]
vdstx, srcx0 = x_ops[0], x_ops[1] if len(x_ops) > 1 else 0
vsrcx1 = x_ops[2] if len(x_ops) > 2 else VGPR(0)
vdsty, srcy0 = y_ops[0], y_ops[1] if len(y_ops) > 1 else 0
vsrcy1 = y_ops[2] if len(y_ops) > 2 else VGPR(0)
# Handle fmaak/fmamk literals (4th operand on x or y side)
lit = None
if 'fmaak' in opx_name.lower() and len(x_ops) > 3: lit = unwrap(x_ops[3])
elif 'fmamk' in opx_name.lower() and len(x_ops) > 3: lit, vsrcx1 = unwrap(x_ops[2]), x_ops[3]
elif 'fmaak' in opy_name.lower() and len(y_ops) > 3: lit = unwrap(y_ops[3])
elif 'fmamk' in opy_name.lower() and len(y_ops) > 3: lit, vsrcy1 = unwrap(y_ops[2]), y_ops[3]
return autogen.VOPD(opx, opy, vdstx=vdstx, vdsty=vdsty, srcx0=srcx0, vsrcx1=vsrcx1, srcy0=srcy0, vsrcy1=vsrcy1, literal=lit)
operands, current, depth, in_pipe = [], "", 0, False
for ch in op_str:
if ch in '[(': depth += 1
elif ch in '])': depth -= 1
elif ch == '|': in_pipe = not in_pipe
if ch == ',' and depth == 0 and not in_pipe: operands.append(current.strip()); current = ""
else: current += ch
if current.strip(): operands.append(current.strip())
parsed = [parse_operand(op) for op in operands]
values = [p[0] for p in parsed]
neg_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[1])
abs_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[2])
opsel_bits = (8 if len(parsed) > 0 and parsed[0][3] else 0) | sum((1 << i) for i, p in enumerate(parsed[1:4]) if p[3])
lit = None
if mnemonic in ('v_fmaak_f32', 'v_fmaak_f16') and len(values) == 4: lit, values = unwrap(values[3]), values[:3]
elif mnemonic in ('v_fmamk_f32', 'v_fmamk_f16') and len(values) == 4: lit, values = unwrap(values[2]), [values[0], values[1], values[3]]
vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32', 'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}
if mnemonic.replace('_e32', '') in vcc_ops and len(values) >= 5: values = [values[0], values[2], values[3]]
if mnemonic.startswith('v_cmp') and len(values) >= 3 and operands[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'):
values = values[1:]
# CMPX instructions with _e64 suffix: prepend implicit EXEC_LO destination (vdst=126)
if 'cmpx' in mnemonic and mnemonic.endswith('_e64') and len(values) == 2:
values = [VGPR(126, 1)] + values
# Recalculate modifiers: parsed[0]=src0, parsed[1]=src1 (no vdst in user input)
neg_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[1])
abs_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[2])
opsel_bits = sum((1 << i) for i, p in enumerate(parsed[:2]) if p[3])
vop3sd_ops = {'v_div_scale_f32', 'v_div_scale_f64'}
if mnemonic in vop3sd_ops and len(parsed) >= 5:
neg_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[1])
abs_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[2])
if mnemonic in SOPK_UNSUPPORTED: raise ValueError(f"unsupported instruction: {mnemonic}")
elif mnemonic in SOP1_SRC_ONLY:
return getattr(autogen, mnemonic)(ssrc0=values[0])
elif mnemonic in SOP1_MSG_IMM:
return getattr(autogen, mnemonic)(sdst=values[0], ssrc0=RawImm(unwrap(values[1])))
elif mnemonic in SOPK_IMM_ONLY:
return getattr(autogen, mnemonic)(simm16=values[0])
elif mnemonic in SOPK_IMM_FIRST:
return getattr(autogen, mnemonic)(simm16=values[0], sdst=values[1])
elif mnemonic in SMEM_OPS and len(operands) >= 3 and re.match(r'^-?[0-9]|^-?0x', operands[2].strip().lower()):
return getattr(autogen, mnemonic)(sdata=values[0], sbase=values[1], offset=values[2], soffset=RawImm(124))
elif mnemonic.startswith('buffer_') and len(operands) >= 2 and operands[1].strip().lower() == 'off':
return getattr(autogen, mnemonic)(vdata=values[0], vaddr=0, srsrc=values[2], soffset=RawImm(unwrap(values[3])) if len(values) > 3 else RawImm(0))
elif (mnemonic.startswith('flat_load') or mnemonic.startswith('global_load') or mnemonic.startswith('scratch_load')) and len(values) >= 3:
offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
return getattr(autogen, mnemonic)(vdst=values[0], addr=values[1], saddr=values[2], offset=offset)
elif (mnemonic.startswith('flat_store') or mnemonic.startswith('global_store') or mnemonic.startswith('scratch_store')) and len(values) >= 3:
offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
return getattr(autogen, mnemonic)(addr=values[0], data=values[1], saddr=values[2], offset=offset)
for suffix in (['_e32', ''] if not (neg_bits or abs_bits or clamp) else ['', '_e32']):
if hasattr(autogen, name := mnemonic.replace('.', '_') + suffix):
use_opsel = 'opsel' in getattr(autogen, name).func._fields
vals = [type(v)(v.idx, v.count, False) if isinstance(v, Reg) and v.hi and use_opsel else v for v in values]
inst = getattr(autogen, name)(*vals, literal=lit, **modifiers)
if neg_bits and 'neg' in inst._fields: inst._values['neg'] = neg_bits
if opsel_bits and use_opsel: inst._values['opsel'] = opsel_bits
if abs_bits and 'abs' in inst._fields: inst._values['abs'] = abs_bits
if clamp and 'clmp' in inst._fields: inst._values['clmp'] = 1
return inst
raise ValueError(f"unknown instruction: {mnemonic}")
File diff suppressed because it is too large Load Diff
+505
View File
@@ -0,0 +1,505 @@
# RDNA3 emulator - pure Python implementation for testing
from __future__ import annotations
import ctypes, struct, math
from typing import Callable
from extra.assembly.rdna3.lib import Inst, Inst32, Inst64, RawImm
Program = dict[int, Inst] # pc (word offset) -> instruction
from extra.assembly.rdna3.autogen import (
SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
)
from extra.assembly.rdna3.alu import (
f32, i32, f16, i16, sext, vopc, SALU, VALU,
SOP1_BASE, SOP2_BASE, SOPC_BASE, SOPK_BASE, VOP1_BASE, VOP2_BASE
)
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
# Pre-computed inline constant table for src operands 128-254 (index = src - 128)
_INLINE_CONSTS = [0] * 127
for _i in range(65): _INLINE_CONSTS[_i] = _i # 128-192 -> 0-64
for _i in range(1, 17): _INLINE_CONSTS[64 + _i] = ((-_i) & 0xffffffff) # 193-208 -> -1 to -16
for _k, _v in {SrcEnum.POS_HALF: 0x3f000000, SrcEnum.NEG_HALF: 0xbf000000, SrcEnum.POS_ONE: 0x3f800000, SrcEnum.NEG_ONE: 0xbf800000,
SrcEnum.POS_TWO: 0x40000000, SrcEnum.NEG_TWO: 0xc0000000, SrcEnum.POS_FOUR: 0x40800000, SrcEnum.NEG_FOUR: 0xc0800000,
SrcEnum.INV_2PI: 0x3e22f983}.items(): _INLINE_CONSTS[_k - 128] = _v
_valid_mem_ranges: list[tuple[int, int]] = []
def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: global _valid_mem_ranges; _valid_mem_ranges = list(ranges)
def _mem_valid(addr: int, size: int) -> bool:
for s, z in _valid_mem_ranges:
if s <= addr and addr + size <= s + z: return True
return not _valid_mem_ranges
def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
def mem_write(addr: int, size: int, val: int) -> None:
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
# Memory op tables - (cnt, sz, sign) for loads, (cnt, sz) for stores
def _mem_ops(ops, suffix_map):
return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
FLAT_LOAD = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP)
FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
DS_LOAD: dict[int, tuple[int,int,int]] = {DSOp.DS_LOAD_B32: (1,4,0), DSOp.DS_LOAD_B64: (2,4,0), DSOp.DS_LOAD_B128: (4,4,0), DSOp.DS_LOAD_U8: (1,1,0), DSOp.DS_LOAD_I8: (1,1,1), DSOp.DS_LOAD_U16: (1,2,0), DSOp.DS_LOAD_I16: (1,2,1)}
DS_STORE: dict[int, tuple[int,int]] = {DSOp.DS_STORE_B32: (1,4), DSOp.DS_STORE_B64: (2,4), DSOp.DS_STORE_B128: (4,4), DSOp.DS_STORE_B8: (1,1), DSOp.DS_STORE_B16: (1,2)}
FLAT_D16_LO = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_U8', (1, 0)), ('LOAD_D16_I8', (1, 1)), ('LOAD_D16_B16', (2, 0))]}
FLAT_D16_HI = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_HI_U8', (1, 0)), ('LOAD_D16_HI_I8', (1, 1)), ('LOAD_D16_HI_B16', (2, 0))]}
FLAT_D16_STORE = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('STORE_D16_HI_B8', 1), ('STORE_D16_HI_B16', 2)]}
SMEM_LOAD: dict[int, int] = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
SOPK_WAIT = {SOPKOp.S_WAITCNT_VSCNT, SOPKOp.S_WAITCNT_VMCNT, SOPKOp.S_WAITCNT_EXPCNT, SOPKOp.S_WAITCNT_LGKMCNT}
class WaveState:
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
def __init__(self):
self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
self.sgpr[EXEC_LO] = 0xffffffff # wave32: all lanes active
self.scc = self.pc = self.literal = 0
self._pend_sgpr = {}
@property
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
@vcc.setter
def vcc(self, v: int) -> None: self.sgpr[VCC_LO] = v & 0xffffffff; self.sgpr[VCC_HI] = (v >> 32) & 0xffffffff
@property
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
@exec_mask.setter
def exec_mask(self, v: int) -> None: self.sgpr[EXEC_LO] = v & 0xffffffff; self.sgpr[EXEC_HI] = (v >> 32) & 0xffffffff
def rsgpr(self, i: int) -> int:
if i == NULL: return 0
if i == SCC: return self.scc
return self.sgpr[i] if i < SGPR_COUNT else 0
def wsgpr(self, i: int, v: int) -> None:
if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & 0xffffffff
def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
def wsgpr64(self, i: int, v: int) -> None: self.wsgpr(i, v & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
def rsrc(self, v: int, lane: int) -> int:
if v < SGPR_COUNT: return self.sgpr[v]
if v == SCC: return self.scc
if v < 255: return _INLINE_CONSTS[v - 128]
if v == 255: return self.literal
return self.vgpr[lane][v - 256] if v <= 511 else 0
def rsrc64(self, v: int, lane: int) -> int:
return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
def pend_sgpr_lane(self, reg: int, lane: int, val: int) -> None:
if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
if val: self._pend_sgpr[reg] |= (1 << lane)
def commit_pends(self) -> None:
for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
self._pend_sgpr.clear()
def decode_format(word: int) -> tuple[type[Inst] | None, bool]:
hi2 = (word >> 30) & 0x3
if hi2 == 0b11:
enc = (word >> 26) & 0xf
if enc == 0b1101: return SMEM, True
if enc == 0b0101:
op = (word >> 16) & 0x3ff
return (VOP3SD, True) if op in (288, 289, 290, 764, 765, 766, 767, 768, 769, 770) else (VOP3, True)
return {0b0011: (VOP3P, True), 0b0110: (DS, True), 0b0111: (FLAT, True), 0b0010: (VOPD, True)}.get(enc, (None, True))
if hi2 == 0b10:
enc = (word >> 23) & 0x7f
return {0b1111101: (SOP1, False), 0b1111110: (SOPC, False), 0b1111111: (SOPP, False)}.get(enc, (SOPK, False) if ((word >> 28) & 0xf) == 0b1011 else (SOP2, False))
enc = (word >> 25) & 0x7f
return (VOPC, False) if enc == 0b0111110 else (VOP1, False) if enc == 0b0111111 else (VOP2, False)
def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
def decode_program(data: bytes) -> Program:
result: Program = {}
i = 0
while i < len(data):
word = int.from_bytes(data[i:i+4], 'little')
inst_class, is_64 = decode_format(word)
if inst_class is None: i += 4; continue
base_size = 8 if is_64 else 4
inst = inst_class.from_bytes(data[i:i+base_size])
for name, val in inst._values.items(): setattr(inst, name, _unwrap(val))
has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
if has_literal: inst._literal = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
inst._words = inst.size() // 4 # cache size for step_wave
result[i // 4] = inst
i += inst._words * 4
return result
# ═══════════════════════════════════════════════════════════════════════════════
# SCALAR EXECUTION
# ═══════════════════════════════════════════════════════════════════════════════
def exec_sop1(st: WaveState, inst: SOP1) -> int:
s0, op = st.rsrc(inst.ssrc0, 0), inst.op
# 64-bit and special ops handled inline
if op == SOP1Op.S_MOV_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0)); return 0
if op == SOP1Op.S_NOT_B64: r = (~st.rsrc64(inst.ssrc0, 0)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP1Op.S_BITSET0_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) & ~(1 << (s0 & 0x1f))); return 0
if op == SOP1Op.S_BITSET1_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) | (1 << (s0 & 0x1f))); return 0
if op == SOP1Op.S_AND_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
if op == SOP1Op.S_OR_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 | old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
if op == SOP1Op.S_AND_NOT1_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & (~old & 0xffffffff); st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
if op == SOP1Op.S_GETPC_B64: return -3
if op == SOP1Op.S_SETPC_B64: return -4
if op == SOP1Op.S_SWAPPC_B64: return -5
if (fn := SALU.get(SOP1_BASE + op)) is None: raise NotImplementedError(f"SOP1 op {op}")
r, scc = fn(s0, 0, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
_SOP2_64: dict[int, Callable[[int, int], int]] = {SOP2Op.S_AND_B64: lambda a, b: a & b, SOP2Op.S_OR_B64: lambda a, b: a | b, SOP2Op.S_XOR_B64: lambda a, b: a ^ b}
def exec_sop2(st: WaveState, inst: SOP2) -> int:
s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
# 64-bit ops handled inline
if op == SOP2Op.S_LSHL_B64: r = (st.rsrc64(inst.ssrc0, 0) << (s1 & 0x3f)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP2Op.S_LSHR_B64: r = st.rsrc64(inst.ssrc0, 0) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP2Op.S_ASHR_I64: r = sext(st.rsrc64(inst.ssrc0, 0), 64) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r & 0xffffffffffffffff); st.scc = int(r != 0); return 0
if (fn := _SOP2_64.get(op)): r = fn(st.rsrc64(inst.ssrc0, 0), st.rsrc64(inst.ssrc1, 0)); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
if op == SOP2Op.S_CSELECT_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0) if st.scc else st.rsrc64(inst.ssrc1, 0)); return 0
if op == SOP2Op.S_FMAC_F32: st.wsgpr(inst.sdst, i32(f32(st.rsgpr(inst.sdst)) + f32(s0) * f32(s1))); return 0
if op == SOP2Op.S_FMAAK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(s1) + f32(inst._literal or 0))); return 0
if op == SOP2Op.S_FMAMK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(inst._literal or 0) + f32(s1))); return 0
if (fn := SALU.get(SOP2_BASE + op)) is None: raise NotImplementedError(f"SOP2 op {op}")
r, scc = fn(s0, s1, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
def exec_sopc(st: WaveState, inst: SOPC) -> int:
s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
if op == SOPCOp.S_CMP_EQ_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) == st.rsrc64(inst.ssrc1, 0)); return 0
if op == SOPCOp.S_CMP_LG_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) != st.rsrc64(inst.ssrc1, 0)); return 0
if (fn := SALU.get(SOPC_BASE + op)) is None: raise NotImplementedError(f"SOPC op {op}")
st.scc = fn(s0, s1, st.scc)[1]; return 0
_SOPK_CMP = frozenset((SOPKOp.S_CMPK_EQ_I32, SOPKOp.S_CMPK_LG_I32, SOPKOp.S_CMPK_GT_I32, SOPKOp.S_CMPK_GE_I32,
SOPKOp.S_CMPK_LT_I32, SOPKOp.S_CMPK_LE_I32, SOPKOp.S_CMPK_EQ_U32, SOPKOp.S_CMPK_LG_U32,
SOPKOp.S_CMPK_GT_U32, SOPKOp.S_CMPK_GE_U32, SOPKOp.S_CMPK_LT_U32, SOPKOp.S_CMPK_LE_U32))
def exec_sopk(st: WaveState, inst: SOPK) -> int:
simm, s0, op = inst.simm16, st.rsgpr(inst.sdst), inst.op
if op in SOPK_WAIT: return 0
if (fn := SALU.get(SOPK_BASE + op)) is None: raise NotImplementedError(f"SOPK op {op}")
r, scc = fn(s0, simm, st.scc)
if op not in _SOPK_CMP: st.wsgpr(inst.sdst, r)
st.scc = scc; return 0
def exec_sopp(st: WaveState, inst: SOPP) -> int:
if inst.op == SOPPOp.S_ENDPGM: return -1
if inst.op == SOPPOp.S_BARRIER: return -2
if inst.op == SOPPOp.S_BRANCH: return sext(inst.simm16, 16)
if inst.op == SOPPOp.S_CBRANCH_SCC0: return sext(inst.simm16, 16) if st.scc == 0 else 0
if inst.op == SOPPOp.S_CBRANCH_SCC1: return sext(inst.simm16, 16) if st.scc == 1 else 0
# In wave32 mode, only VCC_LO is used for lane masks; VCC_HI is a free SGPR
if inst.op == SOPPOp.S_CBRANCH_VCCZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) == 0 else 0
if inst.op == SOPPOp.S_CBRANCH_VCCNZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) != 0 else 0
if inst.op == SOPPOp.S_CBRANCH_EXECZ: return sext(inst.simm16, 16) if st.exec_mask == 0 else 0
if inst.op == SOPPOp.S_CBRANCH_EXECNZ: return sext(inst.simm16, 16) if st.exec_mask != 0 else 0
# Scheduling hints and wait instructions are no-ops in emulation
if inst.op <= 31: return 0 # S_NOP, S_CLAUSE, S_DELAY_ALU, S_WAITCNT, etc.
# S_WAKEUP(52), S_SETPRIO(53), S_SENDMSG(54), S_SENDMSGHALT(55), perf counters, S_ICACHE_INV(60) are no-ops
if inst.op in (52, 53, 54, 55, 56, 57, 60): return 0
raise NotImplementedError(f"SOPP op {inst.op}")
def exec_smem(st: WaveState, inst: SMEM) -> int:
addr = st.rsgpr64(inst.sbase * 2) + sext(inst.offset, 21)
if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & 0xffffffffffffffff, 4))
return 0
# ═══════════════════════════════════════════════════════════════════════════════
# VECTOR EXECUTION
# ═══════════════════════════════════════════════════════════════════════════════
def f64(hi: int, lo: int) -> float: return struct.unpack('<d', struct.pack('<Q', (hi << 32) | lo))[0]
def i64_parts(f: float) -> tuple[int, int]:
if math.isnan(f): val = 0x7ff8000000000000
elif math.isinf(f): val = 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
else: val = struct.unpack('<Q', struct.pack('<d', f))[0]
return val & 0xffffffff, (val >> 32) & 0xffffffff
def exec_vop1(st: WaveState, inst: VOP1, lane: int) -> None:
if inst.op == VOP1Op.V_NOP: return
V, s0 = st.vgpr[lane], st.rsrc(inst.src0, lane)
if inst.op == VOP1Op.V_READFIRSTLANE_B32:
first = (st.exec_mask & -st.exec_mask).bit_length() - 1 if st.exec_mask else 0
st.wsgpr(inst.vdst, st.rsrc(inst.src0, first) if inst.src0 >= 256 else s0); return
# F64 ops handled inline
if inst.op == VOP1Op.V_CVT_F64_F32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(f32(s0))); return
if inst.op == VOP1Op.V_CVT_F64_I32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(sext(s0, 32))); return
if inst.op == VOP1Op.V_CVT_F64_U32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(s0)); return
if inst.op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64):
src = inst.src0 - 256 if inst.src0 >= 256 else inst.src0
lo, hi = (V[src], V[src+1]) if inst.src0 >= 256 else (st.sgpr[src], st.sgpr[src+1])
v = f64(hi, lo)
if inst.op == VOP1Op.V_CVT_F32_F64: V[inst.vdst] = i32(v)
elif inst.op == VOP1Op.V_CVT_I32_F64: V[inst.vdst] = (max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff) if math.isfinite(v) else 0
else: V[inst.vdst] = max(0, min(0xffffffff, int(v))) if math.isfinite(v) and v == v else 0
return
if (fn := VALU.get(VOP1_BASE + inst.op)): V[inst.vdst] = fn(s0, 0, 0); return
raise NotImplementedError(f"VOP1 op {inst.op}")
def exec_vop2(st: WaveState, inst: VOP2, lane: int) -> None:
V, s0, s1, op = st.vgpr[lane], st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1], inst.op
if op == VOP2Op.V_CNDMASK_B32: V[inst.vdst] = s1 if (st.vcc >> lane) & 1 else s0; return
if op == VOP2Op.V_FMAC_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(V[inst.vdst])); return
if op == VOP2Op.V_FMAMK_F32: V[inst.vdst] = i32(f32(s0)*f32(st.literal)+f32(s1)); return
if op == VOP2Op.V_FMAAK_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(st.literal)); return
if op == VOP2Op.V_FMAC_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(V[inst.vdst])); return
if op == VOP2Op.V_FMAMK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(st.literal)+f16(s1)); return
if op == VOP2Op.V_FMAAK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(st.literal)); return
if op == VOP2Op.V_PK_FMAC_F16:
lo = i16(f16(s0 & 0xffff) * f16(s1 & 0xffff) + f16(V[inst.vdst] & 0xffff))
hi = i16(f16((s0 >> 16) & 0xffff) * f16((s1 >> 16) & 0xffff) + f16((V[inst.vdst] >> 16) & 0xffff))
V[inst.vdst] = lo | (hi << 16); return
if op == VOP2Op.V_ADD_CO_CI_U32: r = s0+s1+((st.vcc>>lane)&1); st.pend_sgpr_lane(VCC_LO, lane, r >= 0x100000000); V[inst.vdst] = r & 0xffffffff; return
if op == VOP2Op.V_SUB_CO_CI_U32: b = (st.vcc>>lane)&1; st.pend_sgpr_lane(VCC_LO, lane, s1+b > s0); V[inst.vdst] = (s0-s1-b) & 0xffffffff; return
if (fn := VALU.get(VOP2_BASE + op)): V[inst.vdst] = fn(s0, s1, 0); return
raise NotImplementedError(f"VOP2 op {op}")
def vop3_mod(val: int, neg: int, abs_: int, idx: int) -> int:
if (abs_ >> idx) & 1: val = i32(abs(f32(val)))
if (neg >> idx) & 1: val = i32(-f32(val))
return val
def exec_vop3(st: WaveState, inst: VOP3, lane: int) -> None:
op, src0, src1, src2, vdst, neg, abs_ = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.neg, getattr(inst, 'abs', 0)
V = st.vgpr[lane]
# VOPC encoded in VOP3 (0-255)
if 0 <= op <= 255:
base = op & 0x7f
# For 64-bit comparisons (I64: 80-87, U64: 88-95), read raw 64-bit values (no float modifiers)
if 80 <= base <= 95:
s0_64, s1_64 = st.rsrc64(src0, lane), st.rsrc64(src1, lane)
result = vopc(op, s0_64 & 0xffffffff, s1_64 & 0xffffffff, (s0_64 >> 32) & 0xffffffff, (s1_64 >> 32) & 0xffffffff)
else:
s0, s1 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1)
result = vopc(op, s0, s1)
is_cmpx = op >= 128
st.pend_sgpr_lane(vdst, lane, result)
if is_cmpx: st.pend_sgpr_lane(EXEC_LO, lane, result)
return
s0, s1, s2 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1), vop3_mod(st.rsrc(src2, lane), neg, abs_, 2)
# Special ops
if op == VOP3Op.V_FMAC_F32: V[vdst] = i32(f32(s0)*f32(s1)+f32(V[vdst])); return
if op == VOP3Op.V_READLANE_B32: st.wsgpr(vdst, st.vgpr[s1 & 0x1f][src0 - 256] if src0 >= 256 else s0); return
if op == VOP3Op.V_WRITELANE_B32: st.vgpr[s1 & 0x1f][vdst] = s0; return
if op == VOP3Op.V_CNDMASK_B32:
mask = st.rsgpr(src2) if src2 < 256 else st.vcc
V[vdst] = s1 if (mask >> lane) & 1 else s0; return
if op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64):
v64 = st.rsrc64(src1, lane)
r = ((v64 << (s0 & 0x3f)) & 0xffffffffffffffff if op == VOP3Op.V_LSHLREV_B64 else
v64 >> (s0 & 0x3f) if op == VOP3Op.V_LSHRREV_B64 else sext(v64, 64) >> (s0 & 0x3f))
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff; return
if op in (VOP3Op.V_ADD_F64, VOP3Op.V_MUL_F64, VOP3Op.V_FMA_F64, VOP3Op.V_MAX_F64, VOP3Op.V_MIN_F64):
a, b = f64(st.rsrc(src0+1, lane), s0), f64(st.rsrc(src1+1, lane), s1)
c = f64(st.rsrc(src2+1, lane), s2) if op == VOP3Op.V_FMA_F64 else 0.0
rf = a + b if op == VOP3Op.V_ADD_F64 else a * b if op == VOP3Op.V_MUL_F64 else a * b + c if op == VOP3Op.V_FMA_F64 else max(a, b) if op == VOP3Op.V_MAX_F64 else min(a, b)
V[vdst], V[vdst+1] = i64_parts(rf); return
if (fn := VALU.get(op)): V[vdst] = fn(s0, s1, s2); return
raise NotImplementedError(f"VOP3 op {op}")
def exec_vopc(st: WaveState, inst: VOPC, lane: int) -> None:
result, is_cmpx = vopc(inst.op, st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1]), inst.op >= 128
st.pend_sgpr_lane(EXEC_LO if is_cmpx else VCC_LO, lane, result)
def exec_vop3sd(st: WaveState, inst: VOP3SD, lane: int) -> None:
op, src0, src1, src2, vdst, sdst, neg = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.sdst, inst.neg
s0, s1, s2 = st.rsrc(src0, lane), st.rsrc(src1, lane), st.rsrc(src2, lane)
if (neg >> 0) & 1: s0 = i32(-f32(s0))
if (neg >> 1) & 1: s1 = i32(-f32(s1))
if (neg >> 2) & 1: s2 = i32(-f32(s2))
V = st.vgpr[lane]
if op == VOP3SDOp.V_ADD_CO_U32: r = s0 + s1; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
elif op == VOP3SDOp.V_SUB_CO_U32: V[vdst] = (s0 - s1) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 > s0)
elif op == VOP3SDOp.V_SUBREV_CO_U32: V[vdst] = (s1 - s0) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s0 > s1)
elif op == VOP3SDOp.V_ADD_CO_CI_U32:
cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
r = s0 + s1 + cin; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
elif op == VOP3SDOp.V_SUB_CO_CI_U32:
cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
V[vdst] = (s0 - s1 - cin) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 + cin > s0)
elif op == VOP3SDOp.V_MAD_U64_U32:
s2_64 = s2 | (st.rsrc(src2+1, lane) << 32); r = s0 * s1 + s2_64
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
elif op == VOP3SDOp.V_MAD_I64_I32:
s2_64 = sext(s2 | (st.rsrc(src2+1, lane) << 32), 64)
r = (sext(s0, 32) * sext(s1, 32) + s2_64) & 0xffffffffffffffff
V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
elif op == VOP3SDOp.V_DIV_SCALE_F32: V[vdst] = 0; st.pend_sgpr_lane(sdst, lane, False)
elif op == VOP3SDOp.V_DIV_SCALE_F64: V[vdst], V[vdst+1] = s0, st.rsrc(src0+1, lane); st.pend_sgpr_lane(VCC_LO, lane, s0 == s2)
else: raise NotImplementedError(f"VOP3SD op {op}")
def exec_flat(st: WaveState, inst: FLAT, lane: int) -> None:
op, addr_reg, data_reg, vdst, offset, saddr, V = inst.op, inst.addr, inst.data, inst.vdst, sext(inst.offset, 13), inst.saddr, st.vgpr[lane]
addr = V[addr_reg] | (V[addr_reg+1] << 32)
addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
if op in FLAT_LOAD:
cnt, sz, sign = FLAT_LOAD[op]
for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
elif op in FLAT_STORE:
cnt, sz = FLAT_STORE[op]
for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
elif op in FLAT_D16_LO: sz, sign = FLAT_D16_LO[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0xffff0000) | ((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff))
elif op in FLAT_D16_HI: sz, sign = FLAT_D16_HI[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0x0000ffff) | (((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff)) << 16)
elif op in FLAT_D16_STORE: mem_write(addr, FLAT_D16_STORE[op], (V[data_reg] >> 16) & ((1 << (FLAT_D16_STORE[op] * 8)) - 1))
else: raise NotImplementedError(f"FLAT op {op}")
def exec_ds(st: WaveState, inst: DS, lane: int, lds: bytearray) -> None:
op, addr, vdst, V = inst.op, (st.vgpr[lane][inst.addr] + inst.offset0) & 0xffff, inst.vdst, st.vgpr[lane]
if op in DS_LOAD:
cnt, sz, sign = DS_LOAD[op]
for i in range(cnt): val = int.from_bytes(lds[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
elif op in DS_STORE:
cnt, sz = DS_STORE[op]
for i in range(cnt): lds[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
else: raise NotImplementedError(f"DS op {op}")
VOPD_OPS: dict[int, Callable[[int, int, int, int, int], int]] = {
VOPDOp.V_DUAL_MUL_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)), VOPDOp.V_DUAL_ADD_F32: lambda a, b, d, l, lit: i32(f32(a)+f32(b)),
VOPDOp.V_DUAL_SUB_F32: lambda a, b, d, l, lit: i32(f32(a)-f32(b)), VOPDOp.V_DUAL_SUBREV_F32: lambda a, b, d, l, lit: i32(f32(b)-f32(a)),
VOPDOp.V_DUAL_MAX_F32: lambda a, b, d, l, lit: i32(max(f32(a), f32(b))), VOPDOp.V_DUAL_MIN_F32: lambda a, b, d, l, lit: i32(min(f32(a), f32(b))),
VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: lambda a, b, d, l, lit: i32(0.0 if f32(a) == 0.0 or f32(b) == 0.0 else f32(a)*f32(b)),
VOPDOp.V_DUAL_MOV_B32: lambda a, b, d, l, lit: a, VOPDOp.V_DUAL_ADD_NC_U32: lambda a, b, d, l, lit: (a + b) & 0xffffffff,
VOPDOp.V_DUAL_LSHLREV_B32: lambda a, b, d, l, lit: (b << (a & 0x1f)) & 0xffffffff, VOPDOp.V_DUAL_AND_B32: lambda a, b, d, l, lit: a & b,
VOPDOp.V_DUAL_FMAC_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(d)), VOPDOp.V_DUAL_FMAAK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(lit)),
VOPDOp.V_DUAL_FMAMK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(lit)+f32(b)), VOPDOp.V_DUAL_CNDMASK_B32: lambda a, b, d, l, lit: b if l else a,
}
def exec_vopd(st: WaveState, inst: VOPD, lane: int) -> None:
V, vdsty, vcc_lane = st.vgpr[lane], (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), (st.vcc >> lane) & 1
sx0, sx1, sy0, sy1, dstx = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], inst.vdstx
if (fn := VOPD_OPS.get(inst.opx)): V[dstx] = fn(sx0, sx1, V[dstx], vcc_lane, st.literal)
else: raise NotImplementedError(f"VOPD opx {inst.opx}")
if (fn := VOPD_OPS.get(inst.opy)): V[vdsty] = fn(sy0, sy1, V[vdsty], vcc_lane, st.literal)
else: raise NotImplementedError(f"VOPD opy {inst.opy}")
def exec_vop3p(st: WaveState, inst: VOP3P, lane: int) -> None:
op, vdst, V = inst.op, inst.vdst, st.vgpr[lane]
s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
opsel, opsel_hi = [(inst.opsel >> i) & 1 for i in range(3)], [(inst.opsel_hi >> i) & 1 for i in range(2)] + [inst.opsel_hi2]
neg, neg_hi = inst.neg, inst.neg_hi
def get_src(src: int, idx: int, for_mix: bool = False) -> float:
if for_mix:
if not opsel_hi[idx]: return abs(f32(src)) if (neg_hi >> idx) & 1 else f32(src)
return float(f16((src >> 16) & 0xffff) if opsel[idx] else f16(src & 0xffff))
use_hi = opsel[idx]
val = ((src >> 16) & 0xffff) if use_hi else (src & 0xffff)
f = f16(val)
if use_hi and (neg >> idx) & 1: f = -f
elif not use_hi and (neg_hi >> idx) & 1: f = -f
return f
if op == VOP3POp.V_FMA_MIX_F32: V[vdst] = i32(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
elif op == VOP3POp.V_FMA_MIXLO_F16: V[vdst] = (V[vdst] & 0xffff0000) | i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
elif op == VOP3POp.V_FMA_MIXHI_F16: V[vdst] = (V[vdst] & 0x0000ffff) | (i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True)) << 16)
else: raise NotImplementedError(f"VOP3P op {op}")
def exec_wmma_f32_16x16x16_f16(st: WaveState, inst: VOP3P, n_lanes: int) -> None:
src0_base, src1_base, src2_base = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0, (inst.src1 - 256) if inst.src1 >= 256 else inst.src1, (inst.src2 - 256) if inst.src2 >= 256 else inst.src2
src0_is_vgpr, src1_is_vgpr, src2_is_vgpr, vdst = inst.src0 >= 256, inst.src1 >= 256, inst.src2 >= 256, inst.vdst
A, B, C = [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)]
for lane in range(min(n_lanes, 16)):
V = st.vgpr[lane]
for reg in range(8):
val = V[src0_base + reg] if src0_is_vgpr else st.sgpr[src0_base + reg]
A[lane][reg * 2], A[lane][reg * 2 + 1] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
val = V[src1_base + reg] if src1_is_vgpr else st.sgpr[src1_base + reg]
B[reg * 2][lane], B[reg * 2 + 1][lane] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
for row in range(16):
for col in range(16):
idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
if lane_idx < n_lanes:
val = st.vgpr[lane_idx][src2_base + reg] if src2_is_vgpr else st.sgpr[src2_base + reg]
C[row][col] = f32(val)
for row in range(16):
for col in range(16):
for k in range(16): C[row][col] += A[row][k] * B[k][col]
for row in range(16):
for col in range(16):
idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
if lane_idx < n_lanes and (st.exec_mask & (1 << lane_idx)): st.vgpr[lane_idx][vdst + reg] = i32(C[row][col])
# ═══════════════════════════════════════════════════════════════════════════════
# MAIN EXECUTION LOOP
# ═══════════════════════════════════════════════════════════════════════════════
SCALAR: dict[type, Callable[..., int]] = {SOP1: exec_sop1, SOP2: exec_sop2, SOPC: exec_sopc, SOPK: exec_sopk, SOPP: exec_sopp, SMEM: exec_smem}
VECTOR: dict[type, Callable[..., None]] = {VOP1: exec_vop1, VOP2: exec_vop2, VOP3: exec_vop3, VOP3SD: exec_vop3sd, VOPC: exec_vopc, FLAT: exec_flat, DS: exec_ds, VOPD: exec_vopd, VOP3P: exec_vop3p}
_WMMA_OPS = frozenset((VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16,
VOP3POp.V_WMMA_BF16_16X16X16_BF16, VOP3POp.V_WMMA_I32_16X16X16_IU8, VOP3POp.V_WMMA_I32_16X16X16_IU4))
def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
inst = program.get(st.pc)
if inst is None: return 1
inst_words, st.literal, inst_type = inst._words, inst._literal or 0, type(inst)
if (handler := SCALAR.get(inst_type)) is not None:
delta = handler(st, inst)
if delta == -1: return -1
if delta == -2: st.pc += inst_words; return -2
if delta == -3: # S_GETPC_B64
sop1 = inst if isinstance(inst, SOP1) else None
assert sop1 is not None
next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc += inst_words; return 0
if delta == -4: # S_SETPC_B64
sop1 = inst if isinstance(inst, SOP1) else None
assert sop1 is not None
st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
if delta == -5: # S_SWAPPC_B64
sop1 = inst if isinstance(inst, SOP1) else None
assert sop1 is not None
next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
st.pc += inst_words + delta
else:
vec_handler, exec_mask = VECTOR[inst_type], st.exec_mask
if inst_type is DS:
for lane in range(n_lanes):
if exec_mask & (1 << lane): vec_handler(st, inst, lane, lds)
elif inst_type is VOP3P:
vop3p = inst if isinstance(inst, VOP3P) else None
assert vop3p is not None
if vop3p.op in _WMMA_OPS:
exec_wmma_f32_16x16x16_f16(st, vop3p, n_lanes)
else:
for lane in range(n_lanes):
if exec_mask & (1 << lane): vec_handler(st, vop3p, lane)
else:
for lane in range(n_lanes):
if exec_mask & (1 << lane): vec_handler(st, inst, lane)
st.commit_pends(); st.pc += inst_words
return 0
def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int, wg_id: tuple[int,int,int]=(0,0,0), local_size: tuple[int,int,int]=(1,1,1), wave_start: int=0) -> int:
while st.pc in program:
result = step_wave(program, st, lds, n_lanes)
if result == -1: return 0
if result == -2: return -2
return 0
def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, dispatch_dim: int) -> None:
lx, ly, lz = local_size
total_threads, lds = lx * ly * lz, bytearray(65536)
waves: list[tuple[WaveState, int, int]] = []
for wave_start in range(0, total_threads, WAVE_SIZE):
n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
st.exec_mask = (1 << n_lanes) - 1
st.wsgpr64(0, args_ptr)
gx, gy, gz = workgroup_id
if dispatch_dim >= 3: st.sgpr[13], st.sgpr[14], st.sgpr[15] = gx, gy, gz
elif dispatch_dim == 2: st.sgpr[14], st.sgpr[15] = gx, gy
else: st.sgpr[15] = gx
for i in range(n_lanes):
tid = wave_start + i
st.vgpr[i][0] = tid if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
waves.append((st, n_lanes, wave_start))
has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
for _ in range(2 if has_barrier else 1):
for st, n_lanes, wave_start in waves: exec_wave(program, st, lds, n_lanes, workgroup_id, local_size, wave_start)
def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
data = (ctypes.c_char * lib_sz).from_address(lib).raw
program = decode_program(data)
if not program: return -1
dispatch_dim = 3 if gz > 1 else (2 if gy > 1 else 1)
for gidz in range(gz):
for gidy in range(gy):
for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, dispatch_dim)
return 0
+191
View File
@@ -0,0 +1,191 @@
#!/usr/bin/env python3
# generates autogen/__init__.py by parsing the AMD RDNA3.5 ISA PDF
import re, pdfplumber, pathlib
from tinygrad.helpers import fetch
PDF_URL = "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"
FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
FIELD_ORDER = {
'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
'4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
def parse_bits(s: str) -> tuple[int, int] | None:
return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
def parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
fields = []
for row in table[1:]:
if not row or not row[0]: continue
name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
if not (bits := parse_bits(bits_str)): continue
enc_val, hi, lo = None, bits[0], bits[1]
if name == 'ENCODING' and row[2] and (m := re.search(r"'b([01_]+)", row[2])):
enc_bits = m.group(1).replace('_', '')
enc_val = int(enc_bits, 2)
declared_width, actual_width = hi - lo + 1, len(enc_bits)
if actual_width > declared_width: lo = hi - actual_width + 1
ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
fields.append((name, hi, lo, enc_val, ftype))
return fields
def generate(output_path: pathlib.Path|str|None = None) -> dict:
"""Generate RDNA3.5 instruction definitions from the AMD ISA PDF. Returns dict with formats for testing."""
pdf = pdfplumber.open(fetch(PDF_URL))
pages = pdf.pages[150:200]
page_texts = [p.extract_text() or '' for p in pages]
page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
full_text = '\n'.join(page_texts)
# parse SSRC encoding from first page with VCC_LO
src_enum = dict(SRC_EXTRAS)
for text in page_texts[:10]:
if 'SSRC0' in text and 'VCC_LO' in text:
for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
val, name = int(m.group(1)), m.group(2).rstrip('.:')
if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
break
# parse opcode tables
enums: dict[str, dict[int, str]] = {}
for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
enums[m.group(1) + "Op"] = ops
if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
enums["VOPDOp"] = ops
enum_names = set(enums.keys())
def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
def has_header_before_fields(text) -> bool:
return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
# find format headers with their page indices
format_headers = [] # (fmt_name, page_idx)
for i, text in enumerate(page_texts):
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
if m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < len(page_texts):
next_text = page_texts[i + 1].lstrip()
if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
format_headers.append((m.group(1), i, m.start()))
# parse instruction formats
formats: dict[str, list] = {}
for fmt_name, page_idx, header_pos in format_headers:
if fmt_name in formats: continue
text, tables = page_texts[page_idx], page_tables[page_idx]
field_pos = text.find('Field Name', header_pos)
# find fields table with ENCODING (same page or up to 2 pages ahead)
fields = None
for offset in range(3):
if page_idx + offset >= len(pages): break
if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
fields = f
break
if fields: break
# for modifier formats (no ENCODING), accept first fields table on same page
if not fields and field_pos > header_pos:
for t in tables:
if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)):
fields = f
break
if not fields: continue
field_names = {f[0] for f in fields}
# check next pages for continuation fields (tables without ENCODING)
for pg_offset in range(1, 3):
if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
for t in page_tables[page_idx + pg_offset]:
if is_fields_table(t) and (extra := parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
for ef in extra:
if ef[0] not in field_names:
fields.append(ef)
field_names.add(ef[0])
break
formats[fmt_name] = fields
# fix known PDF errors (verified against LLVM test vectors)
# SMEM: PDF says DLC=bit14, GLC=bit16 but actual encoding is DLC=bit13, GLC=bit14
if 'SMEM' in formats:
formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
for n, h, l, e, t in formats['SMEM']]
# generate output
def enum_lines(name, items):
return [f"class {name}(IntEnum):"] + [f" {n} = {v}" for v, n in sorted(items.items())] + [""]
def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
lines = ["# autogenerated from AMD RDNA3.5 ISA PDF by gen.py - do not edit", "from enum import IntEnum",
"from typing import Annotated",
"from extra.assembly.rdna3.lib import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
"import functools", ""]
lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
# Format-specific field defaults (verified against LLVM test vectors)
format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
lines.append("# instruction formats")
for fmt_name, fields in sorted(formats.items()):
base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
order = FIELD_ORDER.get(fmt_name, [])
lines.append(f"class {fmt_name}({base}):")
if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
lines.append(f" encoding = {enc_str}")
if defaults := format_defaults.get(fmt_name):
lines.append(f" _defaults = {defaults}")
for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
# Wrap IntEnum types (ending in Op) with Annotated[BitField, ...] for correct typing
if ftype and ftype.endswith('Op'):
ann = f":Annotated[BitField, {ftype}]"
else:
ann = f":{ftype}" if ftype else ""
lines.append(f" {name.lower()}{ann} = bits[{hi}]" if hi == lo else f" {name.lower()}{ann} = bits[{hi}:{lo}]")
lines.append("")
lines.append("# instruction helpers")
for cls_name, ops in sorted(enums.items()):
fmt = cls_name[:-2]
for op_val, name in sorted(ops.items()):
seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=2"}.get(fmt, "")
tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
# VOP1/VOP2/VOPC get _e32 suffix, VOP3 promoted ops (< 512) get _e64 suffix
if fmt in ("VOP1", "VOP2", "VOPC"):
suffix = "_e32"
elif fmt == "VOP3" and op_val < 512:
suffix = "_e64"
else:
suffix = ""
lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
# export SrcEnum values, but skip DPP8/DPP16 which conflict with class names
skip_exports = {'DPP8', 'DPP16'}
lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in skip_exports] + ["OFF = NULL\n"]
if output_path is not None: pathlib.Path(output_path).write_text('\n'.join(lines))
return {"formats": formats, "enums": enums, "src_enum": src_enum}
if __name__ == "__main__":
result = generate("extra/assembly/rdna3/autogen/__init__.py")
print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")
-136
View File
@@ -1,136 +0,0 @@
import os, sys, struct
sys.path.append(os.getcwd())
# PROFILE=1 to use
#os.environ["PROFILE"] = "1"
os.environ["SQTT"] = "1"
os.environ["SQTT_ITRACE_SE_MASK"] = "1"
os.environ["SQTT_LIMIT_SE"] = "1"
import xml.etree.ElementTree as ET
from tinygrad import nn, Tensor, Device
from tinygrad.helpers import get_single_element
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.ops_amd import ProfileSQTTEvent
from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
def disassemble(text, root:ET.Element):
i = 0
while i < len(text):
ins = struct.unpack("I", text[i:i+4])[0]
# 1. Get the encoding
did_match = False
for enc_el in root.findall("./ISA/Encodings/Encoding"):
mask = enc_el.findtext("EncodingIdentifierMask")
assert len(mask)%32 == 0
bit_mask = int(mask, 2)
iden = [int(x.text, 2) for x in enc_el.find("EncodingIdentifiers").findall("EncodingIdentifier")]
for ide in iden:
if ins&bit_mask == ide:
did_match = True
break
if did_match: break
if not did_match: raise RuntimeError(f"unknown instruction {ins:08X}")
if len(mask) >= 64: ins = (struct.unpack("I", text[i+4:i+8])[0]<<32) | ins
if len(mask) >= 96: ins = (struct.unpack("I", text[i+8:i+12])[0]<<64) | ins
encoding_name = enc_el.findtext("EncodingName")
#print(ET.tostring(enc_el).decode())
# 2. Parse the Fields for this Encoding
field_data = {}
for field in enc_el.findall("MicrocodeFormat/BitMap/Field"):
# Fields can be split into multiple ranges (RangeCount > 1)
ranges = sorted(field.findall("BitLayout/Range"), key=lambda x: int(x.attrib.get('Order')))
val = 0
current_shift = 0
for rng in ranges:
width = int(rng.find("BitCount").text)
chunk = (ins >> int(rng.find("BitOffset").text)) & ((1 << width) - 1)
val |= (chunk << current_shift)
current_shift += width
field_data[field.find("FieldName").text] = val
# this is already used
del field_data["ENCODING"]
# 3. Extract the instruction
did_match = False
for ins_el in root.findall("./ISA/Instructions/Instruction"):
ins_name = ins_el.findtext("InstructionName")
for ins_enc in ins_el.findall("InstructionEncodings/InstructionEncoding"):
if ins_enc.findtext("EncodingName") == encoding_name:
opcode = int(ins_enc.findtext("Opcode"))
if "OP" in field_data and opcode == field_data["OP"]:
did_match = True
del field_data["OP"]
break
if did_match: break
if did_match: break
#print(ET.tostring(ins_enc).decode())
#print()
#print(field_data)
if not did_match:
print(f"{i:4X} : {ins:16x} -- {encoding_name}")
elif did_match:
params = []
#print(ET.tostring(ins_el).decode())
# 4. Extract the opcodes
for op_ins in ins_enc.findall("Operands/Operand"):
op_type = op_ins.findtext("OperandType")
op_size = op_ins.findtext("OperandSize")
op_fmt = op_ins.findtext("DataFormatName")
op_field_name = op_ins.findtext("FieldName")
if op_field_name is None: continue
assert op_field_name in field_data
# loop through operands for compare
for op_el in root.findall("./ISA/OperandTypes/OperandType"):
test_op_type = op_el.findtext("OperandTypeName")
val_dict = {}
for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
if op_type == test_op_type:
if field_data[op_field_name] in val_dict:
print(op_type, op_size, op_fmt)
params.append(val_dict[field_data[op_field_name]])
else:
params.append(f"{op_type}({field_data[op_field_name]})")
del field_data[op_field_name]
#print(op_type, op_size, op_fmt, op_el, op_field_name,
# field_data[op_field_name],
# val_dict.get(field_data[op_field_name], "<UNK>"))
#print(ET.tostring(op_el).decode())
print(f"{i:4X} : {ins:16x} -- {ins_name.lower()} {', '.join(params)}", field_data)
# advance
i += len(mask) // 8
#print(ET.tostring(root).decode())
if __name__ == "__main__":
# human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
with open("/tmp/rdna35.xml", "wb") as f: f.write(bytes(xml_str))
root = ET.fromstring(xml_str)
a = Tensor.empty(16)+1
for ei in a.schedule():
ei.lower()
# get text
_, hdr, _ = elf_loader(ei.prg.lib)
text = get_single_element([x for x in hdr if x.name==".text"]).content
# llvm disassembler
Device["AMD"].compiler.disassemble(ei.prg.lib)
# run program
ei.run()
sqtt_events = [e for e in Device["AMD"].profile_events if isinstance(e, ProfileSQTTEvent)]
for e in sqtt_events[0:1]: # only the first SE
parse_sqtt_print_packets(e.blob)
disassemble(text[:0x40], root)
+254
View File
@@ -0,0 +1,254 @@
# library for RDNA3 assembly DSL
from __future__ import annotations
from enum import IntEnum
from typing import overload, Annotated, TypeVar, Generic
# Bit field DSL
class BitField:
def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name = hi, lo, name
def __set_name__(self, owner, name): self.name, self._owner = name, owner
def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val) # type: ignore
def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
@property
def marker(self) -> type | None:
# Get marker from Annotated type hint if present
import typing
if hasattr(self, '_owner') and self.name:
hints = typing.get_type_hints(self._owner, include_extras=True)
if self.name in hints:
hint = hints[self.name]
if typing.get_origin(hint) is Annotated:
args = typing.get_args(hint)
return args[1] if len(args) > 1 else None
return None
@overload
def __get__(self, obj: None, objtype: type) -> BitField: ...
@overload
def __get__(self, obj: object, objtype: type | None = None) -> int: ...
def __get__(self, obj, objtype=None):
if obj is None: return self
val = unwrap(obj._values.get(self.name, 0))
# Convert to IntEnum if marker is an IntEnum subclass
if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
try: return self.marker(val)
except ValueError: pass
return val
class _Bits:
def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
bits = _Bits()
# Register types
class Reg:
def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False): self.idx, self.count, self.hi, self.neg = idx, count, hi, neg
def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
def __neg__(self): return self.__class__(self.idx, self.count, self.hi, neg=not self.neg)
T = TypeVar('T', bound=Reg)
class _RegFactory(Generic[T]):
def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
@overload
def __getitem__(self, key: int) -> Reg: ...
@overload
def __getitem__(self, key: slice) -> Reg: ...
def __getitem__(self, key: int | slice) -> Reg:
return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
def __repr__(self): return f"<{self._name} factory>"
class SGPR(Reg): pass
class VGPR(Reg): pass
class TTMP(Reg): pass
s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
# Field type markers (runtime classes for validation)
class _SSrc: pass
class _Src: pass
class _Imm: pass
class _SImm: pass
class _VDSTYEnc: pass # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
class _SGPRField: pass
class _VGPRField: pass
# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
SSrc = Annotated[BitField, _SSrc]
Src = Annotated[BitField, _Src]
Imm = Annotated[BitField, _Imm]
SImm = Annotated[BitField, _SImm]
VDSTYEnc = Annotated[BitField, _VDSTYEnc]
SGPRField = Annotated[BitField, _SGPRField]
VGPRField = Annotated[BitField, _VGPRField]
class RawImm:
def __init__(self, val: int): self.val = val
def __repr__(self): return f"RawImm({self.val})"
def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
def unwrap(val) -> int:
return val.val if isinstance(val, RawImm) else val.value if hasattr(val, 'value') else val.idx if hasattr(val, 'idx') else val
# Encoding helpers
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata'}
def _encode_reg(val) -> int:
if isinstance(val, TTMP): return 108 + val.idx
return val.idx | (0x80 if val.hi else 0)
def encode_src(val) -> int:
if isinstance(val, VGPR): return 256 + _encode_reg(val)
if isinstance(val, Reg): return _encode_reg(val)
if hasattr(val, 'value'): return val.value
if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
# Instruction base class
class Inst:
_fields: dict[str, BitField]
_encoding: tuple[BitField, int] | None = None
_defaults: dict[str, int] = {}
_values: dict[str, int | RawImm]
_words: int # size in 32-bit words, set by decode_program
_literal: int | None
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
def __init__(self, *args, literal: int | None = None, **kwargs):
self._values, self._literal = dict(self._defaults), literal
# Map positional args to field names
field_names = [n for n in self._fields if n != 'encoding']
orig_args = dict(zip(field_names, args))
orig_args.update(kwargs)
self._values.update(orig_args)
# Validate register counts for SMEM instructions (before encoding)
if self.__class__.__name__ == 'SMEM':
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
if op_val is not None:
if hasattr(op_val, 'value'): op_val = op_val.value
expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
sdata_val = orig_args.get('sdata')
if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
# Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
if self.__class__.__name__ == 'SOP1':
op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
if op_val is not None and hasattr(op_val, 'name'):
expected = 2 if op_val.name.endswith('_B64') else 1
sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
if isinstance(sdst_val, Reg) and sdst_val.count != expected:
raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
# Type check and encode values
for name, val in list(self._values.items()):
if name == 'encoding': continue
# For RawImm, only process RAW_FIELDS to unwrap to int
if isinstance(val, RawImm):
if name in RAW_FIELDS: self._values[name] = val.val
continue
field = self._fields.get(name)
marker = field.marker if field else None
# Type validation
if marker is _SGPRField:
if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
if not isinstance(val, (SGPR, TTMP, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
if marker is _VGPRField:
if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
# Encode source fields as RawImm for consistent disassembly
if name in SRC_FIELDS:
encoded = encode_src(val)
self._values[name] = RawImm(encoded)
# Handle negation modifier for VOP3 instructions
if isinstance(val, Reg) and val.neg and 'neg' in self._fields:
neg_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
cur_neg = self._values.get('neg', 0)
self._values['neg'] = (cur_neg.val if isinstance(cur_neg, RawImm) else cur_neg) | neg_bit
# Track literal value if needed (encoded as 255)
if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum):
self._literal = val
elif encoded == 255 and self._literal is None and isinstance(val, float):
import struct
self._literal = struct.unpack('<I', struct.pack('<f', val))[0]
# Encode raw register fields for consistent repr
elif name in RAW_FIELDS:
if isinstance(val, Reg): self._values[name] = _encode_reg(val)
elif hasattr(val, 'value'): self._values[name] = val.value # IntEnum like SrcEnum.NULL
# Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
elif name == 'sbase' and isinstance(val, Reg):
self._values[name] = val.idx // 2
elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
self._values[name] = val.idx // 4
# VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
elif marker is _VDSTYEnc and isinstance(val, VGPR):
self._values[name] = val.idx >> 1
def _encode_field(self, name: str, val) -> int:
if isinstance(val, RawImm): return val.val
if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val
if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
return val.value if hasattr(val, 'value') else val
def to_int(self) -> int:
word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
for n, bf in self._fields.items():
if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
return word
def _get_literal(self) -> int | None:
for n in SRC_FIELDS:
if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
return None
def to_bytes(self) -> bytes:
result = self.to_int().to_bytes(self._size(), 'little')
return result + (lit & 0xffffffff).to_bytes(4, 'little') if (lit := self._get_literal() or getattr(self, '_literal', None)) else result
@classmethod
def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
def size(self) -> int: return self._size() + (4 if self._literal is not None else 0)
@classmethod
def from_int(cls, word: int):
inst = object.__new__(cls)
inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
inst._literal = None
return inst
@classmethod
def from_bytes(cls, data: bytes):
inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
op_val = inst._values.get('op', 0)
has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
for n in SRC_FIELDS:
if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
if has_literal and len(data) >= cls._size() + 4: inst._literal = int.from_bytes(data[cls._size():cls._size()+4], 'little')
return inst
def __repr__(self):
# Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
and not (is_zero(self._values[k]) and k not in {'op'})]
lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
def __eq__(self, other):
if not isinstance(other, Inst): return NotImplemented
return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
def disasm(self) -> str:
from extra.assembly.rdna3.asm import disasm
return disasm(self)
class Inst32(Inst): pass
class Inst64(Inst): pass
-15
View File
@@ -1,15 +0,0 @@
from tinygrad import Tensor, nn
import xml.etree.ElementTree as ET
if __name__ == "__main__":
# human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
root = ET.fromstring(xml_str)
for op_el in root.findall("./ISA/OperandTypes/OperandType"):
op_name = op_el.findtext("OperandTypeName")
val_dict = {}
for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
print(op_name, val_dict)
+294
View File
@@ -0,0 +1,294 @@
#!/usr/bin/env python3
"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
import ctypes, time, os, struct, cProfile, pstats, io
from pathlib import Path
from typing import Callable
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.rdna3.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
def get_rust_remu():
"""Load the Rust libremu shared library."""
if not REMU_PATH.exists(): return None
remu = ctypes.CDLL(str(REMU_PATH))
remu.run_asm.restype = ctypes.c_int32
remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
return remu
def count_instructions(kernel: bytes) -> int:
"""Count instructions in a kernel."""
return len(decode_program(kernel))
def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
"""Allocate buffers and return args pointer + valid ranges."""
if init_data is None: init_data = {}
buffers = []
for i, size in enumerate(buf_sizes):
padded = ((size + 15) // 16) * 16 + 16
data = init_data.get(i, b'\x00' * padded)
data_list = list(data) + [0] * (padded - len(data))
buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
buffers.append(buf)
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
args_ptr = ctypes.addressof(args)
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
ranges.add((args_ptr, ctypes.sizeof(args)))
return buffers, args, args_ptr, ranges
def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
"""Benchmark an emulator and return average time."""
gx, gy, gz = global_size
lx, ly, lz = local_size
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
lib_ptr = ctypes.addressof(kernel_buf)
# Warmup
run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
# Timed runs
times = []
for _ in range(iterations):
start = time.perf_counter()
result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
end = time.perf_counter()
if result != 0:
print(f" {name} returned error: {result}")
return None
times.append(end - start)
return sum(times) / len(times)
def create_synthetic_kernel(n_ops: int) -> bytes:
"""Create a synthetic kernel with n_ops vector operations."""
instructions = []
# VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
ops = [
(0b0000011 << 25) | (1 << 17) | (0 << 9) | 256, # v_add_f32 v0, v0, v1
(0b0001000 << 25) | (1 << 17) | (0 << 9) | 256, # v_mul_f32 v0, v0, v1
(0b0010000 << 25) | (1 << 17) | (0 << 9) | 256, # v_max_f32 v0, v0, v1
(0b0001111 << 25) | (1 << 17) | (0 << 9) | 256, # v_min_f32 v0, v0, v1
]
for i in range(n_ops):
instructions.append(ops[i % len(ops)])
# S_ENDPGM
instructions.append((0b101111111 << 23) | (48 << 16) | 0)
return b''.join(struct.pack('<I', inst) for inst in instructions)
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
try:
from tinygrad import Tensor
from tinygrad.runtime.support.elf import elf_loader
import numpy as np
np.random.seed(42)
ops = {
"add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
"mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
"matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
"matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
"reduce_sum": lambda: Tensor.empty(4096).sum(),
"reduce_max": lambda: Tensor.empty(4096).max(),
"softmax": lambda: Tensor.empty(256).softmax(),
"layernorm": lambda: Tensor.empty(32, 64).layernorm(),
"conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
"gelu": lambda: Tensor.empty(1024).gelu(),
"exp": lambda: Tensor.empty(1024).exp(),
"sin": lambda: Tensor.empty(1024).sin(),
}
if op_name not in ops: return None
out = ops[op_name]()
sched = out.schedule()
for ei in sched:
lowered = ei.lower()
if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
lib = bytes(lowered.prg.p.lib)
_, sections, _ = elf_loader(lib)
for sec in sections:
if sec.name == '.text':
buf_sizes = [b.nbytes for b in lowered.bufs]
# Get initial data from numpy arrays if available
buf_data = {}
for i, buf in enumerate(lowered.bufs):
if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
try: buf_data[i] = bytes(buf.base._buf)
except: pass
return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
return None
except Exception as e:
print(f" Error getting kernel: {e}")
return None
def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
"""Profile the Python emulator to find bottlenecks."""
gx, gy, gz = global_size
lx, ly, lz = local_size
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
lib_ptr = ctypes.addressof(kernel_buf)
pr = cProfile.Profile()
pr.enable()
for _ in range(n_runs):
python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
pr.disable()
s = io.StringIO()
ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
ps.print_stats(20)
return s.getvalue()
def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
"""Measure raw step_wave() performance (steps per second)."""
program = decode_program(kernel)
if not program: return 0.0
st = WaveState()
st.exec_mask = 0xffffffff
lds = bytearray(65536)
n_lanes = 32
# Reset PC for each measurement
start = time.perf_counter()
for _ in range(n_steps):
st.pc = 0
while st.pc in program:
result = step_wave(program, st, lds, n_lanes)
if result == -1: break
elapsed = time.perf_counter() - start
return n_steps / elapsed if elapsed > 0 else 0
# Test configurations
SYNTHETIC_TESTS = [
("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
]
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
def main():
import argparse
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
args = parser.parse_args()
rust_remu = get_rust_remu()
if rust_remu is None:
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
print("Running Python-only benchmarks...\n")
print("=" * 90)
print("RDNA3 Emulator Benchmark: Python vs Rust")
print("=" * 90)
results = []
# Synthetic workloads
if not args.tinygrad_only:
print("\n[SYNTHETIC WORKLOADS]")
print("-" * 90)
for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
kernel = create_synthetic_kernel(n_ops)
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buf_sizes = [4096]
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
set_valid_mem_ranges(ranges)
# Benchmark
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
# Tinygrad kernels
if not args.synthetic_only:
print("\n[TINYGRAD KERNELS]")
print("-" * 90)
for op_name in TINYGRAD_TESTS:
print(f"\n{op_name}:", end=" ", flush=True)
kernel_info = get_tinygrad_kernel(op_name)
if kernel_info is None:
print("failed to compile")
continue
kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
n_insts = count_instructions(kernel)
n_workgroups = global_size[0] * global_size[1] * global_size[2]
n_threads = local_size[0] * local_size[1] * local_size[2]
total_work = n_insts * n_workgroups * n_threads
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
if py_time:
py_rate = total_work / py_time / 1e6
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
if rust_time:
rust_rate = total_work / rust_time / 1e6
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
# Optional profiling
if args.profile and py_time:
print("\n [PROFILE - Top 10 functions]")
profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
for line in profile_output.split('\n')[5:15]:
if line.strip(): print(f" {line}")
# Summary table
print("\n" + "=" * 90)
print("SUMMARY")
print("=" * 90)
print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print("-" * 90)
for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
if rust_time:
rust_ms = f"{rust_time*1000:.3f}"
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
else:
rust_ms, speedup = "N/A", "N/A"
print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
if __name__ == "__main__":
main()
@@ -0,0 +1,466 @@
# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
import unittest, ctypes, os
from dataclasses import dataclass
from pathlib import Path
# Set environment before any tinygrad imports to use MOCKGPU
# This allows generating AMD GPU kernels without requiring real hardware
os.environ["AMD"] = "1"
os.environ["MOCKGPU"] = "1"
from extra.assembly.rdna3.emu import WaveState, decode_program, step_wave, WAVE_SIZE
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
@dataclass
class KernelInfo:
code: bytes
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
buf_sizes: list[int] # sizes for each buffer index
@dataclass
class StateSnapshot:
pc: int
scc: int
vcc: int
exec_mask: int
sgpr: list[int]
vgpr: list[list[int]]
def diff(self, other: 'StateSnapshot', n_lanes: int) -> list[str]:
"""Return list of differences between two states."""
diffs = []
if self.pc != other.pc: diffs.append(f"pc: {self.pc} vs {other.pc}")
if self.scc != other.scc: diffs.append(f"scc: {self.scc} vs {other.scc}")
if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x} vs 0x{other.vcc:08x}")
if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x} vs 0x{other.exec_mask:08x}")
for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
# Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
if i in (106, 107, 126, 127): continue
if a != b: diffs.append(f"sgpr[{i}]: 0x{a:08x} vs 0x{b:08x}")
for lane in range(n_lanes):
for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
if a != b: diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x} vs 0x{b:08x}")
return diffs
class CStateSnapshot(ctypes.Structure):
_fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
def to_snapshot(self) -> StateSnapshot:
return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
class RustEmulator:
def __init__(self):
self.lib = ctypes.CDLL(str(REMU_PATH))
self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
self.lib.wave_create.restype = ctypes.c_void_p
self.lib.wave_step.argtypes = [ctypes.c_void_p]
self.lib.wave_step.restype = ctypes.c_int32
self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
self.lib.wave_free.argtypes = [ctypes.c_void_p]
self.ctx = None
def create(self, kernel: bytes, n_lanes: int):
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
self._kernel_buf = kernel_buf
def step(self) -> int: return self.lib.wave_step(self.ctx)
def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
def get_snapshot(self) -> StateSnapshot:
snap = CStateSnapshot()
self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
return snap.to_snapshot()
def free(self):
if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
class PythonEmulator:
def __init__(self):
self.state: WaveState | None = None
self.program: dict | None = None
self.lds: bytearray | None = None
self.n_lanes = 0
def create(self, kernel: bytes, n_lanes: int):
self.program = decode_program(kernel)
self.state = WaveState()
self.state.exec_mask = (1 << n_lanes) - 1
self.lds = bytearray(65536)
self.n_lanes = n_lanes
def step(self) -> int:
assert self.program is not None and self.state is not None and self.lds is not None
return step_wave(self.program, self.state, self.lds, self.n_lanes)
def set_sgpr(self, idx: int, val: int):
assert self.state is not None
self.state.sgpr[idx] = val & 0xffffffff
def set_vgpr(self, lane: int, idx: int, val: int):
assert self.state is not None
self.state.vgpr[lane][idx] = val & 0xffffffff
def get_snapshot(self) -> StateSnapshot:
assert self.state is not None
return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
max_workgroups: int = 64) -> tuple[bool, str, int]:
"""Run a single kernel through both emulators. Returns (success, message, total_steps)."""
gx, gy, gz = global_size
total_steps = 0
wg_count = 0
for gidz in range(gz):
for gidy in range(gy):
for gidx in range(gx):
if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
wg_count += 1
rust = RustEmulator()
python = PythonEmulator()
rust.create(kernel, n_lanes)
python.create(kernel, n_lanes)
# Initialize LDS (64KB, standard size for AMD GPUs)
rust.init_lds(65536)
for emu in (rust, python):
emu.set_sgpr(0, args_ptr & 0xffffffff)
emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
emu.set_sgpr(13, gidx)
emu.set_sgpr(14, gidy)
emu.set_sgpr(15, gidz)
step = 0
trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
try:
while step < max_steps:
rust_before = rust.get_snapshot()
python_before = python.get_snapshot()
inst = program.get(python_before.pc)
inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
trace.append((step, python_before.pc, inst_str, rust_before, python_before))
if len(trace) > trace_len: trace.pop(0)
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
diffs = rust_before.diff(python_before, n_lanes)
if diffs:
trace_lines = []
for s, pc, d, rb, pb in trace[:-1]:
trace_lines.append(f" step {s}: PC={pc:3d} {d}")
if trace.index((s, pc, d, rb, pb)) < len(trace) - 2:
next_rb, next_pb = trace[trace.index((s, pc, d, rb, pb)) + 1][3:5]
inst_diffs = rb.diff(next_rb, n_lanes)
if inst_diffs: trace_lines.append(f" rust changes: {', '.join(inst_diffs[:3])}")
trace_str = "\n".join(trace_lines)
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ:\n " + "\n ".join(diffs[:10]) + f"\n Recent instructions:\n{trace_str}", total_steps
rust_result = rust.step()
python_result = python.step()
if rust_result != python_result:
trace_str = "\n".join(f" step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n Recent instructions:\n{trace_str}", total_steps
if rust_result == -1:
total_steps += step + 1
break
if rust_result == 1:
total_steps += step + 1
break
if rust_result < 0 and rust_result != -2:
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
step += 1
else:
return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
finally:
rust.free()
return True, f"Completed {gx*gy*gz} workgroups", total_steps
def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
"""Run all kernels through both emulators with shared buffer pool."""
from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
if buf_data is None: buf_data = {}
# Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
buf_id_to_ptr: dict[int, int] = {}
buffers = []
for buf_id, size in buf_pool.items():
padded_size = ((size + 15) // 16) * 16 + 16 # round up to 16 bytes + extra padding
# Initialize with data from COPY if available
init_data = buf_data.get(buf_id, b'\x00' * padded_size)
init_list = list(init_data) + [0] * (padded_size - len(init_data))
buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
buffers.append((buf, padded_size))
buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
# Set up valid memory ranges
ranges = {(ctypes.addressof(b), size) for b, size in buffers}
total_steps = 0
for ki, kernel in enumerate(kernels):
# Create args array for this kernel's buffers
args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
args_ptr = ctypes.addressof(args)
# Update valid ranges to include this args array
kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
set_valid_mem_ranges(kernel_ranges)
program = decode_program(kernel.code)
n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
ok, msg, steps = run_single_kernel(
kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
program, max_steps, debug, trace_len, ki
)
total_steps += steps
if not ok:
return False, msg
return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
"""Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
# Allocate buffers
buffers = []
for size in buf_sizes:
buf = (ctypes.c_uint8 * size)(*[0] * size)
buffers.append(buf)
# Create args array with buffer pointers
args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
args_ptr = ctypes.addressof(args)
# Set up valid memory ranges for Python emulator
ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
ranges.add((args_ptr, ctypes.sizeof(args)))
set_valid_mem_ranges(ranges)
program = decode_program(kernel)
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
return ok, msg
def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
"""Compile a tinygrad operation and extract all kernels with their buffer mappings."""
from tinygrad import Tensor
from tinygrad.runtime.support.elf import elf_loader
out = op_fn(Tensor)
sched = out.schedule()
kernels = []
buf_pool: dict[int, int] = {} # buffer id -> size
buf_data: dict[int, bytes] = {} # buffer id -> initial data from COPY
for ei in sched:
lowered = ei.lower()
if ei.ast.op.name == 'COPY':
# Handle COPY: extract source data to initialize destination buffer
if len(lowered.bufs) >= 2:
dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
dst_id = id(dst_buf)
if dst_id not in buf_pool:
buf_pool[dst_id] = dst_buf.nbytes
# Get source data if it's from numpy/CPU
if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
src_data = bytes(src_buf.base._buf)
buf_data[dst_id] = src_data
elif ei.ast.op.name == 'SINK':
if lowered.prg and lowered.prg.p.lib:
lib = bytes(lowered.prg.p.lib)
_, sections, _ = elf_loader(lib)
for sec in sections:
if sec.name == '.text':
buf_idxs = []
buf_sizes = []
for b in lowered.bufs:
buf_id = id(b)
if buf_id not in buf_pool:
buf_pool[buf_id] = b.nbytes
buf_idxs.append(buf_id)
buf_sizes.append(b.nbytes)
kernels.append(KernelInfo(
code=bytes(sec.content),
global_size=tuple(lowered.prg.p.global_size),
local_size=tuple(lowered.prg.p.local_size),
buf_idxs=buf_idxs,
buf_sizes=buf_sizes
))
if not kernels: raise RuntimeError("No kernel found")
return kernels, buf_pool, buf_data
def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
"""Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
k = kernels[-1]
return k.code, k.global_size, k.local_size, k.buf_sizes
@unittest.skipUnless(REMU_PATH.exists(), "libremu.so not found")
class TestTinygradKernels(unittest.TestCase):
"""Compare emulators on real tinygrad-compiled kernels."""
def _test_kernel(self, op_fn, max_steps=10000):
kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
self.assertTrue(ok, msg)
# Basic unary ops
def test_neg(self): self._test_kernel(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Reductions
def test_sum_reduce(self): self._test_kernel(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel(lambda T: T.empty(32).mean())
# Matmul - various sizes
def test_gemm_4x4(self): self._test_kernel(lambda T: T.empty(4, 4) @ T.empty(4, 4), max_steps=100000)
def test_gemm_8x8(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=200000)
@unittest.skip("too slow")
def test_gemm_16x16(self): self._test_kernel(lambda T: T.empty(16, 16) @ T.empty(16, 16), max_steps=500000)
def test_gemv(self): self._test_kernel(lambda T: T.empty(1, 16) @ T.empty(16, 16), max_steps=100000)
# Complex ops
def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
# Cast ops
def test_cast_int(self): self._test_kernel(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel(lambda T: T.empty(16).half().float())
# Min/max (uses comparison internally)
def test_min_binary(self): self._test_kernel(lambda T: T([1.0, 5.0, 3.0]).minimum(T([3.0, 2.0, 4.0])))
# Comparison ops (test VOPC instructions) - use 32+ elements to force vector instructions
def test_cmp_lt(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_cmp_eq(self): self._test_kernel(lambda T: (T.empty(64) == T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Bitwise ops
def test_bitwise_and(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0xFF]).int() & T([0x0F, 0x0F, 0x00]).int())
def test_bitwise_or(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0x00]).int() | T([0x0F, 0x0F, 0xFF]).int())
def test_bitwise_xor(self): self._test_kernel(lambda T: T([0xFF, 0x0F, 0xF0]).int() ^ T([0x0F, 0xF0, 0xF0]).int())
# Integer ops - use 32+ elements to force vector instructions
def test_int_add(self): self._test_kernel(lambda T: (T.empty(64).int() + T.empty(64).int()).float())
def test_int_mul(self): self._test_kernel(lambda T: (T.empty(64).int() * T.empty(64).int()).float())
def test_int_mod(self): self._test_kernel(lambda T: (T.empty(64).int().abs() % (T.empty(64).int().abs() + 1)).float())
# More math ops - use 32+ elements to force vector instructions
def test_abs(self): self._test_kernel(lambda T: T.empty(64).abs())
def test_floor(self): self._test_kernel(lambda T: T.empty(64).floor())
def test_ceil(self): self._test_kernel(lambda T: T.empty(64).ceil())
def test_trunc(self): self._test_kernel(lambda T: T.empty(64).trunc())
# Fused ops
def test_fma(self): self._test_kernel(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
# Argmax/argmin (tests different reduction pattern) - use 32+ elements to force vector instructions
def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
def test_argmin(self): self._test_kernel(lambda T: T.empty(64).argmin())
# Exact value tests - use 32+ elements to force vector instructions (small tensors use scalar ops which Rust emu doesn't fully support)
def test_abs_exact(self): self._test_kernel(lambda T: T([-1., 0., 1.]*11).abs()) # 33 elements
def test_neg_exact(self): self._test_kernel(lambda T: -T([-1., 0., 1.]*11))
def test_log_special(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).log())
def test_exp_exact(self): self._test_kernel(lambda T: T([0., 1., -1.]*11).exp())
def test_reciprocal_exact(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).reciprocal())
# Integer division and mod - use 32+ elements
def test_int_div(self): self._test_kernel(lambda T: (T([10, 20, 30]*11).int() // T([3, 4, 5]*11).int()).float())
def test_int_neg(self): self._test_kernel(lambda T: (-T([1, -2, 3]*11).int()).float())
# Mixed precision - use 32+ elements
def test_half_add(self): self._test_kernel(lambda T: (T([1., 2.]*16).half() + T([3., 4.]*16).half()).float())
def test_half_mul(self): self._test_kernel(lambda T: (T([2., 3.]*16).half() * T([4., 5.]*16).half()).float())
# Matrix ops - patterns from test_ops.py failures
def test_cat(self): self._test_kernel(lambda T: T.empty(32, 64).cat(T.empty(32, 64), dim=1))
def test_gather(self): self._test_kernel(lambda T: T.empty(64).gather(0, T.arange(32).int()))
# Tests from test_ops.py that are failing
def test_permute(self): self._test_kernel(lambda T: T.empty(3, 4, 5, 6).permute((3, 2, 1, 0)).contiguous())
def test_cat_large(self): self._test_kernel(lambda T: T.empty(45, 65, 9).cat(T.empty(45, 65, 9), T.empty(45, 65, 9), dim=1))
def test_gather_small(self): self._test_kernel(lambda T: T.empty(10).gather(0, T.arange(5).int()))
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
def test_cross_entropy(self): self._test_kernel(lambda T: T.randn(32, 10).softmax().log().sum())
def test_cross_entropy_class(self):
import numpy as np
np.random.seed(0)
classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
x_np = np.random.randn(32, 10).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()).reshape(32,10) + 0).cross_entropy((T(classes).int().reshape(32) + 0)))
# Regression tests for BFE operations with width=0 (walrus operator bug)
def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
def test_interpolate_uint8(self): self._test_kernel(lambda T: T.empty(2,3,64,64).relu().cast('uint8').interpolate((10,10), mode="linear"))
# Regression test for 64-bit comparison (V_CMP_GT_I64, V_CMP_LT_U64, etc.) with rsrc64
def test_index_int64(self):
from tinygrad import dtypes
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
@unittest.skip("only works with mock GPU")
def test_index_int64_2d(self):
from tinygrad import dtypes
# Tests 64-bit compare with inline constants (comparing against 0)
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), T.arange(4).cast(dtypes.int64)])
# Pooling operations - regression test for VCC wave32 mode (S_CBRANCH_VCCZ should only check VCC_LO)
def test_avg_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4), stride=2))
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
def test_avg_pool3d(self):
import numpy as np
np.random.seed(0)
self._test_kernel(lambda T: T(np.random.randn(1, 1, 16, 16, 16).astype(np.float32).tolist()).avg_pool2d(kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False))
def test_max_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4), stride=2))
# Convolution operations - multi-kernel tests
def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv2d(T.empty(4, 4, 3, 3)), max_steps=100000)
def test_conv_transpose2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv_transpose2d(T.empty(4, 4, 3, 3)), max_steps=200000)
@unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
def test_conv_transpose3d(self):
import numpy as np
np.random.seed(0)
self._test_kernel(lambda T: T(np.random.randn(2, 4, 9, 9, 9).astype(np.float32).tolist()).conv_transpose2d(
T(np.random.randn(4, 4, 3, 3, 3).astype(np.float32).tolist())), max_steps=500000)
if __name__ == "__main__":
unittest.main()
+845
View File
@@ -0,0 +1,845 @@
# Unit tests for RDNA3 Python emulator
import unittest
import ctypes
import struct
import math
from extra.assembly.rdna3.emu import (
WaveState, decode_program, exec_wave, exec_workgroup, run_asm,
i32, f32, sext, WAVE_SIZE, set_valid_mem_ranges
)
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import RawImm
def run_kernel(kernel: bytes, n_threads: int = 1, n_outputs: int = 1) -> list[int]:
"""Helper to run a kernel and return output values."""
output = (ctypes.c_uint32 * (n_threads * n_outputs))(*[0xdead] * (n_threads * n_outputs))
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
# Register valid memory ranges for bounds checking
set_valid_mem_ranges({
(output_ptr, ctypes.sizeof(output)),
(args_ptr, ctypes.sizeof(args)),
(kernel_ptr, len(kernel)),
})
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
assert result == 0, f"run_asm failed with {result}"
return [output[i] for i in range(n_threads * n_outputs)]
def make_store_kernel(setup_instrs: list, store_vreg: int = 1) -> bytes:
"""Create a kernel that runs setup instructions then stores v[store_vreg] to output[tid]."""
kernel = b''
# Load output pointer
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
# Run setup instructions
for instr in setup_instrs:
kernel += instr.to_bytes()
# Compute offset: v3 = tid * 4
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
# Store result
kernel += global_store_b32(addr=v[3], data=v[store_vreg], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
return kernel
class TestScalarOps(unittest.TestCase):
def test_s_mov_b32(self):
state = WaveState()
kernel = s_mov_b32(s[5], 42).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[5], 42)
def test_s_add_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 100, 50
kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 150)
self.assertEqual(state.scc, 0) # no carry
def test_s_add_u32_carry(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 0xffffffff, 1
kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 0)
self.assertEqual(state.scc, 1) # carry
def test_s_sub_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 100, 30
kernel = s_sub_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 70)
self.assertEqual(state.scc, 0) # no borrow
def test_s_and_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 0xff00, 0x0ff0
kernel = s_and_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 0x0f00)
def test_s_or_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 0xff00, 0x00ff
kernel = s_or_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 0xffff)
def test_s_lshl_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 1, 4
kernel = s_lshl_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 16)
def test_s_lshr_b32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 256, 4
kernel = s_lshr_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 16)
def test_s_mul_i32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 7, 6
kernel = s_mul_i32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[2], 42)
def test_s_cmp_eq_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 42, 42
kernel = s_cmp_eq_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.scc, 1)
def test_s_cmp_lg_u32(self):
state = WaveState()
state.sgpr[0], state.sgpr[1] = 42, 43
kernel = s_cmp_lg_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.scc, 1)
class TestVectorOps(unittest.TestCase):
def test_v_mov_b32(self):
kernel = make_store_kernel([v_mov_b32_e32(v[1], 42)])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_add_nc_u32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 10),
v_mov_b32_e32(v[2], 32),
v_add_nc_u32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_sub_nc_u32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 50),
v_mov_b32_e32(v[2], 8),
v_sub_nc_u32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_mul_lo_u32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 6),
v_mov_b32_e32(v[2], 7),
v_mul_lo_u32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [42])
def test_v_and_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 0xff0f),
v_mov_b32_e32(v[2], 0x0fff),
v_and_b32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [0x0f0f])
def test_v_or_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 0xf000),
v_mov_b32_e32(v[2], 0x000f),
v_or_b32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [0xf00f])
def test_v_lshlrev_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 1),
v_lshlrev_b32_e32(v[1], 5, v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [32])
def test_v_lshrrev_b32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 128),
v_lshrrev_b32_e32(v[1], 3, v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out, [16])
def test_v_add_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(1.5)),
v_mov_b32_e32(v[2], i32(2.5)),
v_add_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 4.0)
def test_v_mul_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.0)),
v_mov_b32_e32(v[2], i32(4.0)),
v_mul_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 12.0)
def test_v_max_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.0)),
v_mov_b32_e32(v[2], i32(5.0)),
v_max_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 5.0)
def test_v_min_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.0)),
v_mov_b32_e32(v[2], i32(5.0)),
v_min_f32_e32(v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 3.0)
class TestThreading(unittest.TestCase):
def test_thread_id(self):
"""Each thread should get its own thread ID in v0."""
kernel = make_store_kernel([v_mov_b32_e32(v[1], v[0])], store_vreg=1)
out = run_kernel(kernel, n_threads=4)
self.assertEqual(out, [0, 1, 2, 3])
def test_thread_local_ops(self):
"""Each thread computes tid * 10."""
kernel = make_store_kernel([
v_mov_b32_e32(v[2], 10),
v_mul_lo_u32(v[1], v[0], v[2]),
])
out = run_kernel(kernel, n_threads=4)
self.assertEqual(out, [0, 10, 20, 30])
def test_exec_mask(self):
"""Test that exec mask controls which lanes execute."""
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], 0, soffset=NULL).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_mov_b32_e32(v[1], 100).to_bytes() # default value
kernel += s_mov_b32(EXEC_LO, 0b0101).to_bytes() # only lanes 0 and 2
kernel += v_mov_b32_e32(v[1], 42).to_bytes() # only for active lanes
kernel += s_mov_b32(EXEC_LO, 0xf).to_bytes() # restore all lanes
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
out = run_kernel(kernel, n_threads=4)
self.assertEqual(out, [42, 100, 42, 100])
class TestBranching(unittest.TestCase):
def test_s_branch(self):
"""Test unconditional branch."""
state = WaveState()
kernel = b''
kernel += s_mov_b32(s[0], 1).to_bytes()
kernel += s_branch(1).to_bytes() # skip next instruction
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
kernel += s_mov_b32(s[1], 3).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[0], 1) # not overwritten
self.assertEqual(state.sgpr[1], 3)
def test_s_cbranch_scc0(self):
"""Test conditional branch on SCC=0."""
state = WaveState()
state.scc = 0
kernel = b''
kernel += s_mov_b32(s[0], 1).to_bytes()
kernel += s_cbranch_scc0(1).to_bytes() # branch if scc=0
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[0], 1)
def test_s_cbranch_scc1(self):
"""Test conditional branch on SCC=1."""
state = WaveState()
state.scc = 1
kernel = b''
kernel += s_mov_b32(s[0], 1).to_bytes()
kernel += s_cbranch_scc1(1).to_bytes() # branch if scc=1
kernel += s_mov_b32(s[0], 2).to_bytes() # should be skipped
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.sgpr[0], 1)
def test_unknown_sopp_opcode(self):
"""Regression test: unknown SOPP opcodes should be ignored, not crash."""
state = WaveState()
# Create a raw SOPP instruction with opcode 8 (undefined in our enum)
# SOPP format: bits[31:23] = 0b101111111, bits[22:16] = op, bits[15:0] = simm16
unknown_sopp = (0b101111111 << 23) | (8 << 16) | 0 # op=8, simm16=0
kernel = unknown_sopp.to_bytes(4, 'little') + s_endpgm().to_bytes()
prog = decode_program(kernel)
# Should not raise an exception
exec_wave(prog, state, bytearray(65536), 1)
class TestMemory(unittest.TestCase):
def test_global_load_store(self):
"""Test global load followed by store."""
# Create input buffer
input_buf = (ctypes.c_uint32 * 4)(10, 20, 30, 40)
input_ptr = ctypes.addressof(input_buf)
output_buf = (ctypes.c_uint32 * 4)(*[0]*4)
output_ptr = ctypes.addressof(output_buf)
args = (ctypes.c_uint64 * 2)(output_ptr, input_ptr)
args_ptr = ctypes.addressof(args)
# Kernel: load from input[tid], add 1, store to output[tid]
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes() # output ptr
kernel += s_load_b64(s[4:5], s[0:1], soffset=NULL, offset=8).to_bytes() # input ptr
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_lshlrev_b32_e32(v[2], 2, v[0]).to_bytes() # offset = tid * 4
kernel += global_load_b32(vdst=v[1], addr=v[2], saddr=s[4]).to_bytes()
kernel += s_waitcnt(vmcnt=0).to_bytes()
kernel += v_add_nc_u32_e32(v[1], 1, v[1]).to_bytes() # add 1
kernel += global_store_b32(addr=v[2], data=v[1], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({
(input_ptr, ctypes.sizeof(input_buf)),
(output_ptr, ctypes.sizeof(output_buf)),
(args_ptr, ctypes.sizeof(args)),
(kernel_ptr, len(kernel)),
})
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, 4, 1, 1, args_ptr)
self.assertEqual(result, 0)
self.assertEqual([output_buf[i] for i in range(4)], [11, 21, 31, 41])
class TestFloatOps(unittest.TestCase):
def test_v_rcp_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(4.0)),
v_rcp_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertAlmostEqual(f32(out[0]), 0.25, places=5)
def test_v_sqrt_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(16.0)),
v_sqrt_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertAlmostEqual(f32(out[0]), 4.0, places=5)
def test_v_floor_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.7)),
v_floor_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 3.0)
def test_v_ceil_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(3.2)),
v_ceil_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 4.0)
def test_v_cvt_f32_i32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 42),
v_cvt_f32_i32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 42.0)
def test_v_cvt_i32_f32(self):
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(42.9)),
v_cvt_i32_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 42)
class TestVOP3(unittest.TestCase):
def test_v_fma_f32(self):
"""Test fused multiply-add: a*b + c"""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(2.0)),
v_mov_b32_e32(v[2], i32(3.0)),
v_mov_b32_e32(v[4], i32(4.0)),
v_fma_f32(v[1], v[1], v[2], v[4]), # 2*3+4 = 10
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 10.0)
def test_v_add3_u32(self):
"""Test 3-operand add."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 10),
v_mov_b32_e32(v[2], 20),
v_mov_b32_e32(v[4], 12),
v_add3_u32(v[1], v[1], v[2], v[4]), # 10+20+12 = 42
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 42)
def test_v_neg_modifier(self):
"""Test VOP3 negation modifier."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(5.0)),
v_mov_b32_e32(v[2], i32(3.0)),
# v_add_f32 with neg on src1: 5 + (-3) = 2
v_add_f32_e64(v[1], v[1], v[2], neg=0b010),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 2.0)
def test_v_ldexp_f32(self):
"""Regression test: V_LDEXP_F32 used by exp()."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(1.5)),
v_mov_b32_e32(v[2], 3), # exponent
v_ldexp_f32(v[1], v[1], v[2]), # 1.5 * 2^3 = 12.0
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(f32(out[0]), 12.0)
def test_v_xad_u32(self):
"""Regression test: V_XAD_U32 (xor-add) used by random number generation."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 3),
v_mov_b32_e32(v[2], 4),
v_mov_b32_e32(v[4], 5),
v_xad_u32(v[1], v[1], v[2], v[4]), # (3^4)+5 = 7+5 = 12
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 12)
def test_v_lshl_or_b32(self):
"""Regression test: V_LSHL_OR_B32 operand order is (s0 << s1) | s2, not (s0 << s2) | s1."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 5), # s0 = value to shift
v_mov_b32_e32(v[2], 2), # s1 = shift amount
v_mov_b32_e32(v[4], 3), # s2 = value to OR
v_lshl_or_b32(v[1], v[1], v[2], v[4]), # (5 << 2) | 3 = 20 | 3 = 23
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0], 23)
def test_v_sqrt_f32_negative(self):
"""Regression test: V_SQRT_F32 should return NaN for negative inputs, not 0."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(-1.0)),
v_sqrt_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertTrue(math.isnan(f32(out[0])))
def test_v_rsq_f32_negative(self):
"""Regression test: V_RSQ_F32 should return NaN for negative inputs, not inf."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i32(-1.0)),
v_rsq_f32_e32(v[1], v[1]),
])
out = run_kernel(kernel, n_threads=1)
self.assertTrue(math.isnan(f32(out[0])))
class TestVOPD(unittest.TestCase):
def test_vopd_add_nc_u32(self):
"""Test VOPD V_DUAL_ADD_NC_U32."""
state = WaveState()
state.vgpr[0][1] = 100
state.vgpr[0][2] = 50
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][3], 100) # MOV result
self.assertEqual(state.vgpr[0][4], 150) # 100 + 50
def test_vopd_lshlrev(self):
"""Test VOPD V_DUAL_LSHLREV_B32."""
state = WaveState()
state.vgpr[0][1] = 0x10
state.vgpr[0][2] = 0
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
opy=VOPDOp.V_DUAL_LSHLREV_B32, srcy0=4, vsrcy1=VGPR(1), vdsty=VGPR(4)).to_bytes() # V4 = V1 << 4
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][3], 0x10) # MOV result
self.assertEqual(state.vgpr[0][4], 0x100) # 0x10 << 4 = 0x100
def test_vopd_and(self):
"""Test VOPD V_DUAL_AND_B32."""
state = WaveState()
state.vgpr[0][1] = 0xff
state.vgpr[0][2] = 0x0f
# vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
opy=VOPDOp.V_DUAL_AND_B32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][3], 0xff)
self.assertEqual(state.vgpr[0][4], 0x0f) # 0xff & 0x0f = 0x0f
def test_vopd_parallel_read(self):
"""Regression: VOPD must read all inputs before writing - Y op reads register that X op writes."""
state = WaveState()
state.vgpr[0][4] = 0
state.vgpr[0][7] = 5 # Y op reads v7 as vsrcy1, X op writes to v7
# X: MOV v7, v0 (v0=0, so v7 becomes 0)
# Y: ADD v6, v4, v7 (should use original v7=5, not the overwritten 0)
# vdsty_enc=3 with vdstx=7 (odd) -> vdsty = (3 << 1) | (7&1)^1 = 6 | 0 = 6
kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[0], vsrcx1=VGPR(0), vdstx=VGPR(7),
opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[4], vsrcy1=VGPR(7), vdsty=VGPR(6)).to_bytes()
kernel += s_endpgm().to_bytes()
prog = decode_program(kernel)
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][7], 0) # X op: v7 = v0 = 0
self.assertEqual(state.vgpr[0][6], 5) # Y op: v6 = v4 + v7 = 0 + 5 (original v7)
class TestDecoder(unittest.TestCase):
def test_vopd_literal_handling(self):
"""Regression test: VOPD srcx0/srcy0 with literal (255) wasn't consuming the literal dword."""
state = WaveState()
# Create VOPD with srcx0=255 (literal), followed by literal value 0x12345678
vopd_bytes = VOPD(opx=8, srcx0=RawImm(255), vsrcx1=VGPR(0), vdstx=VGPR(1), # MOV: V1 = literal
opy=8, srcy0=RawImm(128), vsrcy1=VGPR(0), vdsty=VGPR(2)).to_bytes() # MOV: V2 = 0
literal_bytes = (0x12345678).to_bytes(4, 'little')
kernel = vopd_bytes + literal_bytes + s_endpgm().to_bytes()
prog = decode_program(kernel)
# Should decode as 3 instructions: VOPD (with literal), then S_ENDPGM
# The literal should NOT be decoded as a separate instruction
self.assertEqual(len(prog), 2) # VOPD + S_ENDPGM
exec_wave(prog, state, bytearray(65536), 1)
self.assertEqual(state.vgpr[0][1], 0x12345678)
def test_s_endpgm_stops_decode(self):
"""Regression test: decoder should stop at S_ENDPGM, not read past into metadata."""
# Create a kernel followed by garbage that looks like an invalid instruction
kernel = s_mov_b32(s[0], 42).to_bytes() + s_endpgm().to_bytes()
garbage = bytes([0xff] * 16) # garbage after kernel
prog = decode_program(kernel + garbage)
# Should only have 2 instructions (s_mov_b32 and s_endpgm)
self.assertEqual(len(prog), 2)
class TestFloatConversion(unittest.TestCase):
"""Unit tests for i32/i16/f32/f16 float conversion functions."""
def test_i32_preserves_nan_sign(self):
"""NaN sign bit should be preserved when converting float to int bits."""
from extra.assembly.rdna3.emu import i32, f32
# 0 * -inf produces a negative NaN
neg_nan = 0.0 * float('-inf')
bits = i32(neg_nan)
# Should have sign bit set (0xffc00000), not canonical positive NaN (0x7fc00000)
self.assertEqual(bits & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{bits:08x}")
self.assertTrue(math.isnan(f32(bits)))
def test_i32_preserves_positive_nan(self):
"""Positive NaN should remain positive."""
from extra.assembly.rdna3.emu import i32, f32
pos_nan = float('nan')
bits = i32(pos_nan)
# Standard Python NaN is positive (0x7fc00000)
self.assertEqual(bits & 0x80000000, 0, f"Expected positive NaN, got 0x{bits:08x}")
self.assertTrue(math.isnan(f32(bits)))
def test_i32_overflow_to_inf(self):
"""Values too large for f32 should become inf."""
from extra.assembly.rdna3.emu import i32, f32
big = 2.0 ** 200
self.assertEqual(i32(big), 0x7f800000) # +inf
self.assertEqual(i32(-big), 0xff800000) # -inf
def test_i32_inf(self):
"""Infinity should be preserved."""
from extra.assembly.rdna3.emu import i32
self.assertEqual(i32(float('inf')), 0x7f800000)
self.assertEqual(i32(float('-inf')), 0xff800000)
def test_i32_normal_values(self):
"""Normal float values should round-trip correctly (within f32 precision)."""
from extra.assembly.rdna3.emu import i32, f32
# Use values exactly representable in float32
for val in [0.0, 1.0, -1.0, 0.5, -0.5, 100.0, -100.0, 1e10]:
bits = i32(val)
self.assertAlmostEqual(f32(bits), val, places=5)
def test_i16_overflow_to_inf(self):
"""Values too large for f16 should become inf."""
from extra.assembly.rdna3.emu import i16
big = 100000.0 # way larger than f16 max (65504)
self.assertEqual(i16(big), 0x7c00) # +inf
self.assertEqual(i16(-big), 0xfc00) # -inf
def test_i16_inf(self):
"""Infinity should be preserved."""
from extra.assembly.rdna3.emu import i16
self.assertEqual(i16(float('inf')), 0x7c00)
self.assertEqual(i16(float('-inf')), 0xfc00)
def test_fma_nan_sign_preserved(self):
"""FMA producing NaN should preserve the correct sign bit."""
from extra.assembly.rdna3.emu import i32, f32
# 0 * (-inf) + 1.0 = NaN (from 0 * -inf)
a, b, c = 0.0, float('-inf'), 1.0
result = i32(a * b + c)
# The NaN should be negative since 0 * -inf produces negative NaN
self.assertEqual(result & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{result:08x}")
class TestMultiWave(unittest.TestCase):
def test_all_waves_execute(self):
"""Regression test: all waves in a workgroup must execute, not just the first."""
n_threads = 64 # 2 waves of 32 threads each
output = (ctypes.c_uint32 * n_threads)(*[0xdead] * n_threads)
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
# Simple kernel: store tid to output[tid]
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_lshlrev_b32_e32(v[1], 2, v[0]).to_bytes() # offset = tid * 4
kernel += global_store_b32(addr=v[1], data=v[0], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({
(output_ptr, ctypes.sizeof(output)),
(args_ptr, ctypes.sizeof(args)),
(kernel_ptr, len(kernel)),
})
result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
self.assertEqual(result, 0)
# All threads should have written their tid
for i in range(n_threads):
self.assertEqual(output[i], i, f"Thread {i} didn't execute")
class TestRegressions(unittest.TestCase):
"""Regression tests for bugs fixed in the emulator."""
def test_v_fmac_f16(self):
"""V_FMAC_F16: fused multiply-add for FP16. Regression for VOP2 op 54."""
from extra.assembly.rdna3.emu import i16, f16
kernel = make_store_kernel([
v_mov_b32_e32(v[1], i16(2.0)), # v1.lo = 2.0 (fp16)
v_mov_b32_e32(v[2], i16(3.0)), # v2.lo = 3.0 (fp16)
# v1 = v1 * v2 + v1 = 2.0 * 3.0 + 2.0 = 8.0
VOP2(VOP2Op.V_FMAC_F16, v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertAlmostEqual(f16(out[0] & 0xffff), 8.0, places=2)
def test_v_cvt_f64_f32(self):
"""V_CVT_F64_F32: convert float32 to float64. Regression for VOP1 op 16."""
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_mov_b32_e32(v[1], i32(3.14159)).to_bytes()
kernel += VOP1(VOP1Op.V_CVT_F64_F32, v[4], v[1]).to_bytes() # v4:v5 = f64(v1)
kernel += v_lshlrev_b32_e32(v[3], 3, v[0]).to_bytes() # offset = tid * 8
kernel += global_store_b64(addr=v[3], data=v[4], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
output = (ctypes.c_double * 1)(0.0)
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
self.assertAlmostEqual(output[0], 3.14159, places=4)
def test_v_add_f64(self):
"""V_ADD_F64: add two float64 values. Regression for VOP3 op 807."""
from extra.assembly.rdna3.emu import i64_parts
kernel = b''
kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
# Load 1.5 into v1:v2
lo, hi = i64_parts(1.5)
kernel += v_mov_b32_e32(v[1], lo).to_bytes()
kernel += v_mov_b32_e32(v[2], hi).to_bytes()
# Load 2.5 into v3:v4
lo, hi = i64_parts(2.5)
kernel += v_mov_b32_e32(v[3], lo).to_bytes()
kernel += v_mov_b32_e32(v[4], hi).to_bytes()
# v5:v6 = v1:v2 + v3:v4 = 1.5 + 2.5 = 4.0
kernel += VOP3(VOP3Op.V_ADD_F64, v[5], v[1], v[3]).to_bytes()
kernel += v_lshlrev_b32_e32(v[7], 3, v[0]).to_bytes()
kernel += global_store_b64(addr=v[7], data=v[5], saddr=s[2]).to_bytes()
kernel += s_endpgm().to_bytes()
output = (ctypes.c_double * 1)(0.0)
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 1)(output_ptr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
self.assertAlmostEqual(output[0], 4.0, places=10)
def test_flat_load_d16_hi_b16(self):
"""FLAT_LOAD_D16_HI_B16: load 16-bit to high half. Regression for FLAT op 35."""
from extra.assembly.rdna3.emu import i16
# Create a buffer with test data
src_data = (ctypes.c_uint16 * 1)(0x1234)
src_ptr = ctypes.addressof(src_data)
output = (ctypes.c_uint32 * 1)(0xABCD0000) # preset low bits
output_ptr = ctypes.addressof(output)
args = (ctypes.c_uint64 * 2)(output_ptr, src_ptr)
args_ptr = ctypes.addressof(args)
kernel = b''
kernel += s_load_b128(s[0:3], s[0:1], soffset=NULL, offset=0).to_bytes()
kernel += s_waitcnt(lgkmcnt=0).to_bytes()
kernel += v_mov_b32_e32(v[1], 0xDEAD).to_bytes() # initial value with low bits set
kernel += v_mov_b32_e32(v[2], 0).to_bytes() # offset = 0
kernel += FLAT(FLATOp.FLAT_LOAD_D16_HI_B16, v[1], v[2], saddr=s[2], offset=0).to_bytes()
kernel += s_waitcnt(vmcnt=0).to_bytes()
kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[0]).to_bytes()
kernel += s_endpgm().to_bytes()
kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
kernel_ptr = ctypes.addressof(kernel_buf)
set_valid_mem_ranges({(output_ptr, 4), (src_ptr, 2), (args_ptr, 16), (kernel_ptr, len(kernel))})
run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
# High 16 bits should be 0x1234, low 16 bits preserved as 0xDEAD
self.assertEqual(output[0], 0x1234DEAD)
def test_v_mad_u16(self):
"""V_MAD_U16: multiply-add unsigned 16-bit. Regression for VOP3 op 577."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 10), # a = 10
v_mov_b32_e32(v[2], 20), # b = 20
v_mov_b32_e32(v[4], 5), # c = 5
VOP3(VOP3Op.V_MAD_U16, v[1], v[1], v[2], v[4]), # v1 = 10*20+5 = 205
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0] & 0xffff, 205)
def test_v_lshrrev_b16(self):
"""V_LSHRREV_B16: logical shift right 16-bit. Regression for VOP3 op 825."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 0x8000), # value to shift
v_mov_b32_e32(v[2], 4), # shift amount
VOP3(VOP3Op.V_LSHRREV_B16, v[1], v[2], v[1]), # v1 = 0x8000 >> 4 = 0x0800
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0] & 0xffff, 0x0800)
def test_v_min_u16(self):
"""V_MIN_U16: minimum of two unsigned 16-bit values. Regression for VOP3 op 779."""
kernel = make_store_kernel([
v_mov_b32_e32(v[1], 100),
v_mov_b32_e32(v[2], 50),
VOP3(VOP3Op.V_MIN_U16, v[1], v[1], v[2]),
])
out = run_kernel(kernel, n_threads=1)
self.assertEqual(out[0] & 0xffff, 50)
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply Accumulate) instructions."""
def test_wmma_f32_16x16x16_f16_identity(self):
"""V_WMMA_F32_16X16X16_F16 with identity matrix. Regression for VOP3P op 64."""
from extra.assembly.rdna3.emu import i16, f16, exec_wmma_f32_16x16x16_f16, WaveState
# Test using direct emulator call rather than full kernel to simplify
st = WaveState()
st.exec_mask = 0xffffffff # all 32 lanes active
# Set up A as identity matrix: A[i][i] = 1.0, rest = 0.0
# Lane i holds row i of A in 8 regs (2 fp16 per reg)
for lane in range(16):
for reg in range(8):
col0, col1 = reg * 2, reg * 2 + 1
val0 = i16(1.0) if col0 == lane else 0
val1 = i16(1.0) if col1 == lane else 0
st.vgpr[lane][0 + reg] = val0 | (val1 << 16) # src0 = v0:v7
# Set up B as identity matrix: lane i holds column i of B
for lane in range(16):
for reg in range(8):
row0, row1 = reg * 2, reg * 2 + 1
val0 = i16(1.0) if row0 == lane else 0
val1 = i16(1.0) if row1 == lane else 0
st.vgpr[lane][8 + reg] = val0 | (val1 << 16) # src1 = v8:v15
# Set up C as zeros
for lane in range(32):
for reg in range(8):
st.vgpr[lane][16 + reg] = 0 # src2 = v16:v23
# Create a fake VOP3P instruction
inst = VOP3P(VOP3POp.V_WMMA_F32_16X16X16_F16, v[24], src0=VGPR(0), src1=VGPR(8), src2=VGPR(16))
# Execute WMMA
exec_wmma_f32_16x16x16_f16(st, inst, 32)
# Check result: C should be identity (since A @ B where both are identity)
# Output i = row*16+col goes to lane (i%32), reg (i//32)
for row in range(16):
for col in range(16):
idx = row * 16 + col
lane, reg = idx % 32, idx // 32
result = st.vgpr[lane][24 + reg]
expected = 1.0 if row == col else 0.0
self.assertAlmostEqual(f32(result), expected, places=3,
msg=f"C[{row},{col}] = {f32(result)}, expected {expected}")
if __name__ == "__main__":
unittest.main()
+332
View File
@@ -0,0 +1,332 @@
#!/usr/bin/env python3
"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
import unittest
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import encode_src
class TestMUBUF(unittest.TestCase):
"""Test MUBUF (buffer) instructions."""
def test_buffer_load_b32_basic(self):
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]))
def test_buffer_load_b32_idxen(self):
# buffer_load_b32 v5, v0, s[8:11], s3 idxen offset:4095
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, idxen=1)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]))
def test_buffer_load_b32_offen(self):
# buffer_load_b32 v5, v0, s[8:11], s3 offen offset:4095
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, offen=1)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]))
def test_buffer_load_b32_glc(self):
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc
# GFX11: encoding: [0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]))
def test_buffer_load_b32_slc(self):
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 slc
# GFX11: encoding: [0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, slc=1)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]))
def test_buffer_load_b32_dlc(self):
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 dlc
# GFX11: encoding: [0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, dlc=1)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]))
def test_buffer_load_b32_all_flags(self):
# buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc slc dlc
# GFX11: encoding: [0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1, slc=1, dlc=1)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]))
def test_buffer_store_b32(self):
# buffer_store_b32 v1, off, s[12:15], s4 offset:4095
# GFX11: encoding: [0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]
inst = buffer_store_b32(vdata=v[1], vaddr=v[0], srsrc=s[12:16], soffset=s[4], offset=4095)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]))
def test_buffer_load_b64(self):
# buffer_load_b64 v[5:6], off, s[8:11], s3 offset:4095
# GFX11: encoding: [0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]
inst = buffer_load_b64(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]))
def test_buffer_load_soffset_m0(self):
# buffer_load_b32 v5, off, s[8:11], m0 offset:4095
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=M0, offset=4095)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]))
def test_buffer_load_soffset_inline_const(self):
# buffer_load_b32 v5, off, s[8:11], 0 offset:4095
# GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=0, offset=4095)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]))
def test_buffer_disasm_roundtrip(self):
inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
decoded = MUBUF.from_bytes(inst.to_bytes())
self.assertEqual(decoded.to_bytes(), inst.to_bytes())
class TestMTBUF(unittest.TestCase):
"""Test MTBUF (typed buffer) instructions."""
def test_tbuffer_load_format_x(self):
# tbuffer_load_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
# BUF_FMT_32_FLOAT = 22
# GFX11: encoding: [0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]
inst = tbuffer_load_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]))
def test_tbuffer_store_format_x(self):
# tbuffer_store_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
# BUF_FMT_32_FLOAT = 22
# GFX11: encoding: [0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]
inst = tbuffer_store_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]))
def test_tbuffer_load_format_xy(self):
# tbuffer_load_format_xy v[5:6], off, s[8:11], s3 format:[BUF_FMT_32_32_FLOAT] offset:4095
# BUF_FMT_32_32_FLOAT = 50
# GFX11: encoding: [0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]
inst = tbuffer_load_format_xy(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=50)
self.assertEqual(inst.to_bytes(), bytes([0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]))
class TestMIMG(unittest.TestCase):
"""Test MIMG (image) instructions."""
def test_image_load_2d(self):
# image_load v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
# GFX11: encoding: [0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
inst = image_load(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1) # dim=1 is SQ_RSRC_IMG_2D
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
def test_image_store_2d(self):
# image_store v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
# GFX11: encoding: [0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]
inst = image_store(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]))
def test_image_load_1d(self):
# image_load v[0:3], v4, s[0:7] dmask:0xf dim:SQ_RSRC_IMG_1D
# GFX11: encoding: [0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
inst = image_load(vdata=v[0:4], vaddr=v[4], srsrc=s[0:8], dmask=0xf, dim=0) # dim=0 is SQ_RSRC_IMG_1D
self.assertEqual(inst.to_bytes(), bytes([0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
def test_image_sample(self):
# image_sample v[0:3], v[4:5], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
# GFX11: encoding: [0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]
inst = image_sample(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], ssamp=s[8:12], dmask=0xf, dim=1)
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]))
def test_image_load_d16(self):
# image_load v[0:1], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D d16
# GFX11: encoding: [0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]
inst = image_load(vdata=v[0:2], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1, d16=1)
self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]))
class TestEXP(unittest.TestCase):
"""Test EXP (export) instructions."""
def test_exp_mrt0(self):
# exp mrt0 v0, v1, v2, v3
# GFX11: encoding: [0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]
inst = EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]))
def test_exp_mrtz(self):
# exp mrtz v4, v3, v2, v1
# GFX11: encoding: [0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]
inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]))
def test_exp_mrtz_done(self):
# exp mrtz v4, v3, v2, v1 done
# GFX11: encoding: [0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x01]
inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[3], done=1)
self.assertEqual(inst.to_bytes(), bytes([0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x03]))
def test_exp_partial_mask(self):
# exp mrt0 v0, v1, off, off (en=0x3, only first two components)
# GFX11: encoding: [0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]
inst = EXP(en=0x3, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[0], vsrc3=v[0])
self.assertEqual(inst.to_bytes(), bytes([0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]))
def test_exp_row_en(self):
# exp mrtz v4, v3, v2, v1 row_en
# GFX11: encoding: [0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]
inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1], row=1)
self.assertEqual(inst.to_bytes(), bytes([0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]))
class TestDS(unittest.TestCase):
"""Test DS (data share / LDS) instructions."""
def test_ds_store_b32(self):
# ds_store_b32 v0, v1
# GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_load_b32(self):
# ds_load_b32 v0, v1
# GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
inst = ds_load_b32(vdst=v[0], addr=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
def test_ds_store_b32_offset(self):
# ds_store_b32 v0, v1 offset:64
# GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_load_b64(self):
# ds_load_b64 v[0:1], v2
# GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
inst = ds_load_b64(vdst=v[0:2], addr=v[2])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
def test_ds_add_u32(self):
# ds_add_u32 v0, v1
# GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
inst = ds_add_u32(addr=v[0], data0=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_store_b32_gds(self):
# ds_store_b32 v0, v1 gds
# GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
class TestVOP3(unittest.TestCase):
"""Test VOP3 (3-operand vector) instructions."""
def test_v_fma_f32(self):
# v_fma_f32 v0, v1, v2, v3
# GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
def test_v_mad_f32(self):
# v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
# Use v_fma_f32 with vdst == src2
inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
def test_v_add3_u32(self):
# v_add3_u32 v0, v1, v2, v3
# GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
class TestFLAT(unittest.TestCase):
"""Test FLAT/GLOBAL/SCRATCH memory instructions."""
def test_global_load_b32(self):
# global_load_b32 v0, v[1:2], off (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
def test_global_store_b32(self):
# global_store_b32 v[0:1], v2, off (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
inst = global_store_b32(addr=v[0:2], data=v[2], saddr=OFF)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
def test_global_load_b32_saddr(self):
# global_load_b32 v0, v1, s[0:1] (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:2])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
def test_global_load_b32_offset(self):
# global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
# GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF, offset=256)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
def test_global_load_b64(self):
# global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
inst = global_load_b64(vdst=v[0:2], addr=v[2:4], saddr=OFF)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
class TestSMEM(unittest.TestCase):
"""Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
def test_smem_dlc_bit_position(self):
# s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
# GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], dlc=1)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_glc_bit_position(self):
# s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
# GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_glc_dlc_combined(self):
# s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
# GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1, dlc=1)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_disasm_roundtrip_dlc(self):
# Test that disassembly/reassembly preserves DLC bit correctly
data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
decoded = SMEM.from_bytes(data)
self.assertEqual(decoded.to_bytes(), data)
def test_smem_disasm_roundtrip_glc_dlc(self):
# Test that disassembly/reassembly preserves GLC+DLC bits correctly
data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
decoded = SMEM.from_bytes(data)
self.assertEqual(decoded.to_bytes(), data)
class TestVOP3Literal(unittest.TestCase):
"""Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
def test_vop3_with_literal(self):
# v_add3_u32 v5, vcc_hi, 0xaf123456, v255
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
from extra.assembly.rdna3.lib import RawImm
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
self.assertEqual(inst.to_bytes(), expected)
def test_vop3_literal_null_operand(self):
# v_add3_u32 v5, null, exec_lo, 0xaf123456
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
from extra.assembly.rdna3.lib import RawImm
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
self.assertEqual(inst.to_bytes(), expected)
def test_vop3p_with_literal(self):
# Test VOP3P literal encoding (also uses Inst64)
from extra.assembly.rdna3.lib import RawImm
inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
self.assertEqual(len(inst.to_bytes()), 12) # 8 bytes + 4 byte literal
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,178 @@
# do not change these tests. we need to fix bugs to make them pass
# the Inst constructor should be looking at the types of the fields to correctly set the value
import unittest, struct
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import Inst
from extra.assembly.rdna3.asm import asm
from extra.assembly.rdna3.test.test_roundtrip import compile_asm
class TestIntegration(unittest.TestCase):
inst: Inst
def tearDown(self):
if not hasattr(self, 'inst'): return
b = self.inst.to_bytes()
st = self.inst.disasm()
reasm = asm(st)
desc = f"{st:25s} {self.inst} {b!r} {reasm}"
self.assertEqual(b, compile_asm(st), desc)
# TODO: this compare should work for valid things
#self.assertEqual(self.inst, reasm)
self.assertEqual(repr(self.inst), repr(reasm))
print(desc)
def test_load_b128(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
def test_load_b128_wrong_size(self):
# this should have to be 4 regs on the loaded to
with self.assertRaises(Exception):
self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
def test_mov_b32(self):
self.inst = s_mov_b32(s[80], s[0])
def test_mov_b64(self):
self.inst = s_mov_b64(s[80:81], s[0:1])
def test_mov_b32_wrong(self):
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80:81], s[0:1])
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80:81], s[0])
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80], s[0:1])
def test_mov_b64_wrong(self):
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80], s[0])
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80], s[0:1])
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80:81], s[0])
def test_load_b128_no_0(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL)
def test_load_b128_s(self):
self.inst = s_load_b128(s[4:7], s[0:1], s[8], 0)
def test_load_b128_v(self):
with self.assertRaises(TypeError):
self.inst = s_load_b128(s[4:7], s[0:1], v[8], 0)
def test_load_b128_off(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 3)
def test_simple_stos(self):
self.inst = s_mov_b32(s[0], s[1])
def test_simple_wrong(self):
with self.assertRaises(TypeError):
self.inst = s_mov_b32(v[0], s[1])
def test_simple_vtov(self):
self.inst = v_mov_b32_e32(v[0], v[1])
def test_simple_stov(self):
self.inst = v_mov_b32_e32(v[0], s[2])
def test_simple_float_to_v(self):
self.inst = v_mov_b32_e32(v[0], 1.0)
def test_simple_v_to_float(self):
with self.assertRaises(TypeError):
self.inst = v_mov_b32_e32(1, v[0])
def test_simple_int_to_v(self):
self.inst = v_mov_b32_e32(v[0], 1)
def test_three_add(self):
self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3])
def test_three_add_v(self):
self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3])
def test_three_add_const(self):
self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3])
def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07)
def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7)
def test_vmad(self):
self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2])
def test_large_imm(self):
self.inst = v_mov_b32_e32(v[0], 0x1234)
def test_dual_mov(self):
self.inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], srcy0=v[4])
def test_dual_mul(self):
self.inst = v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
def test_simple_int_to_s(self):
self.inst = s_mov_b32(s[0], 3)
def test_complex_int_to_s(self):
self.inst = s_mov_b32(s[0], 0x235646)
def test_simple_float_to_s(self):
self.inst = s_mov_b32(s[0], 1.0)
def test_complex_float_to_s(self):
self.inst = s_mov_b32(s[0], 1337.0)
int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
self.assertEqual(self.inst, int_inst)
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,343 @@
#!/usr/bin/env python3
"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
import unittest, re, io, sys
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.asm import waitcnt, asm
def get_amd_toolchain():
"""Check if AMD toolchain is available."""
try:
from tinygrad.runtime.support.compiler_amd import HIPCompiler
HIPCompiler("gfx1100").compile(".text\ns_endpgm")
return True
except Exception:
return False
def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
"""Disassemble ELF binary using tinygrad's compiler, return raw output."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
old_stdout = sys.stdout
sys.stdout = io.StringIO()
HIPCompiler(arch).disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
return output
def parse_disassembly(raw: str) -> list[str]:
"""Parse disassembly output to list of instruction mnemonics."""
lines = []
for line in raw.splitlines():
if line.startswith('\t'):
instr = line.split('//')[0].strip()
if instr: lines.append(instr)
return lines
def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
"""Assemble instructions with our DSL, then disassemble with AMD toolchain."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
# Generate bytes from our DSL
code_bytes = b''.join(inst.to_bytes() for inst in instructions)
# Wrap in minimal ELF-compatible assembly with .byte directives
byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
# Assemble with AMD COMGR and disassemble
lib = HIPCompiler(arch).compile(asm_src)
return parse_disassembly(disassemble(lib, arch))
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
class TestIntegration(unittest.TestCase):
"""Test our assembler output matches LLVM disassembly."""
def test_simple_sop1(self):
"""Test SOP1 instructions round-trip."""
instructions = [
s_mov_b32(s[0], s[1]),
s_mov_b32(s[2], 0),
s_not_b32(s[3], s[4]),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_mov_b32', disasm[0])
self.assertIn('s_mov_b32', disasm[1])
self.assertIn('s_not_b32', disasm[2])
def test_simple_sop2(self):
"""Test SOP2 instructions round-trip."""
instructions = [
s_add_u32(s[0], s[1], s[2]),
s_sub_u32(s[3], s[4], 10),
s_and_b32(s[5], s[6], s[7]),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_add_u32', disasm[0])
self.assertIn('s_sub_u32', disasm[1])
self.assertIn('s_and_b32', disasm[2])
def test_simple_vop2(self):
"""Test VOP2 instructions round-trip."""
instructions = [
v_add_f32_e32(v[0], v[1], v[2]),
v_mul_f32_e32(v[3], 1.0, v[4]), # 1.0 is inline constant
v_and_b32_e32(v[5], 10, v[6]), # small inline constant
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('v_add_f32', disasm[0])
self.assertIn('v_mul_f32', disasm[1])
def test_control_flow(self):
"""Test control flow instructions."""
instructions = [
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_waitcnt', disasm[0])
self.assertIn('s_endpgm', disasm[1])
def test_memory_ops(self):
"""Test memory instructions."""
instructions = [
s_load_b32(s[0], s[0:2], NULL),
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_load_b32', disasm[0])
self.assertIn('s_waitcnt', disasm[1])
self.assertIn('global_store_b32', disasm[2])
def test_full_kernel(self):
"""Test a complete kernel similar to tinygrad output."""
# Simple kernel: load value, add 1, store back
instructions = [
# Get thread ID
v_mov_b32_e32(v[0], s[0]), # base addr low
v_mov_b32_e32(v[1], s[1]), # base addr high
# Load value
global_load_b32(vdst=v[2], addr=v[0:2], saddr=OFF),
s_waitcnt(simm16=waitcnt(vmcnt=0)),
# Add 1.0
v_add_f32_e32(v[2], 1.0, v[2]),
# Store result
global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
s_endpgm(),
]
disasm = assemble_and_disassemble(instructions)
# Verify key instructions are present
self.assertTrue(any('global_load' in d for d in disasm))
self.assertTrue(any('v_add_f32' in d for d in disasm))
self.assertTrue(any('global_store' in d for d in disasm))
self.assertTrue(any('s_endpgm' in d for d in disasm))
def test_bytes_roundtrip(self):
"""Test that our bytes match what AMD assembler produces."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
# Simple instruction
inst = s_mov_b32(s[0], s[1])
our_bytes = inst.to_bytes()
# Assemble same instruction with AMD toolchain
asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
compiler = HIPCompiler("gfx1100")
lib = compiler.compile(asm_src)
raw = disassemble(lib)
for line in raw.splitlines():
if 's_mov_b32' in line and '//' in line:
# Extract hex bytes from comment: "// 000000001300: BE800001"
comment = line.split('//')[1].strip()
hex_str = comment.split(':')[1].strip()
# Convert big-endian hex string to little-endian bytes
amd_bytes = bytes.fromhex(hex_str)[::-1] # reverse for little-endian
self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
return
self.fail("Could not find s_mov_b32 in disassembly")
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
class TestAsm(unittest.TestCase):
"""Test asm() string parsing."""
def test_asm_basic(self):
"""Test basic instruction parsing."""
inst = asm('s_mov_b32 s0, s1')
self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], s[1]).to_bytes())
def test_asm_with_immediates(self):
"""Test parsing with immediate values."""
inst = asm('s_add_u32 s0, s1, 10')
self.assertEqual(inst.to_bytes(), s_add_u32(s[0], s[1], 10).to_bytes())
def test_asm_float_const(self):
"""Test parsing float constants."""
inst = asm('v_mul_f32_e32 v0, 1.0, v1')
self.assertEqual(inst.to_bytes(), v_mul_f32_e32(v[0], 1.0, v[1]).to_bytes())
def test_asm_hex_immediate(self):
"""Test parsing hex immediates."""
inst = asm('s_waitcnt 0xfc07')
self.assertEqual(inst.to_bytes(), s_waitcnt(simm16=0xfc07).to_bytes())
def test_asm_special_regs(self):
"""Test parsing special registers."""
inst = asm('s_mov_b32 s0, vcc_lo')
self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], VCC_LO).to_bytes())
def test_asm_register_range(self):
"""Test parsing register ranges."""
inst = asm('s_load_b128 s[4:7], s[0:1], null')
self.assertEqual(inst.to_bytes(), s_load_b128(s[4:7], s[0:1], NULL).to_bytes())
def test_asm_matches_llvm(self):
"""Test asm() output matches LLVM assembler."""
from tinygrad.runtime.support.compiler_amd import HIPCompiler
compiler = HIPCompiler('gfx1100')
def get_llvm_bytes(instr: str) -> bytes:
src = f'.text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n{instr}\n'
lib = compiler.compile(src)
raw = disassemble(lib)
for line in raw.splitlines():
if instr.split()[0] in line and '//' in line:
hex_str = line.split('//')[1].strip().split(':')[1].strip()
return bytes.fromhex(hex_str)[::-1]
return b''
tests = ['s_mov_b32 s0, s1', 's_endpgm', 'v_add_f32_e32 v0, v1, v2']
for t in tests:
self.assertEqual(asm(t).to_bytes(), get_llvm_bytes(t), f"mismatch for: {t}")
def test_asm_vop3_modifiers(self):
"""Test asm() with VOP3 modifiers (neg, abs, clamp)."""
import subprocess, re
def get_llvm_encoding(instr: str) -> str:
result = subprocess.run(['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
input=instr, capture_output=True, text=True)
if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
return m.group(1).replace('0x','').replace(',','').replace(' ','')
return ''
tests = [
'v_fma_f32 v0, -v1, v2, v3', # neg on src0
'v_fma_f32 v0, v1, |v2|, v3', # abs on src1
'v_fma_f32 v0, v1, v2, v3 clamp', # clamp
'v_fma_f32 v0, -v1, |v2|, v3 clamp', # all modifiers
'v_fma_f32 v0, -|v1|, v2, v3', # neg+abs on same operand
]
for t in tests:
our_hex = asm(t).to_bytes().hex()
llvm_hex = get_llvm_encoding(t)
self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
class TestTinygradIntegration(unittest.TestCase):
"""Test that we can parse disassembled tinygrad kernels."""
def test_simple_add_kernel(self):
"""Generate a simple add kernel from tinygrad and verify disassembly."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Create a computation that generates a real kernel
a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
c = a + b
# Get schedule and find SINK
schedule = c.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
# Generate program
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
self.assertIsNotNone(prg.src)
# Compile and disassemble
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
instrs = parse_disassembly(raw_disasm)
# Verify we got some instructions
self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
# Should have an endpgm
self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
def test_matmul_kernel(self):
"""Generate a matmul kernel and verify disassembly has expected patterns."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Create a small matmul
a = Tensor.rand(4, 4).realize()
b = Tensor.rand(4, 4).realize()
c = a @ b
# Get schedule
schedule = c.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
self.assertTrue(len(sink_items) > 0)
# Generate and compile
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
instrs = parse_disassembly(raw_disasm)
# Matmul should have multiply and add instructions
has_mul = any('mul' in i.lower() for i in instrs)
has_add = any('add' in i.lower() for i in instrs)
self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
def test_disasm_to_bytes_roundtrip(self):
"""Parse disassembled instructions and verify we can re-encode some of them."""
from tinygrad import Tensor
from tinygrad.codegen import get_program
from tinygrad.renderer.cstyle import AMDHIPRenderer
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.uop.ops import Ops
# Simple kernel
a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
b = (a * 2.0)
schedule = b.schedule()
sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
if not sink_items: return # skip if no kernel
renderer = AMDHIPRenderer('gfx1100')
prg = get_program(sink_items[0].ast, renderer)
compiler = HIPCompiler('gfx1100')
lib = compiler.compile(prg.src)
raw_disasm = disassemble(lib)
# Find s_endpgm and verify we can encode it
for line in raw_disasm.splitlines():
if 's_endpgm' in line and '//' in line:
# Extract bytes from comment
comment = line.split('//')[1].strip()
hex_str = comment.split(':')[1].strip()
amd_bytes = bytes.fromhex(hex_str)[::-1]
# Our encoding
our_inst = s_endpgm()
our_bytes = our_inst.to_bytes()
self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
return
if __name__ == "__main__":
unittest.main()
+173
View File
@@ -0,0 +1,173 @@
#!/usr/bin/env python3
"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
import unittest, re
from tinygrad.helpers import fetch
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.asm import asm
from extra.assembly.rdna3.test.test_roundtrip import compile_asm, disassemble_lib
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
# Format info: (filename, format_class, op_enum)
LLVM_TEST_FILES = {
# Scalar ALU
'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
# Vector ALU
'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp), # VOP3SD shares file with VOP3
'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp), # VOPCX uses VOPC format
# VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
# Memory
'ds': ('gfx11_asm_ds.s', DS, DSOp),
'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
# WMMA (matrix multiply)
'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
# Additional features
'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
# Alias files (alternative mnemonics)
'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
}
def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
"""Parse LLVM test format into (asm, expected_bytes) pairs."""
tests, lines = [], text.split('\n')
for i, line in enumerate(lines):
line = line.strip()
if not line or line.startswith(('//', '.', ';')): continue
asm_text = line.split('//')[0].strip()
if not asm_text: continue
for j in range(i, min(i + 3, len(lines))):
# Match GFX11, W32, or W64 encodings (all valid for gfx11)
if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
if hex_bytes:
try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
except ValueError: pass
break
return tests
def try_assemble(text: str):
"""Try to assemble instruction text, return bytes or None on failure."""
try: return asm(text).to_bytes()
except: return None
class TestLLVM(unittest.TestCase):
"""Test assembler and disassembler against all LLVM test vectors."""
tests: dict[str, list[tuple[str, bytes]]] = {}
@classmethod
def setUpClass(cls):
for name, (filename, _, _) in LLVM_TEST_FILES.items():
try:
data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
except Exception as e:
print(f"Warning: couldn't fetch {filename}: {e}")
cls.tests[name] = []
# Generate test methods dynamically for each format
def _make_asm_test(name):
def test(self):
passed, failed, skipped = 0, 0, 0
for asm_text, expected in self.tests.get(name, []):
result = try_assemble(asm_text)
if result is None: skipped += 1
elif result == expected: passed += 1
else: failed += 1
print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
self.assertEqual(failed, 0)
return test
def _make_disasm_test(name):
def test(self):
from tinygrad.runtime.support.compiler_amd import HIPCompiler
compiler = HIPCompiler('gfx1100')
_, fmt_cls, op_enum = LLVM_TEST_FILES[name]
passed, failed, skipped, failures = 0, 0, 0, []
# VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
# Note: opcodes 0-255 are VOPC promoted to VOP3, never VOP3SD
vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
# vop3_from_vopc/vopcx tests have VOPC opcodes 0-255, not VOP3SD - don't detect as VOP3SD
is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
# Undocumented opcodes not in AMD ISA PDF - skip these
undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}} # s_atc_probe*, s_subvector_loop*, s_waitcnt_depctr, unknown
for asm_text, data in self.tests.get(name, []):
if len(data) > fmt_cls._size(): continue # skip literals (need different handling)
# Skip undocumented opcodes
temp_inst = fmt_cls.from_bytes(data)
temp_op = temp_inst._values.get('op', 0)
temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
if temp_op in undocumented.get(name, set()): skipped += 1; continue
# Skip SOPP no-imm instructions with non-zero simm16 (can't roundtrip through LLVM)
if name == 'sopp':
simm16 = temp_inst._values.get('simm16', 0)
simm16 = simm16.val if hasattr(simm16, 'val') else simm16
sopp_no_imm = {48, 54, 53, 55, 60, 61, 62} # s_endpgm, s_barrier, s_wakeup, s_icache_inv, s_wait_idle, s_endpgm_saved, s_code_end
if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
try:
# VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use
if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
temp = VOP3.from_bytes(data)
op_val = temp._values.get('op', 0)
op_val = op_val.val if hasattr(op_val, 'val') else op_val
is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
# Validate opcode with appropriate enum
if is_vop3sd:
VOP3SDOp(op_val)
else:
VOP3Op(op_val)
else:
decoded = fmt_cls.from_bytes(data)
op_val = decoded._values.get('op', 0)
op_val = op_val.val if hasattr(op_val, 'val') else op_val
op_enum(op_val) # validate opcode
if decoded.to_bytes()[:len(data)] != data:
failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue
disasm_str = decoded.disasm()
# Test: LLVM should assemble our disasm output to the same bytes
llvm_bytes = compile_asm(disasm_str, compiler)
if llvm_bytes is None:
failed += 1; failures.append(f"LLVM failed to assemble: '{disasm_str}' (from '{asm_text}')")
elif llvm_bytes == data: passed += 1
else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
except Exception as e:
failed += 1; failures.append(f"exception for {data.hex()}: {e}")
print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
if failures[:10]: print(" " + "\n ".join(failures[:10]))
self.assertEqual(failed, 0)
return test
for name in LLVM_TEST_FILES:
setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,54 @@
#!/usr/bin/env python3
"""Test that invalid instructions raise exceptions through the mock GPU stack."""
import unittest, subprocess, os, time
class TestMockGPUInvalidInstruction(unittest.TestCase):
def test_unsupported_instruction_raises(self):
"""Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
test_code = '''
import struct
from tinygrad import Device, Tensor
from tinygrad.engine.realize import get_runner
from tinygrad.runtime.ops_amd import AMDProgram
dev = Device["AMD"]
a = Tensor([1.0]).realize()
b = a + 1
si = b.schedule()[-1]
runner = get_runner(dev.device, si.ast)
prg = runner._prg
lib = bytearray(prg.lib)
# Find s_endpgm (0xBFB00000) and replace with invalid SOPP op=127 (0xBFFF0000)
found = False
for i in range(0, len(lib) - 4, 4):
if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
lib[i:i+4] = struct.pack("<I", 0xBFFF0000)
found = True
break
assert found, "s_endpgm not found"
patched_prg = AMDProgram(dev, "patched", bytes(lib))
b.uop.buffer.allocate()
patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
dev.synchronize()
'''
env = os.environ.copy()
env["AMD"] = "1"
env["MOCKGPU"] = "1"
env["PYTHON_REMU"] = "1"
env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
st = time.perf_counter()
result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
elapsed = time.perf_counter() - st
self.assertNotEqual(result.returncode, 0, "should have raised")
self.assertIn("NotImplementedError", result.stderr)
# Should exit immediately, not wait for the full timeout
self.assertLess(elapsed, 5.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,158 @@
#!/usr/bin/env python3
"""Test that PDF parser correctly extracts format fields."""
import unittest
from extra.assembly.rdna3.autogen import (
SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
SOP1Op, SOP2Op, VOP1Op, VOP3Op
)
# expected formats with key fields and whether they have ENCODING
EXPECTED_FORMATS = {
'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
'LDSDIR': (['VDST', 'OP'], True),
'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
'SOP1': (['OP', 'SDST', 'SSRC0'], True),
'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
'SOPK': (['OP', 'SDST', 'SIMM16'], True),
'SOPP': (['OP', 'SIMM16'], True),
'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOP1': (['OP', 'VDST', 'SRC0'], True),
'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
}
class TestPDFParserGenerate(unittest.TestCase):
"""Test the PDF parser by running generate() and checking results."""
result: dict
@classmethod
def setUpClass(cls):
from extra.assembly.rdna3.gen import generate
cls.result = generate()
def test_all_formats_present(self):
"""All expected formats should be parsed."""
for fmt_name in EXPECTED_FORMATS:
self.assertIn(fmt_name, self.result["formats"], f"missing format {fmt_name}")
def test_format_count(self):
"""Should have exactly 23 formats."""
self.assertEqual(len(self.result["formats"]), 23)
def test_no_duplicate_fields(self):
"""No format should have duplicate field names."""
for fmt_name, fields in self.result["formats"].items():
field_names = [f[0] for f in fields]
self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
def test_expected_fields(self):
"""Each format should have its expected key fields."""
for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
fields = {f[0] for f in self.result["formats"].get(fmt_name, [])}
for field in expected_fields:
self.assertIn(field, fields, f"{fmt_name} missing {field}")
if has_encoding:
self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
else:
self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
def test_vopd_no_dpp16_fields(self):
"""VOPD should not have DPP16-specific fields (parser boundary bug)."""
vopd_fields = {f[0] for f in self.result["formats"].get("VOPD", [])}
for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
def test_dpp16_no_vinterp_fields(self):
"""DPP16 should not have VINTERP-specific fields."""
dpp16_fields = {f[0] for f in self.result["formats"].get("DPP16", [])}
for field in ['VDST', 'WAITEXP']:
self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
def test_sopp_no_smem_fields(self):
"""SOPP should not have SMEM fields (page break bug)."""
sopp_fields = {f[0] for f in self.result["formats"].get("SOPP", [])}
for field in ['SBASE', 'SDATA']:
self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
class TestPDFParser(unittest.TestCase):
"""Verify format classes have correct fields from PDF parsing."""
def test_sop2_fields(self):
"""SOP2 should have op, sdst, ssrc0, ssrc1."""
for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
self.assertIn(field, SOP2._fields)
self.assertEqual(SOP2._fields['op'].hi, 29)
self.assertEqual(SOP2._fields['op'].lo, 23)
def test_sop1_fields(self):
"""SOP1 should have op, sdst, ssrc0 with correct bit positions."""
for field in ['op', 'sdst', 'ssrc0']:
self.assertIn(field, SOP1._fields)
self.assertNotIn('simm16', SOP1._fields)
self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
assert SOP1._encoding is not None
self.assertEqual(SOP1._encoding[0].hi, 31)
self.assertEqual(SOP1._encoding[1], 0b101111101)
def test_vop3sd_fields(self):
"""VOP3SD should have all fields including src0/src1/src2 from page continuation."""
for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
self.assertIn(field, VOP3SD._fields)
self.assertEqual(VOP3SD._fields['src0'].hi, 40)
self.assertEqual(VOP3SD._fields['src0'].lo, 32)
self.assertEqual(VOP3SD._size(), 8)
def test_flat_has_vdst(self):
"""FLAT should have vdst field."""
self.assertIn('vdst', FLAT._fields)
self.assertEqual(FLAT._fields['vdst'].hi, 63)
self.assertEqual(FLAT._fields['vdst'].lo, 56)
def test_encoding_bits(self):
"""Verify encoding bits are correct for major formats."""
tests = [
(SOP2, 31, 30, 0b10),
(SOPK, 31, 28, 0b1011),
(SOPP, 31, 23, 0b101111111),
(VOP1, 31, 25, 0b0111111),
(VOP2, 31, 31, 0b0),
(VOPC, 31, 25, 0b0111110),
(FLAT, 31, 26, 0b110111),
]
for cls, hi, lo, val in tests:
assert cls._encoding is not None
self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
def test_opcode_enums_exist(self):
"""Verify opcode enums are generated with expected counts."""
self.assertGreater(len(SOP1Op), 50)
self.assertGreater(len(SOP2Op), 50)
self.assertGreater(len(VOP1Op), 50)
self.assertGreater(len(VOP3Op), 200)
def test_vopd_no_duplicate_fields(self):
"""VOPD should not have duplicate fields and should not include DPP16 fields."""
field_names = list(VOPD._fields.keys())
self.assertEqual(len(field_names), len(set(field_names)))
for field in ['srcx0', 'srcy0', 'opx', 'opy']:
self.assertIn(field, VOPD._fields)
for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
self.assertNotIn(field, VOPD._fields)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,94 @@
#!/usr/bin/env python3
import unittest, subprocess
from extra.assembly.rdna3.autogen import *
def llvm_assemble(asm: str) -> bytes:
"""Assemble using llvm-mc and return bytes."""
result = subprocess.run(
["llvm-mc", "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
input=asm, capture_output=True, text=True
)
out = b''
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
enc = enc.strip('[]').replace('0x', '').replace(',', '')
out += bytes.fromhex(enc)
if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
return out
class TestRDNA3Asm(unittest.TestCase):
def test_full_program(self):
"""Test the full program from rdna3fun.py matches llvm-mc output."""
program = [
v_bfe_u32(v[1], v[0], 10, 10),
s_load_b128(s[4:7], s[0:1], NULL),
v_and_b32_e32(v[0], 0x3FF, v[0]),
s_mulk_i32(s[3], 0x87),
v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
v_mul_u32_u24_e32(v[0], 45, v[0]),
v_ashrrev_i32_e32(v[2], 31, v[1]),
v_add3_u32(v[0], v[0], s[3], v[1]),
v_lshlrev_b64(v[2:3], 2, v[1:2]),
v_ashrrev_i32_e32(v[1], 31, v[0]),
v_lshlrev_b64(v[0:1], 2, v[0:1]),
s_waitcnt(0xfc07), # lgkmcnt(0)
v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
v_add_co_ci_u32_e32(v[3], s[7], v[3]),
v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
global_load_b32(vdst=v[2], addr=v[2], saddr=OFF),
v_add_co_ci_u32_e32(v[1], s[5], v[1]),
s_waitcnt(0x03f7), # vmcnt(0)
global_store_b32(addr=v[0], data=v[2], saddr=OFF),
s_endpgm(),
]
asm = """
v_bfe_u32 v1, v0, 10, 10
s_load_b128 s[4:7], s[0:1], null
v_and_b32_e32 v0, 0x3FF, v0
s_mulk_i32 s3, 0x87
v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
v_mul_u32_u24_e32 v0, 45, v0
v_ashrrev_i32_e32 v2, 31, v1
v_add3_u32 v0, v0, s3, v1
v_lshlrev_b64 v[2:3], 2, v[1:2]
v_ashrrev_i32_e32 v1, 31, v0
v_lshlrev_b64 v[0:1], 2, v[0:1]
s_waitcnt lgkmcnt(0)
v_add_co_u32 v2, vcc_lo, s6, v2
v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
v_add_co_u32 v0, vcc_lo, s4, v0
global_load_b32 v2, v[2:3], off
v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
s_waitcnt vmcnt(0)
global_store_b32 v[0:1], v2, off
s_endpgm
"""
expected = llvm_assemble(asm)
for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
actual = b''.join(inst.to_bytes() for inst in program)
self.assertEqual(actual, expected)
def test_sop2_s_add_u32(self):
inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
expected = llvm_assemble("s_add_u32 s3, s0, s1")
self.assertEqual(inst.to_bytes(), expected)
def test_vop2_v_and_b32_inline_const(self):
inst = v_and_b32_e32(v[0], 10, v[0])
expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
self.assertEqual(inst.to_bytes(), expected)
def test_sopp_s_endpgm(self):
inst = s_endpgm()
expected = llvm_assemble("s_endpgm")
self.assertEqual(inst.to_bytes(), expected)
def test_sop1_s_mov_b32(self):
inst = s_mov_b32(s[0], s[1])
expected = llvm_assemble("s_mov_b32 s0, s1")
self.assertEqual(inst.to_bytes(), expected)
if __name__ == "__main__":
unittest.main()
+234
View File
@@ -0,0 +1,234 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.lib import Inst
from extra.assembly.rdna3.asm import asm
# Instruction format detection based on encoding bits
def detect_format(data: bytes) -> type[Inst] | None:
"""Detect instruction format from machine code bytes."""
if len(data) < 4: return None
word = int.from_bytes(data[:4], 'little')
enc_9bit = (word >> 23) & 0x1FF # 9-bit encoding for SOP1/SOPC/SOPP
enc_8bit = (word >> 24) & 0xFF
# Check 9-bit encodings first (most specific)
if enc_9bit == 0x17D: return SOP1 # bits 31:23 = 101111101
if enc_9bit == 0x17E: return SOPC # bits 31:23 = 101111110
if enc_9bit == 0x17F: return SOPP # bits 31:23 = 101111111
# SOPK: bits 31:28 = 1011, bits 27:23 = opcode (check after SOP1/SOPC/SOPP)
if enc_8bit in range(0xB0, 0xC0): return SOPK
# SOP2: bits 31:23 in range 0x100-0x17C (0x80-0xBE in bits 31:24, but not SOPK)
if 0x80 <= enc_8bit <= 0x9F: return SOP2
# VOP1: bits 31:25 = 0111111 (0x3F)
if (word >> 25) == 0x3F: return VOP1
# VOPC: bits 31:25 = 0111110 (0x3E)
if (word >> 25) == 0x3E: return VOPC
# VOP2: bits 31:30 = 00
if (word >> 30) == 0: return VOP2
# Check 64-bit formats
if len(data) >= 8:
if enc_8bit in (0xD4, 0xD5, 0xD7): return VOP3
if enc_8bit == 0xD6: return VOP3SD
if enc_8bit == 0xCC: return VOP3P
if enc_8bit == 0xCD: return VINTERP
if enc_8bit in (0xC8, 0xC9): return VOPD
if enc_8bit == 0xF4: return SMEM
if enc_8bit == 0xD8: return DS
if enc_8bit in (0xDC, 0xDD, 0xDE, 0xDF): return FLAT
if enc_8bit in (0xE0, 0xE1, 0xE2, 0xE3): return MUBUF
if enc_8bit in (0xE8, 0xE9, 0xEA, 0xEB): return MTBUF
return None
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
old_stdout = sys.stdout
sys.stdout = io.StringIO()
compiler.disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
results = []
for line in output.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if not instr: continue
comment = line.split('//')[1].strip()
if ':' not in comment: continue
hex_str = comment.split(':')[1].strip().split()[0]
try:
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
results.append((instr, machine_bytes))
except ValueError:
continue
return results
def compile_asm(instr: str, compiler=None) -> bytes | None:
"""Compile a single instruction with llvm-mc and return the machine code bytes."""
import subprocess
try:
result = subprocess.run(
['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
input=f".text\n{instr}\n", capture_output=True, text=True)
if result.returncode != 0: return None
# Parse encoding: [0x01,0x39,0x0a,0x7e]
for line in result.stdout.split('\n'):
if 'encoding:' in line:
enc = line.split('encoding:')[1].strip()
if enc.startswith('[') and enc.endswith(']'):
hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
return bytes.fromhex(hex_vals)
except Exception:
pass
return None
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
def _test_kernel_roundtrip(self, op_fn):
"""Generate kernel from op_fn, test:
1. decode -> reencode matches original bytes
2. asm(disasm()) matches LLVM output
3. our disasm() matches LLVM's disassembly string exactly
"""
from extra.assembly.rdna3.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.compiler_amd import HIPCompiler
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
compiler = HIPCompiler('gfx1100')
decode_passed, decode_failed, decode_skipped = 0, 0, 0
asm_passed, asm_failed, asm_skipped = 0, 0, 0
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
decode_failures, asm_failures, disasm_failures = [], [], []
for ki, kernel in enumerate(kernels):
offset = 0
while offset < len(kernel.code):
remaining = kernel.code[offset:]
fmt = detect_format(remaining)
if fmt is None:
decode_skipped += 1
asm_skipped += 1
disasm_skipped += 1
offset += 4
continue
size = fmt._size()
if len(remaining) < size:
break
orig_bytes = remaining[:size]
# Test 1: decode -> reencode roundtrip
try:
decoded = fmt.from_bytes(orig_bytes)
reencoded = decoded.to_bytes()
if reencoded[:size] == orig_bytes:
decode_passed += 1
else:
decode_failed += 1
decode_failures.append(f"K{ki}@{offset}: {decoded.disasm()}: orig={orig_bytes.hex()} reenc={reencoded[:size].hex()}")
our_disasm = decoded.disasm()
# Test 2: asm(disasm()) matches LLVM output
try:
our_bytes = asm(our_disasm).to_bytes()
llvm_bytes = compile_asm(our_disasm, compiler)
if llvm_bytes is None:
asm_skipped += 1
elif our_bytes[:len(llvm_bytes)] == llvm_bytes:
asm_passed += 1
else:
asm_failed += 1
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
except Exception:
asm_skipped += 1
# Test 3: our disasm() matches LLVM's disassembly string exactly
# Skip if instruction uses op_XX (unknown opcode) or looks malformed (many raw field values)
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm):
disasm_skipped += 1
else:
try:
# Get LLVM's disassembly of our instruction
src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n {our_disasm}\n"
lib = compiler.compile(src)
llvm_instrs = disassemble_lib(lib, compiler)
if llvm_instrs:
llvm_disasm = llvm_instrs[0][0]
if our_disasm == llvm_disasm:
disasm_passed += 1
else:
disasm_failed += 1
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
else:
disasm_skipped += 1
except Exception:
disasm_skipped += 1
except Exception:
decode_skipped += 1
asm_skipped += 1
disasm_skipped += 1
offset += size
print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
self.assertEqual(disasm_failed, 0, f"Disasm failures:\n" + "\n".join(disasm_failures[:20]))
# Basic unary ops
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Reductions
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
# Matmul
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
# Complex ops
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
# Cast ops
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
# Comparison ops
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
if __name__ == "__main__":
unittest.main()
-1
View File
@@ -1,3 +1,2 @@
*.s
*.ll
fp32_sgemm_amd
+15338
View File
File diff suppressed because it is too large Load Diff
+65
View File
@@ -0,0 +1,65 @@
# Run assembly on the AMD runtime and check correctness
# VIZ=2 to profile
import pathlib
from tinygrad import Tensor, Device, dtypes
from tinygrad.engine.realize import ExecItem, CompiledRunner
from tinygrad.renderer import ProgramSpec
from tinygrad.uop.ops import track_rewrites, UOp
from tinygrad.helpers import TracingKey
fp = pathlib.Path(__file__).parent/"gemm.s"
# ** generate inputs on CPU
N = 8192
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 baseline gemm
C_torch = A@Bt
# ** 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 = Tensor.matmul(from_torch(A), from_torch(Bt), dtype=dtypes.float32).cast(dtypes.bfloat16)
C_asm = Tensor.empty_like(C_tiny)
C_asm.uop.buffer.allocate()
# ** run gemms
@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name,), ret=ret))
def get_asm_gemm(ast:UOp, fp:pathlib.Path) -> ProgramSpec:
src = fp.read_text()
lib = Device[Device.DEFAULT].compiler.compile(src)
return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[1024, 1, 1], local_size=[256, 1, 1], globals=[0, 1, 2])
sched = C_tiny.schedule()
assert len(sched) == 1
eis:list[ExecItem] = [sched[-1].lower()]
ast = eis[0].ast
prg = get_asm_gemm(ast, fp)
eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], prg=CompiledRunner(prg)))
for ei in eis:
et = ei.run(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)
+179
View File
@@ -0,0 +1,179 @@
# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
import struct, pathlib
from tinygrad.runtime.support.elf import elf_loader
def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
with open(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
lib = f.read()
image, sections, relocs = elf_loader(lib)
rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
# rodata is exactly 64 bytes
kd = image[rodata_entry:rodata_entry+64]
desc = int.from_bytes(kd, byteorder="little")
group_segment_fixed_size = bits(desc, 0, 31)
private_segment_fixed_size = bits(desc, 32, 63)
kernarg_size = bits(desc, 64, 95)
reserved_127_96 = bits(desc, 96, 127)
assert reserved_127_96 == 0
print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
print("KERNARG_SIZE:", kernarg_size)
print("RESERVED 127:96:", reserved_127_96)
entry_off = bits(desc, 128, 191)
# sign-extend manually if needed
if entry_off & (1 << 63):
entry_off -= 1 << 64
print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
kd_addr = 0x1840
entry_addr = kd_addr + entry_off
print("Computed entry address: 0x%016x" % entry_addr)
print("256B aligned:", entry_addr % 256 == 0)
pgm_rsrc3 = bits(desc, 352, 383)
pgm_rsrc1 = bits(desc, 384, 415)
pgm_rsrc2 = bits(desc, 416, 447)
print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
# rsrc 3
accum_offset_raw = bits(pgm_rsrc3, 0, 5)
assert_zero(pgm_rsrc3, 6, 15)
tg_split = bits(pgm_rsrc3, 16, 16)
accum_offset_vgprs = (accum_offset_raw + 1) * 4
print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
print("RSRC3.TG_SPLIT:", tg_split)
# rsrc 1
vgpr_gran = bits(pgm_rsrc1, 0, 5)
sgpr_gran = bits(pgm_rsrc1, 6, 9)
assert_zero(pgm_rsrc1, 27, 28)
# NOTE: this is vgprs + agprs
vgprs_used = (vgpr_gran + 1) * 8
assert 0 <= vgprs_used <= 512
k = sgpr_gran // 2
sgprs_used = (k + 1) * 16
print("RSRC1.VGPRS:", vgprs_used)
print("RSRC1.SGPRS:", sgprs_used)
assert_zero(pgm_rsrc1, 10, 11)
float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
priv = bits(pgm_rsrc1, 20, 20)
assert priv == 0
enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
debug_mode = bits(pgm_rsrc1, 22, 22)
enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
bulky = bits(pgm_rsrc1, 24, 24)
assert bulky == 0
cdbg_user = bits(pgm_rsrc1, 25, 25)
assert cdbg_user == 0
fp16_ovfl = bits(pgm_rsrc1, 26, 26)
assert_zero(pgm_rsrc1, 27, 28) # reserved
assert_zero(pgm_rsrc1, 29, 29) # WGP_MODE (reserved on gfx9)
assert_zero(pgm_rsrc1, 30, 30) # MEM_ORDERED (reserved on gfx9)
assert_zero(pgm_rsrc1, 31, 31) # FWD_PROGRESS (reserved on gfx9)
# rsrc 2
enable_private_segment = bits(pgm_rsrc2, 0, 0) # SCRATCH_EN
user_sgpr_count = bits(pgm_rsrc2, 1, 5) # USER_SGPR
enable_trap_handler = bits(pgm_rsrc2, 6, 6) # TRAP_PRESENT (must be 0 here)
assert enable_trap_handler == 0
enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12) # TIDIG_CMP_CNT enum (0..3)
enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
assert enable_exception_address_watch == 0
enable_exception_memory = bits(pgm_rsrc2, 14, 14)
assert enable_exception_memory == 0
granulated_lds_size = bits(pgm_rsrc2, 15, 23)
assert granulated_lds_size == 0 # spec: must be 0; CP uses dispatch packet rounding
enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
assert_zero(pgm_rsrc2, 31, 31)
print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
# user sgprs
enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
enable_sgpr_queue_ptr = bits(desc, 450, 450)
enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
enable_sgpr_dispatch_id = bits(desc, 452, 452)
enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
enable_sgpr_private_segment_size = bits(desc, 454, 454)
assert_zero(desc, 455, 457)
print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
assert_zero(desc, 458, 459)
uses_dynamic_stack = bits(desc, 459, 460)
print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
assert_zero(desc, 460, 463)
kernarg_preload_spec_length = bits(desc, 464, 470)
print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
kernarg_preload_spec_offset = bits(desc, 471, 479)
print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
assert_zero(desc, 480, 511)
+46 -1
View File
@@ -1,4 +1,5 @@
use crate::work_group::WorkGroup;
use crate::state::StateSnapshot;
use crate::work_group::{WaveContext, WorkGroup};
use std::os::raw::c_char;
use std::slice;
mod helpers;
@@ -30,3 +31,47 @@ pub extern "C" fn run_asm(lib: *const c_char, lib_sz: u32, gx: u32, gy: u32, gz:
}
0
}
// FFI functions for single-stepping comparison tests
#[no_mangle]
pub extern "C" fn wave_create(lib: *const c_char, lib_sz: u32, n_lanes: u32) -> *mut WaveContext {
if lib.is_null() || (lib_sz % 4) != 0 { return std::ptr::null_mut(); }
let kernel = unsafe { slice::from_raw_parts(lib as *const u32, (lib_sz / 4) as usize).to_vec() };
Box::into_raw(Box::new(WaveContext::new(kernel, n_lanes as usize)))
}
#[no_mangle]
pub extern "C" fn wave_step(ctx: *mut WaveContext) -> i32 {
if ctx.is_null() { return -99; }
unsafe { (*ctx).step() }
}
#[no_mangle]
pub extern "C" fn wave_get_snapshot(ctx: *const WaveContext, out: *mut StateSnapshot) {
if ctx.is_null() || out.is_null() { return; }
unsafe { *out = (*ctx).get_snapshot(); }
}
#[no_mangle]
pub extern "C" fn wave_set_sgpr(ctx: *mut WaveContext, idx: u32, val: u32) {
if ctx.is_null() || idx >= 128 { return; }
unsafe { (*ctx).scalar_reg[idx as usize] = val; }
}
#[no_mangle]
pub extern "C" fn wave_set_vgpr(ctx: *mut WaveContext, lane: u32, idx: u32, val: u32) {
if ctx.is_null() || lane >= 32 || idx >= 256 { return; }
unsafe { (*ctx).vec_reg.get_lane_mut(lane as usize)[idx as usize] = val; }
}
#[no_mangle]
pub extern "C" fn wave_init_lds(ctx: *mut WaveContext, size: u32) {
if ctx.is_null() { return; }
unsafe { (*ctx).lds.data.resize(size as usize, 0); }
}
#[no_mangle]
pub extern "C" fn wave_free(ctx: *mut WaveContext) {
if !ctx.is_null() { unsafe { drop(Box::from_raw(ctx)); } }
}
+18
View File
@@ -96,6 +96,24 @@ impl WaveValue {
}
}
/// C-compatible state snapshot for FFI - used for comparing emulator states
#[repr(C)]
#[derive(Clone, Debug)]
pub struct StateSnapshot {
pub pc: u32,
pub scc: u32,
pub vcc: u32,
pub exec_mask: u32,
pub sgpr: [u32; 128],
pub vgpr: [[u32; 256]; 32],
}
impl StateSnapshot {
pub fn new() -> Self {
Self { pc: 0, scc: 0, vcc: 0, exec_mask: 0, sgpr: [0; 128], vgpr: [[0; 256]; 32] }
}
}
#[derive(Clone, Debug)]
pub struct VecDataStore {
pub data: Vec<u8>,
+91 -1
View File
@@ -1,5 +1,5 @@
use crate::helpers::{colored, DEBUG};
use crate::state::{Register, VecDataStore, WaveValue, VGPR};
use crate::state::{Register, StateSnapshot, VecDataStore, WaveValue, VGPR};
use crate::thread::{Thread, END_PRG, SGPR_COUNT};
use std::collections::HashMap;
@@ -28,6 +28,96 @@ struct WaveState {
const SYNCS: [u32; 4] = [0xBF89FC07, 0xBC7C0000, 0xBF890007, 0xbFB60003];
const S_BARRIER: u32 = 0xBFBD0000;
/// Context for single-stepping through a wave - holds all mutable state
pub struct WaveContext {
pub kernel: Vec<u32>,
pub scalar_reg: [u32; SGPR_COUNT],
pub scc: u32,
pub pc: usize,
pub vec_reg: VGPR,
pub vcc: WaveValue,
pub exec: WaveValue,
pub lds: VecDataStore,
pub sds: HashMap<usize, VecDataStore>,
pub n_lanes: usize,
}
impl WaveContext {
pub fn new(kernel: Vec<u32>, n_lanes: usize) -> Self {
let active = (!0u32).wrapping_shr(32 - (n_lanes as u32));
Self {
kernel,
scalar_reg: [0; SGPR_COUNT],
scc: 0,
pc: 0,
vec_reg: VGPR::new(),
vcc: WaveValue::new(0, n_lanes),
exec: WaveValue::new(active, n_lanes),
lds: VecDataStore::new(),
sds: (0..=31).map(|i| (i, VecDataStore::new())).collect(),
n_lanes,
}
}
/// Execute a single instruction. Returns: 0=continue, -1=endpgm, -2=barrier, 1=done (pc past program), negative=error
pub fn step(&mut self) -> i32 {
if self.pc >= self.kernel.len() { return 1; }
if self.kernel[self.pc] == END_PRG { return -1; }
if self.kernel[self.pc] == S_BARRIER { self.pc += 1; return -2; }
// Skip sync/nop instructions
if SYNCS.contains(&self.kernel[self.pc]) || self.kernel[self.pc] >> 20 == 0xbf8 || self.kernel[self.pc] == 0x7E000000 {
self.pc += 1;
return 0;
}
let mut sgpr_co = None;
for lane_id in 0..self.n_lanes {
self.vec_reg.default_lane = Some(lane_id);
self.vcc.default_lane = Some(lane_id);
self.exec.default_lane = Some(lane_id);
let mut thread = Thread {
scalar_reg: &mut self.scalar_reg,
scc: &mut self.scc,
vec_reg: &mut self.vec_reg,
vcc: &mut self.vcc,
exec: &mut self.exec,
lds: &mut self.lds,
sds: &mut self.sds.get_mut(&lane_id).unwrap(),
pc_offset: 0,
stream: self.kernel[self.pc..].to_vec(),
scalar: false,
simm: None,
warp_size: self.n_lanes,
sgpr_co: &mut sgpr_co,
};
if let Err(e) = thread.interpret() { return e; }
if thread.scalar {
self.pc = ((self.pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
break;
}
if lane_id == self.n_lanes - 1 {
self.pc = ((self.pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
}
}
if self.vcc.mutations.is_some() { self.vcc.apply_muts(); self.vcc.mutations = None; }
if self.exec.mutations.is_some() { self.exec.apply_muts(); self.exec.mutations = None; }
if let Some((idx, mut wv)) = sgpr_co.take() { wv.apply_muts(); self.scalar_reg[idx] = wv.value; }
0
}
pub fn get_snapshot(&self) -> StateSnapshot {
let mut snap = StateSnapshot::new();
snap.pc = self.pc as u32;
snap.scc = self.scc;
snap.vcc = self.vcc.value;
snap.exec_mask = self.exec.value;
snap.sgpr = self.scalar_reg;
for lane in 0..32 { snap.vgpr[lane] = self.vec_reg.get_lane(lane); }
snap
}
}
impl<'a> WorkGroup<'a> {
pub fn new(dispatch_dim: u32, id: [u32; 3], launch_bounds: [u32; 3], kernel: &'a Vec<u32>, kernel_args: *const u64) -> Self {
Self { dispatch_dim, id, kernel, launch_bounds, kernel_args, lds: VecDataStore::new(), wave_state: HashMap::new() }
+105 -91
View File
@@ -1,32 +1,37 @@
# ruff: noqa: F405, F403
# allow define from star imports
import numpy as np
import unittest
import subprocess, struct, math
import subprocess, struct, math, textwrap
from tinygrad import Tensor, dtypes, Device, UOp
from tinygrad.uop.ops import Ops
from tinygrad.helpers import getenv
from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
from tinygrad.renderer import ProgramSpec
from tinygrad.engine.realize import CompiledRunner
def get_output(asm:str, n_threads:int=1):
input_asm = "\n".join([ln if ln.strip().startswith('asm volatile') else f'asm volatile("{ln.strip().lstrip()}" : "+v"(a), "+v"(b));'
for ln in asm.strip().splitlines() if ln.strip()])
src = f"""
typedef long unsigned int size_t;
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {n_threads}))) test(unsigned int* data0_1) {{
int l = __ockl_get_local_id(0);
unsigned a = 0, b = 0, c = 0;
{input_asm}
unsigned res;
asm volatile("v_mov_b32 %0, %1" : "=v"(res) : "v"(a));
*(data0_1+l) = res;
}}"""
t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
from extra.assembly.rdna3.autogen import *
from extra.assembly.rdna3.asm import waitcnt
from test.testextra.test_cfg_viz import template
def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
out = Tensor([0]*n_threads, dtype=dtypes.uint32).realize()
src = "\n".join(inst.disasm() for inst in [
s_load_b64(s[0:1], s[0:1], NULL),
*asm,
v_lshlrev_b32_e32(v[0], 2, v[0]),
s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
#global_store_b32(v[0], v[1], s[0:1]),
global_store_b32(addr=v[0], data=vdst, saddr=s[0:1]),
s_endpgm()
])
prg = ProgramSpec("test", template.replace("fn_name", "test").replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
global_size=[1, 1, 1], local_size=[n_threads, 1, 1], globals=[0])
car = CompiledRunner(prg)
if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
car([t.uop.buffer], {}, wait=True)
return t.numpy()
car([out.uop.buffer], {}, wait=True)
return out.tolist()
def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
@@ -37,83 +42,89 @@ class TestHW(unittest.TestCase):
def setUp(self):
if getenv("MOCKGPU"): subprocess.run(["cargo", "build", "--release", "--manifest-path", "./extra/remu/Cargo.toml"], check=True)
def test_simple(self):
out = get_output("""
v_mov_b32_e32 %1 42
v_mov_b32_e32 %2 %1
""")[0]
np.testing.assert_equal(out, 42)
def test_simple_v_mov(self):
out = get_output([
v_mov_b32_e32(v[1], 2),
])
self.assertEqual(out, [2])
def test_simple_s_mov(self):
out = get_output([
s_mov_b32(s[7], 0x7fffffff),
v_mov_b32_e32(v[1], s[7]),
])
self.assertEqual(out, [0x7fffffff])
def test_exec_mov(self):
out = get_output("""
v_mov_b32_e32 %1 42
s_mov_b32_e32 exec_lo 0b10
v_mov_b32_e32 %1 10
s_mov_b32_e32 exec_lo 0b11
v_mov_b32_e32 %2 %1
""", n_threads=2)
out = get_output([
v_mov_b32_e32(v[1], 42),
s_mov_b32(EXEC_LO, 0b10),
v_mov_b32_e32(v[1], 10),
s_mov_b32(EXEC_LO, 0b11),
], n_threads=2)
np.testing.assert_equal(out, [42, 10])
def test_exec_cmp_vopc(self):
out = get_output("""
s_mov_b32 vcc_lo 0 // reset vcc
v_mov_b32_e32 %1 42
v_mov_b32_e32 %2 10
s_mov_b32_e32 exec_lo 0b01
v_cmp_ne_u32 %1 %2
s_mov_b32_e32 exec_lo 0b11
v_mov_b32_e32 %2 vcc_lo
""", n_threads=2)
np.testing.assert_equal(out, 0b01)
out = get_output([
s_mov_b32(VCC_LO, 0), # reset vcc
v_mov_b32_e32(v[1], 42),
v_mov_b32_e32(v[2], 10),
s_mov_b32(EXEC_LO, 0b01),
v_cmp_ne_u32_e32(v[1], v[2]),
s_mov_b32(EXEC_LO, 0b11),
v_mov_b32_e32(v[1], VCC_LO),
], n_threads=2)[0]
np.testing.assert_equal(out, 1)
def test_exec_cmpx_vop3(self):
out = get_output("""
s_mov_b32_e32 exec_lo 0b11
v_mov_b32_e32 %1 42
v_mov_b32_e32 %2 10
s_mov_b32_e32 exec_lo 0b01
v_cmpx_ne_u32 %1 %2
s_mov_b32_e32 s10 exec_lo
s_mov_b32_e32 exec_lo 0b11
v_mov_b32_e32 %2 s10
""", n_threads=2)[0]
out = get_output([
s_mov_b32(EXEC_LO, 0b11),
v_mov_b32_e32(v[1], 42),
v_mov_b32_e32(v[2], 10),
s_mov_b32(EXEC_LO, 0b01),
v_cmpx_ne_u32_e32(v[1], v[2]),
s_mov_b32(s[10], EXEC_LO),
s_mov_b32(EXEC_LO, 0b11),
v_mov_b32_e32(v[1], s[10]),
], n_threads=2)[0]
np.testing.assert_equal(out & 0b11, 0b01)
def test_fmac_vop3_modifier(self):
init_state = f"""
asm volatile("v_mov_b32_e32 %1, {f16_to_bits(4.0)}" : "+v"(a));
asm volatile("v_mov_b32_e32 %1, {f16_to_bits(3.0)}" : "+v"(b));
asm volatile("v_mov_b32_e32 %1, {f16_to_bits(2.0)}" : "+v"(c));
"""
mov = """asm volatile("v_mov_b32_e32 %1, %2" : "+v"(c), "+v"(a));"""
def fmac(a, b, c): return f"""asm volatile("v_fmac_f16_e64 {c}, {a}, {b}" : "+v"(c) : "v"(a), "v"(b));"""+"\n"+mov
self.assertEqual(get_output(init_state+"\n"+fmac("%1", "%2", "%3")), f16_to_bits(14.))
self.assertEqual(get_output(init_state+"\n"+fmac("%1", "-%2", "%3")), f16_to_bits(-10.))
self.assertEqual(get_output(init_state+"\n"+fmac("-%1", "-%2", "%3")), f16_to_bits(14.))
init_state = [
v_mov_b32_e32(a:=v[1], f16_to_bits(4.0)),
v_mov_b32_e32(b:=v[2], f16_to_bits(3.0)),
v_mov_b32_e32(c:=v[3], f16_to_bits(2.0)),
]
def run_fmac(a, b): return get_output(init_state+[v_fmac_f16_e64(c, a, b)], vdst=c)[0]
self.assertEqual(run_fmac(a, b), f16_to_bits(14.0))
self.assertEqual(run_fmac(a, -b), f16_to_bits(-10.0))
self.assertEqual(run_fmac(-a, -b), f16_to_bits(14.0))
def test_s_abs_i32(self):
def s_abs_i32(x, y, dst="s10", scc=0):
for reg,val in [(dst, y), ("scc", scc)]:
self.assertEqual(get_output(f"""
s_mov_b32_e32 {dst} {x}
s_abs_i32 {dst} {dst}
v_mov_b32_e32 %2 {reg}
""")[0], val)
s_abs_i32(0x00000001, 0x00000001, scc=1)
s_abs_i32(0x7fffffff, 0x7fffffff, scc=1)
s_abs_i32(0x80000000, 0x80000000, scc=1)
s_abs_i32(0x80000001, 0x7fffffff, scc=1)
s_abs_i32(0x80000002, 0x7ffffffe, scc=1)
s_abs_i32(0xffffffff, 0x00000001, scc=1)
s_abs_i32(0, 0, scc=0)
def check(x, y, dst=s[10], scc=0):
for reg,val in [(dst, y), (SCC, scc)]:
self.assertEqual(get_output([
s_mov_b32(dst, x),
s_abs_i32(dst, dst),
v_mov_b32_e32(v[1], reg)
])[0], val)
check(0x00000001, 0x00000001, scc=1)
check(0x7fffffff, 0x7fffffff, scc=1)
check(0x80000000, 0x80000000, scc=1)
check(0x80000001, 0x7fffffff, scc=1)
check(0x80000002, 0x7ffffffe, scc=1)
check(0xffffffff, 0x00000001, scc=1)
check(0, 0, scc=0)
def test_v_rcp_f32_neg_vop3(self):
def v_neg_rcp_f32(x:float, y:float):
out = get_output(f"""
v_mov_b32_e32 %2 {f32_to_bits(x)}
v_rcp_f32_e64 %2, -%2
""")[0]
out = get_output([
v_mov_b32_e32(v[2], f32_to_bits(x)),
v_rcp_f32_e64(v[2], -v[2]),
], vdst=v[2])[0]
assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
v_neg_rcp_f32(math.inf, -0.0)
v_neg_rcp_f32(-math.inf, 0.0)
v_neg_rcp_f32(0.0, -math.inf)
@@ -122,25 +133,28 @@ class TestHW(unittest.TestCase):
v_neg_rcp_f32(2.0, -0.5)
def test_v_cndmask_b32_neg(self):
def v_neg(x:int|float, y:float):
# always pick -v1
out = get_output(f"""
v_mov_b32_e32 %2 {f32_to_bits(x)}
s_mov_b32_e32 s10 1
v_cndmask_b32 %2, %2, -%2 s10
""")[0]
def v_neg(x:float, y:float):
out = get_output([
v_mov_b32_e32(v[1], f32_to_bits(x)),
s_mov_b32(s[10], 1),
v_cndmask_b32_e64(v[1], v[1], -v[1], s[10]),
])[0]
assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
v_neg(-0.0, 0.0)
v_neg(0.0, -0.0)
v_neg(2.0, -2.0)
v_neg(math.inf, -math.inf)
v_neg(-math.inf, math.inf)
@unittest.skip("how does VOPD work in the dsl")
def test_v_subrev_wrap(self):
out = get_output("""
v_dual_mov_b32 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
v_subrev_co_u32 %2, vcc_lo, %2, %1
""")[0]
out = get_output([
#v_dual_mov_b32(v[1], 0xffffffff, v[2], 0x0),
#v_dual_mov_b32(vdstx=v[1], srcx=0xffffffff, vdsty=v[2], srcy=0x0),
#VOPD(opx=VOPDOp.V_DUAL_MOV_B32, opy=VOPDOp.V_DUAL_MOV_B32, vdstx=v[1], srcx=0xffffffff, vdsty=v[2], srcy=0x0),
v_subrev_co_u32(v[2], VCC_LO, v[2], v[1]),
], vdst=v[2])[0]
self.assertEqual(out, 0xffff_ffff)
if __name__ == "__main__":
+2 -2
View File
@@ -534,8 +534,8 @@ def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -
def parse(fn:str):
with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
if getenv("ROCM", 0):
with Timing(f"decode {fn}: "): ctx = decode(dat)
#if getenv("ROCM", 0):
# with Timing(f"decode {fn}: "): ctx = decode(dat)
dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
print(f"got {len(dat_sqtt)} SQTT events in {fn}")
return dat_sqtt
+6 -20
View File
@@ -48,17 +48,11 @@ class OccEvent(WaveSlot):
RunKey = tuple[str, int]
class _ROCParseCtx:
def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
self.disasms:dict[str, dict[int, tuple[str, int]]] = {}
def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[str, int]]]):
self.sqtt_evs, self.disasms = iter(sqtt_evs), disasms
self.inst_execs:dict[RunKey, list[WaveExec]] = {}
self.occ_events:dict[RunKey, list[OccEvent]] = {}
for prog in prog_evs:
arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
base = unwrap(prog.base)
self.disasms[prog.name] = asm = {base+addr:info for addr,info in llvm_disasm(arch, unwrap(prog.lib)).items()}
def next_sqtt(self):
x = next(self.sqtt_evs, None)
self.active_run = (x.kern, x.exec_tag) if x is not None else None
@@ -81,16 +75,8 @@ class _ROCParseCtx:
self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
ev.end_time, insts_blob))
def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
dev_events:dict[str, ProfileDeviceEvent] = {}
sqtt_events:list[ProfileSQTTEvent] = []
prog_events:list[ProfileProgramEvent] = []
for e in profile:
if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, tuple[str, int]]]) -> _ROCParseCtx:
ROCParseCtx = _ROCParseCtx(sqtt_evs, disasms)
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, _):
@@ -150,7 +136,7 @@ if __name__ == "__main__":
args = parser.parse_args()
with args.profile.open("rb") as f: profile = pickle.load(f)
rctx = decode(profile)
print('SQTT:', rctx.inst_execs.keys())
#rctx = decode(profile, disasm)
#print('SQTT:', rctx.inst_execs.keys())
print_pmc([ev for ev in profile if isinstance(ev, ProfilePMCEvent)])
+2
View File
@@ -185,6 +185,8 @@ class PM4Executor(AMDQueue):
if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
# Pass valid memory ranges to Python emulator for bounds checking
if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
err = remu.run_asm(prg_addr, prg_sz, *gl, *lc, args_addr)
if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
+14
View File
@@ -1,4 +1,5 @@
import ctypes, ctypes.util
from tinygrad.helpers import getenv
def _try_dlopen_gpuocelot():
GPUOCELOT_PATHS = [ctypes.util.find_library("gpuocelot")] if ctypes.util.find_library("gpuocelot") is not None else []
@@ -14,7 +15,20 @@ def _try_dlopen_gpuocelot():
print("Could not find libgpuocelot.so")
return None
class PythonRemu:
"""Python RDNA3 emulator wrapper that matches the libremu.so interface."""
valid_mem_ranges: set[tuple[int, int]] = set()
def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
from extra.assembly.rdna3.emu import run_asm, set_valid_mem_ranges
# Pad ranges to handle GPU loads that may read past small buffers (e.g. s_load_b128 on 12-byte buffer)
set_valid_mem_ranges({(start, size + 4096) for start, size in self.valid_mem_ranges})
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr)
def _try_dlopen_remu():
# Use Python emulator only if PYTHON_REMU=1
if getenv("PYTHON_REMU"):
return PythonRemu()
REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
"extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
for path in REMU_PATHS:
+6
View File
@@ -57,12 +57,18 @@ class TestMultiTensor(unittest.TestCase):
assert lb.shape == (128,)
(X + X).realize()
@unittest.expectedFailure # TODO: fix
def test_shard_empty(self):
GlobalCounters.reset()
X = Tensor.empty(256).shard(devices_2, 0).realize()
assert GlobalCounters.kernel_count == 0
(X + X).realize()
def test_arange_shrink(self):
x = Tensor.arange(4)
self.assertEqual(x.shard(devices_2, 0).realize().shrink(((2, 4),)).tolist(), [2, 3])
self.assertEqual(x.shard(devices_2, 0).realize().shrink(((0, 2),)).tolist(), [0, 1])
def test_shard_like(self):
X = Tensor.ones(256).shard(devices_2, 0)
Y = Tensor.zeros(256).shard_like(X)
+6
View File
@@ -544,6 +544,12 @@ class TestUopsObject(unittest.TestCase):
self.assertEqual(a.device, Device.DEFAULT)
class TestUOpRender(unittest.TestCase):
def test_render_vectorize_empty(self):
u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(0), src=())
self.assertEqual(u.render(simplify=False), "{}")
def test_render_vectorize_empty_simplified(self):
u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(0), src=())
self.assertEqual(u.render(), "{}")
def test_render_vectorize_same(self):
u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
self.assertEqual(u.render(simplify=False), "{0, ...}")
+101 -93
View File
@@ -1,12 +1,17 @@
# ruff: noqa: F405, F403
# allow define from star imports
import unittest
import textwrap
from tinygrad import Device, Tensor
from tinygrad.uop.ops import UOp, Ops, track_rewrites
from tinygrad.renderer import ProgramSpec
from tinygrad.helpers import TracingKey
from tinygrad.helpers import TracingKey, getenv
from tinygrad.engine.realize import ExecItem, CompiledRunner
from extra.assembly.rdna3.autogen import *
# TODO: use the RDNA3 renderer when it's in master
template = """.text
.globl fn_name
@@ -18,6 +23,7 @@ fn_name:
.rodata
.p2align 6
.amdhsa_kernel fn_name
.amdhsa_kernarg_size 8
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
.amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
@@ -52,13 +58,15 @@ amdhsa.kernels:
"""
@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, ret=ret))
def run_asm(name:str, src:str) -> ProgramSpec:
prg = ProgramSpec(name, template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK))
def run_asm(name:str, insts:list) -> ProgramSpec:
src = "\n".join([inst if isinstance(inst, str) else inst.disasm() for inst in insts])
prg = ProgramSpec(name, src:=template.replace("fn_name", name).replace("INSTRUCTION", textwrap.dedent(src)), Device.DEFAULT, UOp(Ops.SINK),
lib=Device[Device.DEFAULT].compiler.compile(src), global_size=[1, 1, 1], local_size=[1, 1, 1], globals=[0])
ei = ExecItem(UOp(Ops.SINK), [Tensor.empty(1).uop.buffer.ensure_allocated()], prg=CompiledRunner(prg))
ei.run()
return prg
@unittest.skipUnless(Device.DEFAULT == "AMD", "only on AMD")
@unittest.skipUnless(Device.DEFAULT == "AMD" and not getenv("AMD_LLVM"), "only on AMD with comgr")
class TestCfg(unittest.TestCase):
def setUp(self):
arch = Device["AMD"].arch
@@ -66,107 +74,107 @@ class TestCfg(unittest.TestCase):
self.skipTest(f"tests written for RDNA, got arch {arch}")
def test_simple(self):
run_asm("simple", """
entry:
s_branch bb1
bb1:
s_endpgm
""")
run_asm("simple", [
"entry:",
s_branch("bb1"),
"bb1:",
s_endpgm(),
])
def test_diamond(self):
run_asm("diamond", """
entry:
s_cmp_eq_i32 s0, 0
s_cbranch_scc1 if
s_branch else
if:
s_nop 1
s_branch end
else:
s_nop 0
end:
s_endpgm
""")
run_asm("diamond", [
"entry:",
s_cmp_eq_i32(s[0], 0),
s_cbranch_scc1("if"),
s_branch("else"),
"if:",
s_nop(1),
s_branch("end"),
"else:",
s_nop(0),
"end:",
s_endpgm(),
])
def test_loop(self):
run_asm("simple_loop", """
entry:
s_mov_b32 s1, 4
loop:
s_add_u32 s1, s1, -1
s_cmp_eq_i32 s1, 0
s_cbranch_scc0 loop
s_endpgm
""")
run_asm("simple_loop", [
"entry:",
s_mov_b32(s[1], 4),
"loop:",
s_add_u32(s[1], s[1], -1),
s_cmp_eq_i32(s[1], 0),
s_cbranch_scc0("loop"),
s_endpgm(),
])
def test_loop_branch(self):
run_asm("loop_if", """
entry:
s_mov_b32 s1, 4
loop:
s_add_u32 s1, s1, -1
s_cmp_eq_i32 s1, 2
s_cbranch_scc1 cond
s_branch cont
cond:
s_add_u32 s1, s1, -2
cont:
s_cmp_eq_i32 s1, 0
s_cbranch_scc0 loop
s_endpgm
""")
run_asm("loop_if", [
"entry:",
s_mov_b32(s[1], 4),
"loop:",
s_add_u32(s[1], s[1], -1),
s_cmp_eq_i32(s[1], 2),
s_cbranch_scc1("cond"),
s_branch("cont"),
"cond:",
s_add_u32(s[1], s[1], -2),
"cont:",
s_cmp_eq_i32(s[1], 0),
s_cbranch_scc0("loop"),
s_endpgm(),
])
def test_loop_break(self):
run_asm("loop_break", """
entry:
s_mov_b32 s1, 8
loop:
s_add_u32 s1, s1, -1
s_cmp_eq_i32 s1, 5
s_cbranch_scc1 break
s_cmp_eq_i32 s1, 0
s_cbranch_scc0 loop
break:
s_endpgm
""")
run_asm("loop_break", [
"entry:",
s_mov_b32(s[1], 8),
"loop:",
s_add_u32(s[1], s[1], -1),
s_cmp_eq_i32(s[1], 5),
s_cbranch_scc1("break"),
s_cmp_eq_i32(s[1], 0),
s_cbranch_scc0("loop"),
"break:",
s_endpgm(),
])
def test_switch(self):
run_asm("switch_case", """
entry:
s_cmp_eq_i32 s0, 0
s_cbranch_scc1 case0
s_cmp_eq_i32 s0, 1
s_cbranch_scc1 case1
s_branch case2
case0:
s_nop 0
s_branch join
case1:
s_nop 1
s_branch join
case2:
s_nop 2
s_branch join
join:
s_endpgm
""")
run_asm("switch_case", [
"entry:",
s_cmp_eq_i32(s[0], 0),
s_cbranch_scc1("case0"),
s_cmp_eq_i32(s[0], 1),
s_cbranch_scc1("case1"),
s_branch("case2"),
"case0:",
s_nop(0),
s_branch("join"),
"case1:",
s_nop(1),
s_branch("join"),
"case2:",
s_nop(2),
s_branch("join"),
"join:",
s_endpgm(),
])
def test_ping_pong(self):
run_asm("ping_pong", """
entry:
s_cmp_eq_i32 s0, 0
s_cbranch_scc1 ping
s_branch pong
ping:
s_cmp_eq_i32 s1, 0
s_cbranch_scc1 pong
s_branch end
pong:
s_cmp_eq_i32 s2, 0
s_cbranch_scc1 ping
end:
s_endpgm
""")
run_asm("ping_pong", [
"entry:",
s_cmp_eq_i32(s[0], 0),
s_cbranch_scc1("ping"),
s_branch("pong"),
"ping:",
s_cmp_eq_i32(s[1], 0),
s_cbranch_scc1("pong"),
s_branch("end"),
"pong:",
s_cmp_eq_i32(s[2], 0),
s_cbranch_scc1("ping"),
"end:",
s_endpgm(),
])
if __name__ == "__main__":
unittest.main()
+14 -2
View File
@@ -1,6 +1,6 @@
import ctypes, gzip, unittest, timeit
import ctypes, gzip, unittest, timeit, pickle
from tinygrad import Variable
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count
from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
from tinygrad.tensor import Tensor, get_shape
import numpy as np
@@ -120,6 +120,18 @@ class TestRoundUp(unittest.TestCase):
self.assertEqual(round_up(232, 24984), 24984)
self.assertEqual(round_up(24984, 232), 25056)
class TestCount(unittest.TestCase):
def test_count_basic(self):
c = count(3)
self.assertEqual(next(c), 3)
self.assertEqual(next(c), 4)
def test_count_step_pickle(self):
c = count(1, 2)
self.assertEqual(next(c), 1)
c2 = pickle.loads(pickle.dumps(c))
self.assertEqual(next(c2), 3)
@unittest.skip("no fetch tests because they need internet")
class TestFetch(unittest.TestCase):
def test_fetch_bad_http(self):
+2 -2
View File
@@ -3,7 +3,7 @@ from unittest.mock import patch
from io import StringIO
from collections import namedtuple
from tqdm import tqdm
from tinygrad.helpers import tqdm as tinytqdm, trange as tinytrange
from tinygrad.helpers import tqdm as tinytqdm, trange
import numpy as np
def _get_iter_per_second(raw:str) -> float:
@@ -165,7 +165,7 @@ class TestProgressBar(unittest.TestCase):
mock_stderr.truncate(0)
# compare bars at each iteration (only when tinytqdm bar has been updated)
for n in (bar := tinytrange(total, desc="Test")):
for n in (bar := trange(total, desc="Test")):
if bar.i % bar.skip != 0: continue
tiny_output = mock_stderr.getvalue().split("\r")[-1].rstrip()
iters_per_sec = float(tiny_output.split("it/s")[-2].split(" ")[-1]) if n>0 else 0
+2 -3
View File
@@ -1,4 +1,3 @@
from __future__ import annotations
import math, itertools
from collections import defaultdict
from typing import cast, Final
@@ -7,7 +6,7 @@ from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
from tinygrad.device import Buffer
from tinygrad.dtype import dtypes, ImageDType
from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
from tinygrad.helpers import ALLOW_TF32
from tinygrad.helpers import ALLOW_TF32, count
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
from tinygrad.codegen.simplify import pm_flatten_range
from tinygrad.renderer import Renderer
@@ -19,7 +18,7 @@ class Scheduler:
self.ast, self.ren = ast, ren
self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
self.opt_range = itertools.count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
self.opt_range = count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
@property
def rngs(self):
+1 -2
View File
@@ -199,8 +199,7 @@ class Buffer:
return mv
def view(self, size:int, dtype:DType, offset:int) -> Buffer:
assert offset < self.nbytes, "offset must be less than nbytes"
if self._base is not None: return Buffer(self.device, size, dtype, base=self._base, offset=self.offset+offset)
return Buffer(self.device, size, dtype, base=self, offset=offset)
return Buffer(self.device, size, dtype, base=self.base, offset=self.offset+offset)
@dataclass(frozen=True)
class DMACPURef:
+3 -3
View File
@@ -67,13 +67,13 @@ class ViewOp(Runner):
class BufferCopy(Runner):
def __init__(self, total_sz, dest_device, src_device):
if total_sz >= 1e6: name = f"{type(self).__name__[6:].lower()} {total_sz/1e6:7.2f}M, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
else: name = f"{type(self).__name__[6:].lower()} {total_sz:8d}, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
sz = f"{total_sz/1e6:7.2f}M" if total_sz >= 1e6 else f"{total_sz:8d}"
name = f"{type(self).__name__[6:].lower()} {sz}, {dest_device[:7]:>7s} <- {src_device[:7]:7s}"
super().__init__(colored(name, "yellow"), dest_device, Estimates(lds=total_sz, mem=total_sz))
def copy(self, dest, src):
disk_supports_fast_copyout = src.device.startswith("DISK") and hasattr(src.allocator.dev, 'io_uring') and \
getattr(src.allocator.dev, 'fd', None) is not None and dest.allocator.supports_copy_from_disk
if src.device.startswith("DISK") and hasattr(dest.allocator, 'copy_from_disk') and disk_supports_fast_copyout and src.nbytes >= 4096:
if disk_supports_fast_copyout and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096:
dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
elif (src.device.startswith("DISK") or src.device.startswith("TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
# fast(ish) path, uses readinto in diskbuffers
+1 -3
View File
@@ -1,9 +1,7 @@
from __future__ import annotations
import time
from typing import cast
from collections import deque
from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites
from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map
from tinygrad.uop.spec import type_verify, tensor_spec
from tinygrad.device import Buffer, MultiBuffer
from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize
+1 -1
View File
@@ -59,7 +59,7 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
grads = {root: root_grad}
for t0 in reversed(_deepwalk(root, targets)):
if t0 not in grads: continue
lgrads: tuple[UOp|None, ...]|None = cast(tuple[UOp, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
lgrads: tuple[UOp|None, ...]|None = cast(tuple[UOp|None, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0]))
if lgrads is None: raise RuntimeError(f"failed to compute gradient for {t0.op}\n\nin {str(t0)[0:1000]}...")
assert len(lgrads) == len(t0.src), f"got {len(lgrads)} gradient, expected {len(t0.src)}"
for k,v in zip(t0.src, lgrads):
+9 -2
View File
@@ -523,8 +523,7 @@ class tqdm(Generic[T]):
@classmethod
def write(cls, s:str): print(f"\r\033[K{s}", flush=True, file=sys.stderr)
class trange(tqdm):
def __init__(self, n:int, **kwargs): super().__init__(iterable=range(n), total=n, **kwargs)
def trange(n:int, **kwargs) -> tqdm[int]: return tqdm(range(n), total=n, **kwargs)
class disable_gc(contextlib.ContextDecorator):
def __enter__(self):
@@ -543,3 +542,11 @@ copyreg.pickle(types.CodeType, _serialize_code)
def _serialize_module(module:types.ModuleType): return importlib.import_module, (module.__name__,)
copyreg.pickle(types.ModuleType, _serialize_module)
class count:
def __init__(self, start:int=0, step:int=1):
self.n, self.step = start, step
def __next__(self) -> int:
cur = self.n
self.n += self.step
return cur
+17 -18
View File
@@ -1,5 +1,4 @@
# mixins add syntactic sugar to Tensor and UOp
import functools
from typing import TypeAlias, TYPE_CHECKING, Self
from tinygrad.uop import Ops
from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv
@@ -282,38 +281,38 @@ class MovementMixin:
```
"""
def parse_formula(formula: str):
tokens = f" {formula} ".replace("", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", " ").replace(" 1 ", " ( ) ").split()
lparens, rparens = map(lambda x: [i for i, ch in enumerate(tokens) if ch == x], ("(", ")"))
def parse_side(s: str) -> tuple[list[str], list[tuple[int, int]]]:
"""Parse one side of formula into (axis_names, dims) where dims are (start, end) index pairs for parens."""
tokens = f" {s} ".replace("", "...").replace("(", " ( ").replace(")", " ) ").replace(" ", " ").replace(" 1 ", " ( ) ").split()
lparens, rparens = [i for i, tok in enumerate(tokens) if tok == "("], [i for i, tok in enumerate(tokens) if tok == ")"]
pairs = list(zip(lparens, rparens))
assert len(lparens) == len(rparens) and sorted(flatten(pairs)) == flatten(pairs), "bracket mismatch"
return [name for name in tokens if name not in ("(", ")")], [(s - 2 * i, e - 1 - 2 * i) for i, (s, e) in enumerate(pairs)]
return [tok for tok in tokens if tok not in ("(", ")")], [(lp - 2*i, rp - 1 - 2*i) for i, (lp, rp) in enumerate(pairs)]
assert formula.count("->") == 1, 'need exactly one "->" in formula'
(lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_side, formula.split("->"))
(lhs, unflatten_dims), (rhs, flatten_dims) = map(parse_formula, formula.split("->"))
for name in sizes:
assert name in lhs, f"axis {name} is not used in transform"
for name in sizes: assert name in lhs, f"axis {name} is not used in transform"
assert sorted(lhs) == sorted(rhs) and len(lhs) == len(set(lhs)), f"name mismatch in {formula}"
for name in flatten((lhs, rhs)):
assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
for name in lhs+rhs: assert name == "..." or (name.isidentifier() and "_" not in (name[0], name[-1])), f"invalid axis name {name}"
assert "..." not in flatten([lhs[s:e] for s, e in unflatten_dims]), f"cannot have collapsed ellipsis (...) in lhs of {formula}"
assert lhs.count("...") <= 1, f"too many ellipses in {formula}"
# resolve ellipsis
if "..." in lhs:
ell_len = len(self.shape) - len(lhs) + 1 + sum(e - s - 1 for s, e in unflatten_dims)
lhs, rhs = map(lambda l: l[: (i := l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1 :] if "..." in l else l, (lhs, rhs))
unflatten_dims = [(s + (ell_len - 1 if "...0" in lhs[:s] else 0), e + (ell_len - 1 if "...0" in lhs[:e] else 0)) for s, e in unflatten_dims]
flatten_dims = [(s + (ell_len - 1 if "...0" in rhs[:s] else 0), e + (ell_len - 1 if "...0" in rhs[:e] else 0)) for s, e in flatten_dims]
lhs, rhs = map(lambda l: l[:(i := l.index("..."))] + [f"...{j}" for j in range(ell_len)] + l[i + 1:] if "..." in l else l, (lhs, rhs))
def newdims(side, s, e): return (s + (ell_len - 1 if "...0" in side[:s] else 0), e + (ell_len - 1 if "...0" in side[:e] else 0))
unflatten_dims, flatten_dims = [newdims(lhs, s, e) for s, e in unflatten_dims], [newdims(rhs, s, e) for s, e in flatten_dims]
# apply movement ops in order unflatten -> permute -> flatten/unsqueeze
t = functools.reduce(lambda x, dims: x.unflatten(dims[0], tuple(sizes.get(lhs[d], -1) for d in range(*dims))), unflatten_dims, self)
# unflatten -> permute -> flatten
t = self
for start, end in unflatten_dims: t = t.unflatten(start, tuple(sizes.get(lhs[i], -1) for i in range(start, end)))
for i, name in enumerate(lhs):
assert (name not in sizes) or sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
if name in sizes: assert sizes[name] == t.shape[i], f"size provided for dimension {name} incorrect"
t = t.permute([lhs.index(name) for name in rhs])
return functools.reduce(lambda x, dims: x.flatten(dims[0], dims[1] - 1) if dims[0] < dims[1] else x.unsqueeze(dims[0]), reversed(flatten_dims), t)
for start, end in reversed(flatten_dims): t = t.flatten(start, end - 1) if start < end else t.unsqueeze(start)
return t
# *** movement ops with expand ***
+12 -5
View File
@@ -382,8 +382,8 @@ class CUDARenderer(CStyleLanguage):
shared_max = 49152
def __init__(self, arch:str):
self.arch = arch
self.tensor_cores = tc.cuda_sm89 if int(arch[3:]) >= 89 else tc.cuda_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else []
self.arch, arch_ver = arch, int(arch[3:])
self.tensor_cores = tc.cuda_sm89 if arch_ver >= 89 else tc.cuda_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
def __reduce__(self): return self.__class__, (self.arch,)
# language options
@@ -446,7 +446,7 @@ class CUDARenderer(CStyleLanguage):
return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix)
class AMDRenderer(CStyleLanguage):
class AMDHIPRenderer(CStyleLanguage):
device = "AMD"
shared_max = 65536
# NOTE: this is only really needed on gfx12, even though gfx11 reports the same limitation
@@ -458,7 +458,8 @@ class AMDRenderer(CStyleLanguage):
@staticmethod
def is_cdna(arch): return arch.split(":")[0] in {"gfx942", "gfx950"}
def __init__(self, arch:str): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700
self.arch = arch
from tinygrad.runtime.support.compiler_amd import HIPCompiler
self.arch, self.compiler = arch, HIPCompiler(arch)
self.tensor_cores = self.get_tensor_cores(arch)
if self.is_cdna(self.arch):
self.string_rewrite = PatternMatcher([
@@ -539,5 +540,11 @@ class AMDRenderer(CStyleLanguage):
return super().render_kernel(function_name, kernel, bufs, uops, prefix)
class NVRenderer(CUDARenderer): device = "NV"
class HIPRenderer(AMDRenderer): device = "HIP"
class HIPRenderer(AMDHIPRenderer): device = "HIP"
class AMDHIPCCRenderer(AMDHIPRenderer):
def __init__(self, arch:str):
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
super().__init__(arch)
self.compiler = HIPCCCompiler(arch)
class QCOMRenderer(OpenCLRenderer): device = "QCOM"
+7 -6
View File
@@ -2,7 +2,7 @@ from typing import cast
import math, struct, sys
from tinygrad.codegen.opt import tc
from tinygrad.renderer import Renderer
from tinygrad.renderer.cstyle import AMDRenderer, create_non_native_float_pats, pm_manual_bf16_cast
from tinygrad.renderer.cstyle import AMDHIPRenderer, create_non_native_float_pats, pm_manual_bf16_cast
from tinygrad.uop.decompositions import xexp2, xlog2
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate
@@ -209,8 +209,8 @@ llvm_intrinsics = {Ops.SQRT: "sqrt", Ops.LOG2: "log2", Ops.EXP2: "exp2"}
class AMDLLVMRenderer(LLVMRenderer):
device = "AMD"
has_local = True
shared_max = AMDRenderer.shared_max
global_max = AMDRenderer.global_max
shared_max = AMDHIPRenderer.shared_max
global_max = AMDHIPRenderer.global_max
abi = "amdgpu_kernel"
code_for_op = {**LLVMRenderer.code_for_op, **{op: lambda: None for op in llvm_intrinsics}}
string_rewrite = PatternMatcher([
@@ -254,9 +254,10 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc
f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
def __init__(self, arch:str):
self.arch = arch
self.tensor_cores = AMDRenderer.get_tensor_cores(arch)
self.is_cdna = AMDRenderer.is_cdna(arch)
from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
self.arch, self.compiler = arch, AMDLLVMCompiler(arch)
self.tensor_cores = AMDHIPRenderer.get_tensor_cores(arch)
self.is_cdna = AMDHIPRenderer.is_cdna(arch)
self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
if self.is_cdna:
self.extra_matcher += PatternMatcher([
+2 -2
View File
@@ -145,8 +145,8 @@ class PTXRenderer(Renderer):
code_for_op = asm_for_op
extra_matcher = ptx_matcher
def __init__(self, arch:str, device="CUDA"):
self.device, self.arch = device, arch
self.tensor_cores = PTXRenderer.tc_sm80 if int(arch[3:]) >= 80 else tc.cuda_sm75 if int(arch[3:]) >= 75 else []
self.device, self.arch, arch_ver = device, arch, int(arch[3:])
self.tensor_cores = PTXRenderer.tc_sm80 if arch_ver >= 80 else tc.cuda_sm75 if arch_ver >= 75 else []
def __reduce__(self): return self.__class__, (self.arch, self.device)
# language options
+9 -6
View File
@@ -9,11 +9,10 @@ from tinygrad.uop.ops import sint
from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerSet, CompilerPair
from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar
from tinygrad.helpers import VIZ, AMD_CC, AMD_LLVM, ceildiv
from tinygrad.renderer.cstyle import AMDRenderer
from tinygrad.renderer.cstyle import AMDHIPRenderer, AMDHIPCCRenderer
from tinygrad.renderer.llvmir import AMDLLVMRenderer
from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
from tinygrad.runtime.autogen.am import am
from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
@@ -649,7 +648,11 @@ class AMDQueueDesc:
# Flush hdp if queue is in dev mem.
if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
try:
for doorbell in self.doorbells: doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
except Exception as e:
dev.error_state = e
raise
class KFDIface:
kfd:FileIOInterface|None = None
@@ -931,9 +934,9 @@ class AMDDevice(HCQCompiled):
max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
compilers = CompilerSet([CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch), AMD_LLVM),
CompilerPair(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))], ctrl_var=AMD_CC)
compilers = CompilerSet([CompilerPair(functools.partial(AMDHIPRenderer, self.arch), None),
CompilerPair(functools.partial(AMDLLVMRenderer, self.arch), None, AMD_LLVM),
CompilerPair(functools.partial(AMDHIPCCRenderer, self.arch), None)], ctrl_var=AMD_CC)
super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
+1 -2
View File
@@ -2,7 +2,6 @@ import ctypes, functools
from tinygrad.helpers import init_c_var, mv_address, init_c_struct_t, getenv
from tinygrad.device import Compiled, LRUAllocator, BufferSpec, CompilerSet, CompilerPair
from tinygrad.runtime.autogen import hip
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.renderer.cstyle import HIPRenderer
if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import
@@ -15,7 +14,7 @@ class HIPDevice(Compiled):
self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device_id))).gcnArchName.decode()
self.time_event_st, self.time_event_en = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), functools.partial(HIPCompiler, self.arch))])
compilers = CompilerSet([CompilerPair(functools.partial(HIPRenderer, self.arch), None)])
super().__init__(device, HIPAllocator(self), compilers, functools.partial(HIPProgram, self))
def synchronize(self):
check(hip.hipSetDevice(self.device_id))
+1 -1
View File
@@ -508,7 +508,7 @@ class PCIIface(PCIIfaceBase):
# PCIIface's MAP_FIXED mmap will overwrite UVM allocations made by NVKIface, so don't try PCIIface if kernel driver was already used.
if NVKIface.root is not None: raise RuntimeError("Cannot use PCIIface after NVKIface has been initialized (would corrupt UVM memory)")
super().__init__(dev, dev_id, vendor=0x10de, devices=[(0xff00, [0x2200, 0x2400, 0x2500, 0x2600, 0x2700, 0x2800, 0x2b00, 0x2c00, 0x2d00, 0x2f00])],
bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
base_class=0x03, bars=[0, 1], vram_bar=1, va_start=NVMemoryManager.va_allocator.base, va_size=NVMemoryManager.va_allocator.size)
if not OSX: System.reserve_hugepages(64)
self.pci_dev.write_config(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
+2 -2
View File
@@ -161,7 +161,7 @@ class AMDev(PCIDevImplBase):
# To enable this, AM uses a separate boot memory that is guaranteed not to be overwritten. This physical memory is utilized for
# all blocks that are initialized only during the initial AM boot.
# To determine if the GPU is in the third state, AM uses regSCRATCH_REG7 as a flag.
self.is_booting = True
self.is_booting = True # During boot only boot memory can be allocated. This flag is to validate this.
self.init_sw(smi_dev=False)
self.partial_boot = (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
@@ -194,7 +194,7 @@ class AMDev(PCIDevImplBase):
if DEBUG >= 2: print(f"am {self.devfmt}: boot done")
def init_sw(self, smi_dev=False):
self.smi_dev = smi_dev # During boot only boot memory can be allocated. This flag is to validate this.
self.smi_dev = smi_dev
# Memory manager & firmware
self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
+2 -2
View File
@@ -449,8 +449,8 @@ class AM_SDMA(AM_IP):
reg, inst = ("regSDMA_GFX", pipe*4+queue) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x1, inst=inst)
self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
self.adev.wreg_pair(f"{reg}_RB_WPTR", "", "_HI", 0, inst=inst)
if not self.adev.partial_boot: self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
if not self.adev.partial_boot: self.adev.wreg_pair(f"{reg}_RB_WPTR", "", "_HI", 0, inst=inst)
self.adev.wreg_pair(f"{reg}_RB_BASE", "", "_HI", ring_addr >> 8, inst=inst)
self.adev.wreg_pair(f"{reg}_RB_RPTR_ADDR", "_LO", "_HI", rptr_addr, inst=inst)
self.adev.wreg_pair(f"{reg}_RB_WPTR_POLL_ADDR", "_LO", "_HI", wptr_addr, inst=inst)
+5 -4
View File
@@ -93,11 +93,12 @@ class _System:
if data is not None: sysmem_view[:len(data)] = data
return sysmem_view, [p + i for p, sz in paddrs for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]
def pci_scan_bus(self, target_vendor:int, target_devices:list[tuple[int, list[int]]]) -> list[str]:
def pci_scan_bus(self, target_vendor:int, target_devices:list[tuple[int, list[int]]], base_class:int|None=None) -> list[str]:
result = []
for pcibus in FileIOInterface("/sys/bus/pci/devices").listdir():
vendor = int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/vendor").read(), 16)
device = int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/device").read(), 16)
if base_class is not None and int(FileIOInterface(f"/sys/bus/pci/devices/{pcibus}/class").read(), 16) >> 16 != base_class: continue
if vendor == target_vendor and any((device & mask) in devlist for mask, devlist in target_devices): result.append(pcibus)
return sorted(result)
@@ -247,9 +248,9 @@ class LNXPCIIfaceBase:
dev_impl:PCIDevImplBase
gpus:ClassVar[list[str]] = []
def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size):
def __init__(self, dev, dev_id, vendor, devices:list[tuple[int, list[int]]], bars, vram_bar, va_start, va_size, base_class:int|None=None):
if len((cls:=type(self)).gpus) == 0:
cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices))
cls.gpus = hcq_filter_visible_devices(System.pci_scan_bus(vendor, devices, base_class))
# Acquire va range to avoid collisions.
FileIOInterface.anon_mmap(va_start, va_size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE | MAP_FIXED, 0)
@@ -290,7 +291,7 @@ class LNXPCIIfaceBase:
self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
class APLPCIIfaceBase(LNXPCIIfaceBase):
def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size):
def __init__(self, dev, dev_id, vendor, devices, bars, vram_bar, va_start, va_size, base_class:int|None=None):
self.pci_dev, self.dev, self.vram_bar = APLPCIDevice(dev.__class__.__name__[:2], pcibus=f'usb4:{dev_id}', bars=bars), dev, vram_bar
assert (read_vendor:=self.pci_dev.read_config(0x00, 2)) == vendor, f"Vendor ID mismatch: expected {vendor:#x}, got {read_vendor:#x}"
def map(self, b:HCQBuffer): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
-2
View File
@@ -102,8 +102,6 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
replace_allreduce = PatternMatcher([
(UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), handle_allreduce_multirank),
(UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), handle_allreduce),
(UPat(Ops.COPY, src=(UPat(Ops.BUFFER, name="buf"), UPat(Ops.DEVICE, name="dev"))),lambda buf,dev: UOp.new_buffer(dev.arg, buf.arg, buf.dtype)
if buf.device not in {"DISK", "NPY"} and isinstance(dev.arg, tuple) and isinstance(buf.device, str) else None),
# BROADCAST: explicitly expand broadcast copies and combine with MSTACK
(UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"), UPat(Ops.DEVICE))), lambda c,x:
UOp(Ops.MSTACK, c.dtype, tuple(x.copy_to_device(d) for d in c.device)) if isinstance(c.device, tuple) and isinstance(x.device, str) else None),
+1 -1
View File
@@ -1353,7 +1353,7 @@ renderer = PatternMatcher([
(UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
(UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
(UPat(Ops.VECTORIZE, name="x"),
lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
lambda ctx,x: f"{{{','.join([ctx[y] for y in x.src])}}}" if not x.src or not all_same(x.src) else f"{{{ctx[x.src[0]]}, ...}}"),
(UPat(GroupOp.All, name="x"), lambda x: str(x)),
])
+1 -1
View File
@@ -12,7 +12,7 @@ try:
def z3_cdiv(a, b):return z3.If((a<0), z3.If(0<b, (a+(b-1))/b, (a-(b+1))/b), a/b)
def z3_xor(a,b):
if isinstance(a, z3.BoolRef): return a^b
assert a==-1 or b==-1, "xor can only be used in indexing if one of the aruments is -1"
assert a==-1 or b==-1, "xor can only be used in indexing if one of the arguments is -1"
return -a-1 if b==-1 else -b-1
z3_alu: dict[Ops, Callable] = python_alu | {Ops.MOD: lambda a,b: a-z3_cdiv(a,b)*b, Ops.IDIV: z3_cdiv, Ops.SHR: lambda a,b: a/(2**b.as_long()),
Ops.SHL: lambda a,b: a*(2**b.as_long()), Ops.AND: lambda a,b: a%(b+1) if isinstance(b, z3.ArithRef) else a&b, Ops.WHERE: z3.If, Ops.XOR: z3_xor,
+17 -15
View File
@@ -140,6 +140,8 @@ def option(s:int|None) -> int: return 0 if s is None else s+1
device_ts_diffs:dict[str, tuple[Decimal, Decimal]] = {}
def cpu_ts_diff(device:str, thread=0) -> Decimal: return device_ts_diffs.get(device, (Decimal(0),))[thread]
device_props:dict[str, dict] = {}
DevEvent = ProfileRangeEvent|ProfileGraphEntry|ProfilePointEvent
def flatten_events(profile:list[ProfileEvent]) -> Generator[tuple[Decimal, Decimal, DevEvent], None, None]:
for e in profile:
@@ -241,13 +243,11 @@ def load_counters(profile:list[ProfileEvent]) -> None:
counter_events:dict[tuple[str, int], dict] = {}
durations:dict[str, list[float]] = {}
prg_events:dict[str, ProfileProgramEvent] = {}
dev_events:dict[str, ProfileDeviceEvent] = {}
for e in profile:
if isinstance(e, (ProfilePMCEvent, ProfileSQTTEvent)): counter_events.setdefault((e.kern, e.exec_tag), {}).setdefault(type(e), []).append(e)
if isinstance(e, ProfileRangeEvent) and e.device.startswith("AMD") and e.en is not None:
durations.setdefault(str(e.name), []).append(float(e.en-e.st))
if isinstance(e, ProfileProgramEvent): prg_events[str(e.name)] = e
if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
if len(counter_events) == 0: return None
ctxs.append({"name":"All Counters", "steps":[create_step("PMC", ("/all-pmc", len(ctxs), 0), (durations, all_counters:={}))]})
run_number = {n:0 for n,_ in counter_events}
@@ -261,7 +261,7 @@ def load_counters(profile:list[ProfileEvent]) -> None:
all_counters[(name, run_number[k], k)] = pmc[0]
if (sqtt:=v.get(ProfileSQTTEvent)):
# to decode a SQTT trace, we need the raw stream, program binary and device properties
steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), ((k, tag), [*sqtt, prg_events[k], dev_events[sqtt[0].device]])))
steps.append(create_step("SQTT", ("/prg-sqtt", len(ctxs), len(steps)), ((k, tag), sqtt, prg_events[k])))
if getenv("SQTT_PARSE"):
# run our decoder on startup, we don't use this since it only works on gfx11
from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
@@ -270,11 +270,12 @@ def load_counters(profile:list[ProfileEvent]) -> None:
# ** SQTT OCC only unpacks wave start, end time and SIMD location
def unpack_sqtt(key:tuple[str, int], profile:list[ProfileEvent]) -> tuple[dict[str, list[ProfileEvent]], list[str], dict[str, dict[str, dict]]]:
def unpack_sqtt(key:tuple[str, int], data:list, p:ProfileProgramEvent) -> tuple[dict[str, list[ProfileEvent]], list[str], dict[str, dict[str, dict]]]:
# * init decoder
from extra.sqtt.roc import decode
rctx = decode(profile)
disasm = rctx.disasms[key[0]]
base = unwrap(p.base)
disasm = {addr+base:inst_disasm for addr,inst_disasm in llvm_disasm(device_props[p.device]["gfx_target_version"], unwrap(p.lib)).items()}
rctx = decode(data, {p.name:disasm})
cu_events:dict[str, list[ProfileEvent]] = {}
# * INST waves
wave_insts:dict[str, dict[str, dict]] = {}
@@ -309,6 +310,7 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
for ev in profile:
if isinstance(ev, ProfileDeviceEvent):
device_ts_diffs[ev.device] = (ev.comp_tdiff,ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff)
if ev.props is not None: device_props[ev.device] = ev.props
if (d:=ev.device.split(":")[0]) == "AMD": device_decoders[d] = load_counters
# load device specific counters
for fxn in device_decoders.values(): fxn(profile)
@@ -358,7 +360,7 @@ def amd_readelf(lib:bytes) -> list[dict]:
".group_segment_fixed_size":"LDS size", ".private_segment_fixed_size":"Scratch size"}
return [{"label":label, "value":v} for k,label in keys.items() if (v:=notes["amdhsa.kernels"][0][k]) > 0]
def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
def llvm_disasm(target:int, lib:bytes) -> dict[int, tuple[str, int]]:
from tinygrad.runtime.autogen import llvm
from tinygrad.runtime.support.elf import elf_loader
llvm.LLVMInitializeAMDGPUTargetInfo()
@@ -367,6 +369,7 @@ def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
llvm.LLVMInitializeAMDGPUDisassembler()
# pass NULL to callbacks
cbs = [ctypes.cast(0, llvm.LLVMCreateDisasmCPUFeatures.argtypes[i]) for i in {5,6}]
arch = "gfx%d%x%x" % (target // 10000, (target // 100) % 100, target % 100)
ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, *cbs)
image, sections, _ = elf_loader(lib)
text = next((sh.header for sh in sections if sh.name == ".text"), None)
@@ -392,9 +395,9 @@ def parse_branch(asm:str) -> int|None:
COND_TAKEN, COND_NOT_TAKEN, UNCOND = range(3)
cfg_colors = {COND_TAKEN: "#3f7564", COND_NOT_TAKEN: "#7a4540", UNCOND: "#3b5f7e"}
def amdgpu_cfg(lib:bytes, arch:str) -> dict:
def amdgpu_cfg(lib:bytes, target:int) -> dict:
# disassemble
pc_table = llvm_disasm(arch, lib)
pc_table = llvm_disasm(target, lib)
# get leaders
leaders:set[int] = {next(iter(pc_table))}
for pc, (asm, sz) in pc_table.items():
@@ -427,13 +430,12 @@ def get_render(i:int, j:int, fmt:str) -> dict:
if fmt == "uops": return {"src":get_stdout(lambda: print_uops(data.uops or [])), "lang":"txt"}
if fmt == "code": return {"src":data.src, "lang":"cpp"}
if fmt == "asm":
compiler = Device[data.device].compiler
disasm_str = get_stdout(lambda: compiler.disassemble(compiler.compile(data.src)))
ret:dict = {"src":disasm_str}
if data.device.startswith("AMD"):
ret:dict = {"metadata":[]}
if data.device.startswith("AMD") and data.lib is not None:
with soft_err(lambda err: ret.update(err)):
metadata = amd_readelf(lib:=compiler.compile(data.src))
ret = {"data":amdgpu_cfg(lib, getattr(compiler, "arch")), "metadata":[metadata]}
ret["data"] = amdgpu_cfg(lib:=data.lib, device_props[data.device]["gfx_target_version"])
with soft_err(lambda err: ret["metadata"].append(err)): ret["metadata"].append(amd_readelf(lib))
else: ret["src"] = get_stdout(lambda: (compiler:=Device[data.device].compiler).disassemble(compiler.compile(data.src)))
return ret
if fmt == "all-pmc":
durations, pmc = data