Compare commits

..
Author SHA1 Message Date
geohot 48dd1d6543 no long 2026-01-21 21:53:23 +09:00
geohot a06023cd36 viz slowness 2026-01-21 21:50:09 +09:00
geohot ac232bceb5 remove the device when we render 2026-01-21 21:41:29 +09:00
geohot f19fbadce4 regression test 2026-01-21 21:30:04 +09:00
geohot 8eb762d6fa add device to local, fix PCONTIG=2 2026-01-21 19:11:49 +09:00
180 changed files with 3290 additions and 26172 deletions
+1 -1
View File
@@ -11,5 +11,5 @@ runs:
git fetch origin $CURRENT_SHA
export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
export CURRENT_HEAD=$(git rev-parse HEAD)
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && CHECK_OOB=0 PYTHONPATH=. python3 process_replay.py
cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
git checkout $CURRENT_HEAD # restore to branch
+3 -3
View File
@@ -40,13 +40,13 @@ jobs:
mesa: 'true'
pydeps: 'pyyaml mako'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
- name: Regenerate autogen files
run: |
find tinygrad/runtime/autogen -type f -name "*.py" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
python3 -c "from tinygrad.runtime.autogen import opencl"
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd"
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
python3 -c "from tinygrad.runtime.autogen import llvm"
@@ -84,7 +84,7 @@ jobs:
- name: Regenerate autogen files
run: |
rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
- name: Check for differences
run: |
if ! git diff --quiet; then
+1 -1
View File
@@ -561,7 +561,7 @@ jobs:
- name: openpilot compile3 0.10.1 driving_policy
run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.10.1 dmonitoring
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
+15 -22
View File
@@ -5,7 +5,7 @@ env:
CAPTURE_PROCESS_REPLAY: 1
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PYTHONPATH: ${{ github.workspace }}
CHECK_OOB: 1
IGNORE_OOB: 0
on:
push:
@@ -38,7 +38,7 @@ jobs:
runs-on: ubuntu-22.04
timeout-minutes: 10
env:
CHECK_OOB: 0
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -241,7 +241,7 @@ jobs:
python -m mypy --lineprecision-report .
cat lineprecision.txt
- name: Run TYPED=1
run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py
run: TYPED=1 python -c "import tinygrad"
unittest:
name: Unit Tests
@@ -367,7 +367,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: gen-dataset
deps: testing
deps: testing_minimal
opencl: 'true'
- name: Generate Dataset
run: CL=1 extra/optimization/generate_dataset.sh
@@ -476,7 +476,7 @@ jobs:
runs-on: ubuntu-24.04
timeout-minutes: 15
env:
CHECK_OOB: 0
IGNORE_OOB: 1
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -650,23 +650,16 @@ jobs:
run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
- name: Run TestOps.test_add with SQTT
run: |
VIZ=-2 DEBUG=5 python3 test/test_ops.py TestOps.test_add
VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
- name: Run AMD emulated mmapeak on NULL backend
env:
AMD: 0
run: PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
testamdasm:
name: AMD ASM IDE
runs-on: ubuntu-24.04
timeout-minutes: 20
env:
AMD: 1
PYTHON_REMU: 1
MOCKGPU: 1
timeout-minutes: 10
steps:
- name: Checkout Code
uses: actions/checkout@v4
@@ -679,7 +672,7 @@ jobs:
python-version: '3.13'
- name: Verify AMD autogen is up to date
run: |
python -m extra.assembly.amd.generate
python -m extra.assembly.amd.amdxml
git diff --exit-code extra/assembly/amd/autogen/
- name: Install LLVM 21
run: |
@@ -692,16 +685,16 @@ jobs:
- name: Install rocprof-trace-decoder
run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
- name: Run RDNA3 emulator tests
run: AMD_LLVM=0 python -m pytest -n=auto extra/assembly/amd/ --durations 20
run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
- name: Run RDNA3 dtype tests
run: AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
run: AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
# TODO: run all once emulator is faster
- name: Run RDNA3 ops tests
run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
testnvidia:
strategy:
@@ -734,8 +727,6 @@ jobs:
- name: Run pytest (cuda)
# skip multitensor because it's slow
run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
- name: Run TestOps.test_add with PMA
run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -792,6 +783,8 @@ jobs:
ocelot: 'true'
llvm: 'true'
- name: Run unit tests
env:
LIBCLANG_PATH: '/opt/homebrew/opt/llvm@20/lib/libclang.dylib'
run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
- name: Run ONNX
run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-196
View File
@@ -1,196 +0,0 @@
from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
from tinygrad.uop.ops import Ops, PatternMatcher, UPat
from tinygrad.uop.symbolic import symbolic
from tinygrad.codegen import Renderer
from tinygrad.codegen.opt import Opt, OptOps
# ************************* implementation of the problem ************************
def myhash(a: Tensor) -> Tensor:
a = (a + 0x7ED55D16) + (a << 12)
a = (a ^ 0xC761C23C) ^ (a >> 19)
a = (a + 0x165667B1) + (a << 5)
a = (a + 0xD3A2646C) ^ (a << 9)
a = (a + 0xFD7046C5) + (a << 3)
a = (a ^ 0xB55A4F09) ^ (a >> 16)
return a
def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
n = values.shape[0]
if n == 1: return values[0].expand(relative_idx.shape)
mid = n // 2
left = select_with_where_tree(values[:mid], relative_idx)
right = select_with_where_tree(values[mid:], relative_idx - mid)
go_left = relative_idx < mid
return go_left.where(left, right)
def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
# All walkers start at idx=0
idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
for r in range(rounds):
level = r % (height + 1)
level_start = (1 << level) - 1
level_size = 1 << level
if level == 0:
# At root (level 0), all walkers are at idx=0
# No gather needed, just broadcast the root value
node_val = forest[0].expand(val.shape)
idx = idx * 0 # Reset to 0
elif level <= where_tree_threshold:
# Small level: use where-tree
level_values = forest[level_start : level_start + level_size]
relative_idx = (idx - level_start)
node_val = select_with_where_tree(level_values, relative_idx)
else:
# Large level: use gather
node_val = forest.gather(0, idx)
val = myhash(val ^ node_val)
idx = (idx << 1) + (1 + (val & 1))
# No wrap check needed! At round 10 (level becomes 0), we reset idx above.
return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
# ************************* renderer for VLIW machine *************************
def loop_unrolling(sink:UOp):
rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
if len(rng) == 0: return None
print(f"unrolling loop with size {rng[0].vmax+1}")
unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
return UOp.sink(*unrolled_sinks, arg=sink.arg)
global_addrs = []
vliw_prepare = PatternMatcher([
# loop unrolling (should be a part of tinygrad)
(UPat(Ops.SINK, name="sink"), loop_unrolling),
# cast is fake
(UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
# rewrites to hardcode the addresses in memory
(UPat(Ops.DEFINE_GLOBAL, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
# INDEX is just plus
(UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
])+symbolic
class VLIWRenderer(Renderer):
has_local = False # TODO: this should be the default / cleaned up
# this says this backend supports MULACC + more. decompositions uses this
code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
# this matcher runs while still in graph form
pre_matcher = vliw_prepare
def render(self, uops:list[UOp]):
# TODO: this is a minimal renderer. for low cycle count, make it good
# to get speed, you need to add VLIW packing
# to get under 1536 regs, you need to add a register allocator
# we left the fun parts to you
print(f"rendering with {len(uops)} uops")
reg, inst = 0, []
r: dict[UOp, int] = {}
for u in uops:
assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
# dumb register allocator
if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
r[u] = reg
reg += u.dtype.count
# render UOps to instructions
match u.op:
case Ops.SINK:
inst.append({"flow": [("halt",)]})
case Ops.CONST:
inst.append({"load": [("const", r[u], u.arg)]})
case Ops.GEP:
# a GEP is just an alias to a special register in the vector
r[u] = r[u.src[0]] + u.arg[0]
case Ops.VECTORIZE:
if all(s == u.src[0] for s in u.src):
# if all sources are the same, we can broadcast
inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
else:
# this is a copy into a contiguous chunk of registers
inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
case Ops.LOAD:
op = "vload" if u.dtype.count > 1 else "load"
inst.append({"load": [(op, r[u], r[u.src[0]])]})
case Ops.STORE:
op = "vstore" if u.src[1].dtype.count > 1 else "store"
inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
case Ops.MULACC:
assert u.dtype.count == 8
inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
case Ops.WHERE:
assert u.dtype.count == 8
inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
case _ if u.op in self.code_for_op:
cat = "valu" if u.dtype.count > 1 else "alu"
inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
case _:
raise NotImplementedError(f"unhandled op {u.op}")
return repr(inst)
# ************************* test and render *************************
import sys, types
PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
sys.modules["problem"] = problem = types.ModuleType("problem")
exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
if __name__ == "__main__":
batch_size = getenv("BS", 256)
height = 10
rounds = getenv("ROUNDS", 16)
# build problem
tree = problem.Tree.generate(height)
inp = problem.Input.generate(tree, batch_size, rounds)
mem = problem.build_mem_image(tree, inp)
global_addrs.extend([mem[6], mem[6], mem[4]]) # output, input, forest
# *** verify the kernel in tinygrad compared to reference ***
forest_t = Tensor(tree.values, dtype=dtypes.uint32)
val_t = Tensor(inp.values, dtype=dtypes.uint32)
if getenv("VERIFY", 1):
# verify on normal tinygrad device
with Context(PCONTIG=2):
out = tree_traversal(forest_t, val_t, height, rounds)
val_out = out.tolist()
problem.reference_kernel(tree, inp)
assert val_out == inp.values
print("verification passed")
# *** render to device ***
from tinygrad.codegen import get_program
with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
out = tree_traversal(forest_t, val_t, height, rounds)
sink = out.schedule()[-1].ast
prg = get_program(sink, VLIWRenderer())
# *** run on Machine and compare ***
# NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
src = eval(prg.src)
max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else " <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
machine.run()
print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else " <-- EVEN CLAUDE GOT 1363"))
# compare to reference
ref_mem = mem.copy()
for _ in problem.reference_kernel2(ref_mem, {}): pass
assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
print("compare passed!")
+22 -39
View File
@@ -1299,11 +1299,6 @@ def train_llama3():
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
SMALL = config["SMALL"] = getenv("SMALL", 0)
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
EVAL_SAMPLES = config["EVAL_SAMPLES"] = getenv("EVAL_SAMPLES", 5760 if not SMALL else 1024)
MAX_STEPS = config["MAX_STEPS"] = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS))
WARMUP_STEPS = config["WARMUP_STEPS"] = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
LR = config["LR"] = getenv("LR", 8e-5 * GBS / 1152)
END_LR = config["END_LR"] = getenv("END_LR", 8e-7)
EVAL_FREQ = config["EVAL_FREQ"] = getenv("EVAL_FREQ", 46080)
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
EVAL_TARGET = config["EVAL_TARGET"] = getenv("EVAL_TARGET", 5.6)
@@ -1317,10 +1312,10 @@ def train_llama3():
opt_adamw_weight_decay = 0.1
opt_gradient_clip_norm = 1.0
opt_learning_rate_warmup_steps = WARMUP_STEPS
opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
opt_base_learning_rate = LR
opt_end_learning_rate = END_LR
opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
opt_end_learning_rate = getenv("END_LR", 8e-7)
# ** init wandb **
WANDB = getenv("WANDB")
@@ -1403,7 +1398,7 @@ def train_llama3():
total_norm += p.grad.float().square().sum()
total_norm = total_norm.sqrt().contiguous()
for p in optim.params:
p.grad = (p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(p.grad.dtype)
p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
optim.step()
scheduler.step()
@@ -1441,54 +1436,42 @@ def train_llama3():
eval_dataset = None
else:
from examples.mlperf.dataloader import get_llama3_dataset
eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
eval_dataset = get_llama3_dataset(1024 if SMALL else 5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
def get_eval_iter():
if eval_dataset is None:
return fake_data(EVAL_BS, EVAL_SAMPLES)
return fake_data(EVAL_BS, 5760)
from examples.mlperf.dataloader import iterate_llama3_dataset
return iterate_llama3_dataset(eval_dataset, EVAL_BS)
num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
train_iter = get_train_iter()
iter = get_train_iter()
i, sequences_seen = resume_ckpt, 0
step_times = []
while i < MAX_STEPS:
for tokens in tqdm(iter, total=SAMPLES//GBS):
GlobalCounters.reset()
if getenv("TRAIN", 1):
st = time.perf_counter()
try: tokens = next(train_iter)
except StopIteration: break
dt = time.perf_counter()
t = time.perf_counter()
loss, lr = train_step(model, tokens)
loss = loss.float().item()
lr = lr.item()
et = time.perf_counter()
step_time = et - st
dev_time = et - dt
data_time = dt - st
if BENCHMARK: step_times.append(step_time)
i += 1
sequences_seen += tokens.shape[0]
sec = time.perf_counter()-t
if BENCHMARK: step_times.append(sec)
mem_gb = GlobalCounters.mem_used / 1e9
gflops = GlobalCounters.global_ops / 1e9 / dev_time
mfu = ((6 * num_params * SEQLEN * BS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
gflops = GlobalCounters.global_ops / 1e9 / sec
tqdm.write(
f"{i:5} {step_time:.3f} s run, {dev_time:.3f} s device, {data_time:.3f} s data, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
if (fname:=getenv("LOSS_FILE", "")):
with open(fname, "a") as f:
f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
if WANDB:
wandb.log({
"lr": lr, "train/loss": loss,
"train/step_time": step_time,
"train/dev_time": dev_time,
"train/data_time": data_time,
"train/GFLOPS": gflops,
"train/MFU": mfu,
"train/sequences_seen": sequences_seen
})
wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
tqdm.write("saving checkpoint")
@@ -1515,10 +1498,10 @@ def train_llama3():
eval_iter = get_eval_iter()
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
for j,tokens in tqdm(enumerate(eval_iter), total=5760//EVAL_BS):
eval_losses += eval_step(model, tokens).tolist()
if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
if BENCHMARK and (j+1) == min(BENCHMARK, 5760//EVAL_BS):
return
log_perplexity = Tensor(eval_losses).mean().float().item()
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
export CHECK_OOB=0
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
@@ -5,7 +5,7 @@ export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000
export BEAM=0 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -11,7 +11,7 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export CHECK_OOB=0
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=500000
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -1,32 +0,0 @@
#!/usr/bin/env bash
export PYTHONPATH="."
export DEV=${DEV:-AMD}
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEBUG=${DEBUG:-2}
export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
export ALL2ALL=${ALL2ALL:-1}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=8 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=1
export GBS=$((BS * GRADIENT_ACC_STEPS))
export MODEL="llama3"
export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEED=5760
export JITBEAM=3
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2
python3 examples/mlperf/model_train.py
@@ -1,26 +1,22 @@
#!/usr/bin/env bash
export PYTHONPATH="."
export DEV=${DEV:-AMD}
export CHECK_OOB=0
export PYTHONPATH="." AMD=1
export IGNORE_OOB=1
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
export DEBUG=${DEBUG:-0}
export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
export ALL2ALL=${ALL2ALL:-1}
export FLASH_ATTENTION=1
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DP=8 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=1
export GBS=$((BS * GRADIENT_ACC_STEPS))
export DP=8 BS=8 EVAL_BS=8
export MODEL="llama3"
export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8b"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export LR="1e-3" END_LR="1e-4" WARMUP_STEPS=1024 MAX_STEPS=1200000
export SAMPLES=$((MAX_STEPS * BS))
export SEED=5760
@@ -22,8 +22,6 @@ FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_
"rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}},
"cdna": {"DS": {152: "DS_GWS_SEMA_RELEASE_ALL", 154: "DS_GWS_SEMA_V", 156: "DS_GWS_SEMA_P"},
"VOP3P": {44: "V_MFMA_LD_SCALE_B32", 62: "V_MFMA_F32_16X16X8_XF32", 63: "V_MFMA_F32_32X32X4_XF32"}}}
# Fields missing from XML but present in hardware (format: {arch: {encoding: [(name, hi, lo), ...]}})
FIELD_FIXES = {"cdna": {"VOP3P": [("opsel_hi2", 14, 14)]}}
# Encoding suffixes to strip (variants we don't generate separate classes for)
_ENC_SUFFIXES = ("_NSA1",)
# Encoding suffix to class suffix mapping (for variants we DO generate)
@@ -126,13 +124,13 @@ def parse_xml(filename: str):
if fmt and fmt not in fmts: fmts[fmt] = 0
if otype: op_types_set.add(otype)
if op_info: types[(name, base_enum)] = op_info
# Find opcodes that only exist in a specific variant encoding (no base format version)
suffix_only_ops: dict[str, dict[str, set[int]]] = {} # {suffix: {base_fmt: {opcodes}}}
# Find opcodes that only exist in _LIT encoding (no base format version)
lit_only_ops: dict[str, set[int]] = {}
for base_fmt, opcodes in opcode_encs.items():
for opcode, encs in opcodes.items():
suffix = next((s for s in _ENC_SUFFIX_MAP.values() if all(s in e for e in encs)), None)
if suffix is not None: suffix_only_ops.setdefault(suffix, {}).setdefault(base_fmt, set()).add(opcode)
return encodings, enums, types, fmts, op_types_set, suffix_only_ops
if all("_LIT" in e for e in encs):
lit_only_ops.setdefault(base_fmt, set()).add(opcode)
return encodings, enums, types, fmts, op_types_set, lit_only_ops
# ═══════════════════════════════════════════════════════════════════════════════
# PDF parsing
@@ -222,10 +220,6 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
def write_common(all_fmts, all_op_types, path):
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
lines.append("class ReprEnum(Enum):")
lines.append(' """Enum with clean repr that roundtrips with eval()."""')
lines.append(' def __repr__(self): return f"{type(self).__name__}.{self.name}"')
lines.append("")
lines.append("class Fmt(Enum):")
for fmt in sorted(all_fmts.keys()): lines.append(f" {fmt} = auto()")
lines.append("")
@@ -238,11 +232,11 @@ def write_common(all_fmts, all_op_types, path):
with open(path, "w") as f: f.write("\n".join(lines))
def write_enum(enums, path):
lines = ["# autogenerated from AMD ISA XML - do not edit", "from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401", ""]
lines = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum", "from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401", ""]
for name, ops in sorted(enums.items()):
if not ops: continue
suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
lines.append(f"class {name}(ReprEnum):" if name in ("HWREG", "MSG") else f"class {name}Op(ReprEnum):")
lines.append(f"class {name}(Enum):" if name in ("HWREG", "MSG") else f"class {name}Op(Enum):")
aliases = []
for op, mem in sorted(ops.items()):
msuf = suffix if name != "VOP3" or op < 512 else ""
@@ -252,8 +246,8 @@ def write_enum(enums, path):
lines.append("")
with open(path, "w") as f: f.write("\n".join(lines))
def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr", "vsrc"}
def write_ins(encodings, enums, lit_only_ops, types, arch, path):
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr"}
_VARIANT_SUFFIXES = ("_LIT", "_DPP16", "_DPP8", "_SDWA_SDST", "_SDWA", "_MFMA")
def get_base_fmt(fmt):
for sfx in _VARIANT_SUFFIXES: fmt = fmt.replace(sfx, "")
@@ -273,8 +267,6 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
# GLOBAL/SCRATCH: offset is 13-bit signed [12:0], FLAT: 12-bit unsigned (XML has 12-bit for all)
if name == "offset" and base_fmt in ("GLOBAL", "SCRATCH"): return f"BitField(12, {lo})"
if base_fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
if base_fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
return f"BitField({hi}, {lo})"
@@ -313,15 +305,12 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
# Generate base classes first
for enc_name, (fields, enc_bits) in sorted(base_encodings.items()):
# Get lit-only ops for this format (these can't be used in base class)
base_lit_ops = lit_only_ops.get(enc_name, set())
all_ops = set(enums.get(enc_name, {}).keys())
# Get suffix-only ops for this format (these can't be used in base class)
base_suffix_ops = set().union(*(d.get(enc_name, set()) for d in suffix_only_ops.values()))
# Exclude SDST ops from base class (they need VOP1_SDST/VOP3_SDST/VOP3B)
base_allowed = all_ops - base_suffix_ops - sdst_opcodes.get(enc_name, set())
# RDNA3 FLAT/GLOBAL/SCRATCH share encoding bits, differentiated by seg field
# RDNA4 VFLAT/VGLOBAL/VSCRATCH have distinct encoding bits, no seg field needed
has_seg_field = any(fn == "seg" for fn, _, _ in fields)
if enc_name in ("FLAT", "VFLAT") and has_seg_field:
base_allowed = all_ops - base_lit_ops - sdst_opcodes.get(enc_name, set())
if enc_name in ("FLAT", "VFLAT"):
prefix = "V" if enc_name == "VFLAT" else ""
for cls, seg, op_enum in [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"), (f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]:
cls_ops = set(enums.get(cls, {}).keys())
@@ -331,7 +320,7 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
elif fn == "op": lines.append(f" op = EnumBitField({hi}, {lo}, {op_enum}, {fmt_allowed(op_enum, cls_ops)})")
else: lines.append(f" {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
lines.append("")
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "DPP", "SDWA"):
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "VGLOBAL", "VSCRATCH", "DPP", "SDWA"):
lines.append(f"class {enc_name}(Inst):")
for fn, hi, lo in sort_fields(fields):
if fn == "op":
@@ -347,19 +336,15 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
if base not in base_encodings: continue # skip if no base class
base_fields = {f[0] for f in base_encodings[base][0]}
extra_fields = [(fn, hi, lo) for fn, hi, lo in fields if fn not in base_fields]
# Check if this is a suffix-only variant
variant_suffix = next((sfx for sfx in _VARIANT_SUFFIXES if enc_name.endswith(sfx)), None)
is_suffix_variant = variant_suffix in suffix_only_ops
is_lit = enc_name.endswith("_LIT")
all_ops = set(enums.get(base, {}).keys())
if extra_fields or is_suffix_variant:
if extra_fields or is_lit:
lines.append(f"class {enc_name}({base}):")
op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
# _LIT classes: override op to allow all opcodes (base excludes lit-only ops)
# other classes override op to only suffix-only opcodes
if op_field and is_suffix_variant:
if op_field and is_lit:
_, hi, lo = op_field
allowed_ops = all_ops if variant_suffix == "_LIT" else suffix_only_ops[variant_suffix][base]
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', allowed_ops)})")
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', all_ops)})")
for fn, hi, lo in sort_fields(extra_fields):
lines.append(f" {fn} = {field_def(fn, hi, lo, enc_name)}")
lines.append("")
@@ -393,14 +378,14 @@ def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
for fmt, ops in sorted(enums.items()):
if fmt not in base_encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
op_to_suffix = {op:suffix for suffix,ops in suffix_only_ops.items() for op in ops.get(fmt, set())}
lit_ops = lit_only_ops.get(fmt, set())
fmt_sdst_ops = sdst_opcodes.get(fmt, set())
for op, name in sorted(ops.items()):
msuf = suffix if fmt != "VOP3" or op < 512 else ""
# Determine class: SDST variants, suffix-specific variants (e.g., _MFMA, _LIT), or base
# Determine class: SDST variants, LIT-only instructions, or base
if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
elif fmt == "VOP3" and (op in fmt_sdst_ops or op < 256): cls = "VOP3_SDST"
elif op_to_suffix.get(op): cls = f"{fmt}{op_to_suffix[op]}"
elif op in lit_ops: cls = f"{fmt}_LIT"
else: cls = fmt
lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
with open(path, "w") as f: f.write("\n".join(lines))
@@ -449,11 +434,9 @@ if __name__ == "__main__":
# First pass: parse XML for all architectures
for arch, cfg in ARCHS.items():
print(f"Parsing XML: {cfg['xml']} -> {arch}")
encodings, enums, types, fmts, op_types_set, suffix_only_ops = parse_xml(cfg["xml"])
encodings, enums, types, fmts, op_types_set, lit_only_ops = parse_xml(cfg["xml"])
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
for fmt, fields in FIELD_FIXES.get(arch, {}).items():
if fmt in encodings: encodings[fmt] = (encodings[fmt][0] + fields, encodings[fmt][1])
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "suffix_only_ops": suffix_only_ops}
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "lit_only_ops": lit_only_ops}
for fmt, bits in fmts.items():
assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
all_fmts[fmt] = bits
@@ -466,7 +449,7 @@ if __name__ == "__main__":
for arch, data in arch_data.items():
base = pathlib.Path(__file__).parent / "autogen" / arch
write_enum(data["enums"], base / "enum.py")
write_ins(data["encodings"], data["enums"], data["suffix_only_ops"], data["types"], arch, base / "ins.py")
write_ins(data["encodings"], data["enums"], data["lit_only_ops"], data["types"], arch, base / "ins.py")
write_operands(data["types"], data["enums"], arch, base / "operands.py")
print(f" {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
# Second pass: parse PDFs and write pcode
+22 -21
View File
@@ -1,7 +1,8 @@
# autogenerated from AMD ISA XML - do not edit
from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(ReprEnum):
class DSOp(Enum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -132,7 +133,7 @@ class DSOp(ReprEnum):
DS_READ_B96 = 254
DS_READ_B128 = 255
class FLATOp(ReprEnum):
class FLATOp(Enum):
FLAT_LOAD_UBYTE = 16
FLAT_LOAD_SBYTE = 17
FLAT_LOAD_USHORT = 18
@@ -188,7 +189,7 @@ class FLATOp(ReprEnum):
FLAT_ATOMIC_INC_X2 = 107
FLAT_ATOMIC_DEC_X2 = 108
class GLOBALOp(ReprEnum):
class GLOBALOp(Enum):
GLOBAL_LOAD_UBYTE = 16
GLOBAL_LOAD_SBYTE = 17
GLOBAL_LOAD_USHORT = 18
@@ -251,7 +252,7 @@ class GLOBALOp(ReprEnum):
GLOBAL_LOAD_LDS_DWORDX4 = 125
GLOBAL_LOAD_LDS_DWORDX3 = 126
class HWREG(ReprEnum):
class HWREG(Enum):
HW_REG_MODE = 1
HW_REG_STATUS = 2
HW_REG_TRAPSTS = 3
@@ -277,7 +278,7 @@ class HWREG(ReprEnum):
HW_REG_SQ_PERF_SNAPSHOT_PC_LO = 23
HW_REG_SQ_PERF_SNAPSHOT_PC_HI = 24
class MTBUFOp(ReprEnum):
class MTBUFOp(Enum):
TBUFFER_LOAD_FORMAT_X = 0
TBUFFER_LOAD_FORMAT_XY = 1
TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -295,7 +296,7 @@ class MTBUFOp(ReprEnum):
TBUFFER_STORE_FORMAT_D16_XYZ = 14
TBUFFER_STORE_FORMAT_D16_XYZW = 15
class MUBUFOp(ReprEnum):
class MUBUFOp(Enum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -371,7 +372,7 @@ class MUBUFOp(ReprEnum):
BUFFER_ATOMIC_INC_X2 = 107
BUFFER_ATOMIC_DEC_X2 = 108
class SCRATCHOp(ReprEnum):
class SCRATCHOp(Enum):
SCRATCH_LOAD_UBYTE = 16
SCRATCH_LOAD_SBYTE = 17
SCRATCH_LOAD_USHORT = 18
@@ -400,7 +401,7 @@ class SCRATCHOp(ReprEnum):
SCRATCH_LOAD_LDS_SSHORT = 41
SCRATCH_LOAD_LDS_DWORD = 42
class SMEMOp(ReprEnum):
class SMEMOp(Enum):
S_LOAD_DWORD = 0
S_LOAD_DWORDX2 = 1
S_LOAD_DWORDX4 = 2
@@ -486,7 +487,7 @@ class SMEMOp(ReprEnum):
S_ATOMIC_INC_X2 = 171
S_ATOMIC_DEC_X2 = 172
class SOP1Op(ReprEnum):
class SOP1Op(Enum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -542,7 +543,7 @@ class SOP1Op(ReprEnum):
S_ANDN2_WREXEC_B64 = 54
S_BITREPLICATE_B64_B32 = 55
class SOP2Op(ReprEnum):
class SOP2Op(Enum):
S_ADD_U32 = 0
S_SUB_U32 = 1
S_ADD_I32 = 2
@@ -597,7 +598,7 @@ class SOP2Op(ReprEnum):
S_PACK_LH_B32_B16 = 51
S_PACK_HH_B32_B16 = 52
class SOPCOp(ReprEnum):
class SOPCOp(Enum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -619,7 +620,7 @@ class SOPCOp(ReprEnum):
S_CMP_EQ_U64 = 18
S_CMP_LG_U64 = 19
class SOPKOp(ReprEnum):
class SOPKOp(Enum):
S_MOVK_I32 = 0
S_CMOVK_I32 = 1
S_CMPK_EQ_I32 = 2
@@ -642,7 +643,7 @@ class SOPKOp(ReprEnum):
S_SETREG_IMM32_B32 = 20
S_CALL_B64 = 21
class SOPPOp(ReprEnum):
class SOPPOp(Enum):
S_NOP = 0
S_ENDPGM = 1
S_BRANCH = 2
@@ -676,7 +677,7 @@ class SOPPOp(ReprEnum):
S_ENDPGM_ORDERED_PS_DONE = 30
S_SET_VALU_COEXEC_MODE = 31
class VOP1Op(ReprEnum):
class VOP1Op(Enum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -854,7 +855,7 @@ class VOP1Op(ReprEnum):
V_PERMLANE32_SWAP_B32 = V_PERMLANE32_SWAP_B32_E32
V_CVT_F32_BF16 = V_CVT_F32_BF16_E32
class VOP2Op(ReprEnum):
class VOP2Op(Enum):
V_CNDMASK_B32_E32 = 0
V_ADD_F32_E32 = 1
V_SUB_F32_E32 = 2
@@ -980,7 +981,7 @@ class VOP2Op(ReprEnum):
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
V_XNOR_B32 = V_XNOR_B32_E32
class VOP3Op(ReprEnum):
class VOP3Op(Enum):
V_CMP_CLASS_F32_E64 = 16
V_CMPX_CLASS_F32_E64 = 17
V_CMP_CLASS_F64_E64 = 18
@@ -1878,7 +1879,7 @@ class VOP3Op(ReprEnum):
V_ADD_LSHL_U32 = V_ADD_LSHL_U32_E64
V_ADD3_U32 = V_ADD3_U32_E64
class VOP3POp(ReprEnum):
class VOP3POp(Enum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1987,11 +1988,11 @@ class VOP3POp(ReprEnum):
V_SMFMAC_F32_32X32X32_FP8_BF8 = 126
V_SMFMAC_F32_32X32X32_FP8_FP8 = 127
class VOP3PX2Op(ReprEnum):
class VOP3PX2Op(Enum):
V_MFMA_SCALE_F32_16X16X128_F8F6F4 = 45
V_MFMA_SCALE_F32_32X32X64_F8F6F4 = 46
class VOP3SDOp(ReprEnum):
class VOP3SDOp(Enum):
V_ADD_CO_U32 = 281
V_SUB_CO_U32 = 282
V_SUBREV_CO_U32 = 283
@@ -2003,7 +2004,7 @@ class VOP3SDOp(ReprEnum):
V_MAD_U64_U32 = 488
V_MAD_I64_I32 = 489
class VOPCOp(ReprEnum):
class VOPCOp(Enum):
V_CMP_CLASS_F32_E32 = 16
V_CMPX_CLASS_F32_E32 = 17
V_CMP_CLASS_F64_E32 = 18
+69 -71
View File
@@ -38,7 +38,7 @@ class GLOBAL(Inst):
addr = VGPRField(39, 32)
data = VGPRField(47, 40)
saddr = SGPRField(54, 48, default=NULL)
offset = BitField(12, 0)
offset = BitField(11, 0)
seg = FixedBitField(15, 14, 2)
acc = BitField(55, 55)
sve = BitField(13, 13)
@@ -53,7 +53,7 @@ class SCRATCH(Inst):
addr = VGPRField(39, 32)
data = VGPRField(47, 40)
saddr = SGPRField(54, 48, default=NULL)
offset = BitField(12, 0)
offset = BitField(11, 0)
seg = FixedBitField(15, 14, 1)
acc = BitField(55, 55)
sve = BitField(13, 13)
@@ -164,7 +164,7 @@ class VOP3(Inst):
class VOP3P(Inst):
encoding = FixedBitField(31, 23, 0b110100111)
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE})
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_PK_MAD_I16, VOP3POp.V_PK_MUL_LO_U16, VOP3POp.V_PK_ADD_I16, VOP3POp.V_PK_SUB_I16, VOP3POp.V_PK_LSHLREV_B16, VOP3POp.V_PK_LSHRREV_B16, VOP3POp.V_PK_ASHRREV_I16, VOP3POp.V_PK_MAX_I16, VOP3POp.V_PK_MIN_I16, VOP3POp.V_PK_MAD_U16, VOP3POp.V_PK_ADD_U16, VOP3POp.V_PK_SUB_U16, VOP3POp.V_PK_MAX_U16, VOP3POp.V_PK_MIN_U16, VOP3POp.V_PK_FMA_F16, VOP3POp.V_PK_ADD_F16, VOP3POp.V_PK_MUL_F16, VOP3POp.V_PK_MIN_F16, VOP3POp.V_PK_MAX_F16, VOP3POp.V_DOT2_F32_BF16, VOP3POp.V_PK_MINIMUM3_F16, VOP3POp.V_PK_MAXIMUM3_F16, VOP3POp.V_MAD_MIX_F32, VOP3POp.V_MAD_MIXLO_F16, VOP3POp.V_MAD_MIXHI_F16, VOP3POp.V_DOT2_F32_F16, VOP3POp.V_DOT2_I32_I16, VOP3POp.V_DOT2_U32_U16, VOP3POp.V_DOT4_I32_I8, VOP3POp.V_DOT4_U32_U8, VOP3POp.V_DOT8_I32_I4, VOP3POp.V_DOT8_U32_U4, VOP3POp.V_MFMA_LD_SCALE_B32, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_PK_FMA_F32, VOP3POp.V_PK_MUL_F32, VOP3POp.V_PK_ADD_F32, VOP3POp.V_PK_MOV_B32, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X8_XF32, VOP3POp.V_MFMA_F32_32X32X4_XF32, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_ACCVGPR_READ, VOP3POp.V_ACCVGPR_WRITE, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
vdst = VGPRField(7, 0)
src0 = SrcField(40, 32)
src1 = SrcField(49, 41)
@@ -174,7 +174,6 @@ class VOP3P(Inst):
clmp = BitField(15, 15)
opsel = BitField(13, 11)
opsel_hi = BitField(60, 59, default=3)
opsel_hi2 = BitField(14, 14, default=1)
class VOP3PX2(Inst):
encoding = FixedBitField(95, 87, 0b110100111)
@@ -310,7 +309,6 @@ class VOP2_SDWA_SDST(VOP2):
s1 = BitField(63, 63)
class VOP3P_MFMA(VOP3P):
op = EnumBitField(22, 16, VOP3POp, {VOP3POp.V_MFMA_F32_16X16X128_F8F6F4, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4, VOP3POp.V_MFMA_F32_16X16X32_BF16, VOP3POp.V_MFMA_I32_16X16X64_I8, VOP3POp.V_MFMA_F32_32X32X16_BF16, VOP3POp.V_MFMA_I32_32X32X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_BF16, VOP3POp.V_SMFMAC_I32_16X16X128_I8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X1_2B_F32, VOP3POp.V_MFMA_F32_16X16X1_4B_F32, VOP3POp.V_MFMA_F32_4X4X1_16B_F32, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X2_F32, VOP3POp.V_MFMA_F32_16X16X4_F32, VOP3POp.V_SMFMAC_F32_32X32X32_BF16, VOP3POp.V_SMFMAC_I32_32X32X64_I8, VOP3POp.V_MFMA_F32_32X32X4_2B_F16, VOP3POp.V_MFMA_F32_16X16X4_4B_F16, VOP3POp.V_MFMA_F32_4X4X4_16B_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X8_F16, VOP3POp.V_MFMA_F32_16X16X16_F16, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8, VOP3POp.V_MFMA_I32_32X32X4_2B_I8, VOP3POp.V_MFMA_I32_16X16X4_4B_I8, VOP3POp.V_MFMA_I32_4X4X4_16B_I8, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8, VOP3POp.V_MFMA_F32_16X16X32_F16, VOP3POp.V_MFMA_F32_32X32X16_F16, VOP3POp.V_MFMA_I32_32X32X16_I8, VOP3POp.V_MFMA_I32_16X16X32_I8, VOP3POp.V_SMFMAC_F32_16X16X64_F16, VOP3POp.V_SMFMAC_F32_32X32X32_F16, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16, VOP3POp.V_MFMA_F32_32X32X8_BF16, VOP3POp.V_MFMA_F32_16X16X16_BF16, VOP3POp.V_SMFMAC_F32_16X16X32_F16, VOP3POp.V_SMFMAC_F32_32X32X16_F16, VOP3POp.V_SMFMAC_F32_16X16X32_BF16, VOP3POp.V_SMFMAC_F32_32X32X16_BF16, VOP3POp.V_SMFMAC_I32_16X16X64_I8, VOP3POp.V_SMFMAC_I32_32X32X32_I8, VOP3POp.V_MFMA_F64_16X16X4_F64, VOP3POp.V_MFMA_F64_4X4X4_4B_F64, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8})
cbsz = BitField(10, 8)
abid = BitField(14, 11)
acc_cd = BitField(15, 15)
@@ -1649,80 +1647,80 @@ v_dot4_u32_u8 = functools.partial(VOP3P, VOP3POp.V_DOT4_U32_U8)
v_dot8_i32_i4 = functools.partial(VOP3P, VOP3POp.V_DOT8_I32_I4)
v_dot8_u32_u4 = functools.partial(VOP3P, VOP3POp.V_DOT8_U32_U4)
v_mfma_ld_scale_b32 = functools.partial(VOP3P, VOP3POp.V_MFMA_LD_SCALE_B32)
v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
v_mfma_f32_16x16x128_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X128_F8F6F4)
v_mfma_f32_32x32x64_f8f6f4 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X64_F8F6F4)
v_pk_fma_f32 = functools.partial(VOP3P, VOP3POp.V_PK_FMA_F32)
v_pk_mul_f32 = functools.partial(VOP3P, VOP3POp.V_PK_MUL_F32)
v_pk_add_f32 = functools.partial(VOP3P, VOP3POp.V_PK_ADD_F32)
v_pk_mov_b32 = functools.partial(VOP3P, VOP3POp.V_PK_MOV_B32)
v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF16)
v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X64_I8)
v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF16)
v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X32_I8)
v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
v_mfma_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF16)
v_mfma_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X64_I8)
v_mfma_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF16)
v_mfma_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X32_I8)
v_smfmac_f32_16x16x64_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF16)
v_smfmac_i32_16x16x128_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X128_I8)
v_smfmac_f32_16x16x128_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_BF8)
v_smfmac_f32_16x16x128_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_BF8_FP8)
v_smfmac_f32_16x16x128_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_BF8)
v_mfma_f32_16x16x8_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X8_XF32)
v_mfma_f32_32x32x4_xf32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_XF32)
v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X2_F32)
v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_F32)
v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_F16)
v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_F16)
v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_F16)
v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_F16)
v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_32X32X16_I8)
v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_I32_16X16X32_I8)
v_mfma_f32_32x32x1_2b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X1_2B_F32)
v_mfma_f32_16x16x1_4b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X1_4B_F32)
v_mfma_f32_4x4x1_16b_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X1_16B_F32)
v_smfmac_f32_16x16x128_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X128_FP8_FP8)
v_mfma_f32_32x32x2_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X2_F32)
v_mfma_f32_16x16x4_f32 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_F32)
v_smfmac_f32_32x32x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF16)
v_smfmac_i32_32x32x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X64_I8)
v_mfma_f32_32x32x4_2b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_F16)
v_mfma_f32_16x16x4_4b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_F16)
v_mfma_f32_4x4x4_16b_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_F16)
v_smfmac_f32_32x32x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_BF8)
v_mfma_f32_32x32x8_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_F16)
v_mfma_f32_16x16x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_F16)
v_smfmac_f32_32x32x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_BF8_FP8)
v_smfmac_f32_32x32x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_BF8)
v_mfma_i32_32x32x4_2b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X4_2B_I8)
v_mfma_i32_16x16x4_4b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X4_4B_I8)
v_mfma_i32_4x4x4_16b_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_4X4X4_16B_I8)
v_smfmac_f32_32x32x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X64_FP8_FP8)
v_mfma_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_F16)
v_mfma_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_F16)
v_mfma_i32_32x32x16_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_32X32X16_I8)
v_mfma_i32_16x16x32_i8 = functools.partial(VOP3P, VOP3POp.V_MFMA_I32_16X16X32_I8)
v_accvgpr_read = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_READ)
v_accvgpr_write = functools.partial(VOP3P, VOP3POp.V_ACCVGPR_WRITE)
v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X8_BF16)
v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X16_BF16)
v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_16X16X4_F64)
v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P_MFMA, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
v_smfmac_f32_16x16x64_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_F16)
v_smfmac_f32_32x32x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_F16)
v_mfma_f32_32x32x4_2b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X4_2B_BF16)
v_mfma_f32_16x16x4_4b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X4_4B_BF16)
v_mfma_f32_4x4x4_16b_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_4X4X4_16B_BF16)
v_mfma_f32_32x32x8_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X8_BF16)
v_mfma_f32_16x16x16_bf16 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X16_BF16)
v_smfmac_f32_16x16x32_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_F16)
v_smfmac_f32_32x32x16_f16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_F16)
v_smfmac_f32_16x16x32_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X32_BF16)
v_smfmac_f32_32x32x16_bf16 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X16_BF16)
v_smfmac_i32_16x16x64_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_16X16X64_I8)
v_smfmac_i32_32x32x32_i8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_I32_32X32X32_I8)
v_mfma_f64_16x16x4_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_16X16X4_F64)
v_mfma_f64_4x4x4_4b_f64 = functools.partial(VOP3P, VOP3POp.V_MFMA_F64_4X4X4_4B_F64)
v_mfma_f32_16x16x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_BF8)
v_mfma_f32_16x16x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_BF8_FP8)
v_mfma_f32_16x16x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_BF8)
v_mfma_f32_16x16x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_16X16X32_FP8_FP8)
v_mfma_f32_32x32x16_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_BF8)
v_mfma_f32_32x32x16_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_BF8_FP8)
v_mfma_f32_32x32x16_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_BF8)
v_mfma_f32_32x32x16_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_MFMA_F32_32X32X16_FP8_FP8)
v_smfmac_f32_16x16x64_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_BF8)
v_smfmac_f32_16x16x64_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_BF8_FP8)
v_smfmac_f32_16x16x64_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_BF8)
v_smfmac_f32_16x16x64_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_16X16X64_FP8_FP8)
v_smfmac_f32_32x32x32_bf8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_BF8)
v_smfmac_f32_32x32x32_bf8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_BF8_FP8)
v_smfmac_f32_32x32x32_fp8_bf8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_BF8)
v_smfmac_f32_32x32x32_fp8_fp8 = functools.partial(VOP3P, VOP3POp.V_SMFMAC_F32_32X32X32_FP8_FP8)
v_mfma_scale_f32_16x16x128_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_16X16X128_F8F6F4)
v_mfma_scale_f32_32x32x64_f8f6f4 = functools.partial(VOP3PX2, VOP3PX2Op.V_MFMA_SCALE_F32_32X32X64_F8F6F4)
v_add_co_u32 = functools.partial(VOP3SD, VOP3SDOp.V_ADD_CO_U32)
-4
View File
@@ -1,10 +1,6 @@
# autogenerated from AMD ISA XML - do not edit
from enum import Enum, auto
class ReprEnum(Enum):
"""Enum with clean repr that roundtrips with eval()."""
def __repr__(self): return f"{type(self).__name__}.{self.name}"
class Fmt(Enum):
FMT_ANY = auto()
FMT_BUF = auto()
+27 -26
View File
@@ -1,7 +1,8 @@
# autogenerated from AMD ISA XML - do not edit
from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(ReprEnum):
class DSOp(Enum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -129,10 +130,10 @@ class DSOp(ReprEnum):
DS_LOAD_B96 = 254
DS_LOAD_B128 = 255
class EXPOp(ReprEnum):
class EXPOp(Enum):
EXP = 0
class FLATOp(ReprEnum):
class FLATOp(Enum):
FLAT_LOAD_U8 = 16
FLAT_LOAD_I8 = 17
FLAT_LOAD_U16 = 18
@@ -190,7 +191,7 @@ class FLATOp(ReprEnum):
FLAT_ATOMIC_MAX_F32 = 82
FLAT_ATOMIC_ADD_F32 = 86
class GLOBALOp(ReprEnum):
class GLOBALOp(Enum):
GLOBAL_LOAD_U8 = 16
GLOBAL_LOAD_I8 = 17
GLOBAL_LOAD_U16 = 18
@@ -253,7 +254,7 @@ class GLOBALOp(ReprEnum):
GLOBAL_ATOMIC_MAX_F32 = 82
GLOBAL_ATOMIC_ADD_F32 = 86
class HWREG(ReprEnum):
class HWREG(Enum):
HW_REG_MODE = 1
HW_REG_STATUS = 2
HW_REG_TRAPSTS = 3
@@ -279,11 +280,11 @@ class HWREG(ReprEnum):
HW_REG_IB_STS2 = 28
HW_REG_SHADER_CYCLES = 29
class LDSDIROp(ReprEnum):
class LDSDIROp(Enum):
LDS_PARAM_LOAD = 0
LDS_DIRECT_LOAD = 1
class MIMGOp(ReprEnum):
class MIMGOp(Enum):
IMAGE_LOAD = 0
IMAGE_LOAD_MIP = 1
IMAGE_LOAD_PCK = 2
@@ -369,7 +370,7 @@ class MIMGOp(ReprEnum):
IMAGE_GATHER4_C_B_CL = 101
IMAGE_GATHER4H = 144
class MSG(ReprEnum):
class MSG(Enum):
MSG_RTN_GET_DOORBELL = 128
MSG_RTN_GET_DDID = 129
MSG_RTN_GET_TMA = 130
@@ -379,7 +380,7 @@ class MSG(ReprEnum):
MSG_RTN_GET_TBA_TO_PC = 134
MSG_RTN_ILLEGAL_MSG = 255
class MTBUFOp(ReprEnum):
class MTBUFOp(Enum):
TBUFFER_LOAD_FORMAT_X = 0
TBUFFER_LOAD_FORMAT_XY = 1
TBUFFER_LOAD_FORMAT_XYZ = 2
@@ -397,7 +398,7 @@ class MTBUFOp(ReprEnum):
TBUFFER_STORE_D16_FORMAT_XYZ = 14
TBUFFER_STORE_D16_FORMAT_XYZW = 15
class MUBUFOp(ReprEnum):
class MUBUFOp(Enum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -478,7 +479,7 @@ class MUBUFOp(ReprEnum):
BUFFER_ATOMIC_MAX_F32 = 82
BUFFER_ATOMIC_ADD_F32 = 86
class SCRATCHOp(ReprEnum):
class SCRATCHOp(Enum):
SCRATCH_LOAD_U8 = 16
SCRATCH_LOAD_I8 = 17
SCRATCH_LOAD_U16 = 18
@@ -507,7 +508,7 @@ class SCRATCHOp(ReprEnum):
SCRATCH_LOAD_LDS_I16 = 48
SCRATCH_LOAD_LDS_B32 = 49
class SMEMOp(ReprEnum):
class SMEMOp(Enum):
S_LOAD_B32 = 0
S_LOAD_B64 = 1
S_LOAD_B128 = 2
@@ -523,7 +524,7 @@ class SMEMOp(ReprEnum):
S_ATC_PROBE = 34
S_ATC_PROBE_BUFFER = 35
class SOP1Op(ReprEnum):
class SOP1Op(Enum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -605,7 +606,7 @@ class SOP1Op(ReprEnum):
S_TRUNC_F16 = 109
S_RNDNE_F16 = 110
class SOP2Op(ReprEnum):
class SOP2Op(Enum):
S_ADD_U32 = 0
S_SUB_U32 = 1
S_ADD_I32 = 2
@@ -674,7 +675,7 @@ class SOP2Op(ReprEnum):
S_MUL_F16 = 77
S_FMAC_F16 = 78
class SOPCOp(ReprEnum):
class SOPCOp(Enum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -722,7 +723,7 @@ class SOPCOp(ReprEnum):
S_CMP_NEQ_F16 = 93
S_CMP_NLT_F16 = 94
class SOPKOp(ReprEnum):
class SOPKOp(Enum):
S_MOVK_I32 = 0
S_VERSION = 1
S_CMOVK_I32 = 2
@@ -751,7 +752,7 @@ class SOPKOp(ReprEnum):
S_WAITCNT_EXPCNT = 26
S_WAITCNT_LGKMCNT = 27
class SOPPOp(ReprEnum):
class SOPPOp(Enum):
S_NOP = 0
S_SETKILL = 1
S_SETHALT = 2
@@ -792,7 +793,7 @@ class SOPPOp(ReprEnum):
S_ICACHE_INV = 60
S_BARRIER = 61
class VINTERPOp(ReprEnum):
class VINTERPOp(Enum):
V_INTERP_P10_F32 = 0
V_INTERP_P2_F32 = 1
V_INTERP_P10_F16_F32 = 2
@@ -800,7 +801,7 @@ class VINTERPOp(ReprEnum):
V_INTERP_P10_RTZ_F16_F32 = 4
V_INTERP_P2_RTZ_F16_F32 = 5
class VOP1Op(ReprEnum):
class VOP1Op(Enum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -974,7 +975,7 @@ class VOP1Op(ReprEnum):
V_CVT_I32_I16 = V_CVT_I32_I16_E32
V_CVT_U32_U16 = V_CVT_U32_U16_E32
class VOP2Op(ReprEnum):
class VOP2Op(Enum):
V_CNDMASK_B32_E32 = 1
V_DOT2ACC_F32_F16_E32 = 2
V_ADD_F32_E32 = 3
@@ -1068,7 +1069,7 @@ class VOP2Op(ReprEnum):
V_LDEXP_F16 = V_LDEXP_F16_E32
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
class VOP3Op(ReprEnum):
class VOP3Op(Enum):
V_CMP_F_F16_E64 = 0
V_CMP_LT_F16_E64 = 1
V_CMP_EQ_F16_E64 = 2
@@ -1808,7 +1809,7 @@ class VOP3Op(ReprEnum):
V_CVT_I32_I16 = V_CVT_I32_I16_E64
V_CVT_U32_U16 = V_CVT_U32_U16_E64
class VOP3POp(ReprEnum):
class VOP3POp(Enum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1844,7 +1845,7 @@ class VOP3POp(ReprEnum):
V_WMMA_I32_16X16X16_IU8 = 68
V_WMMA_I32_16X16X16_IU4 = 69
class VOP3SDOp(ReprEnum):
class VOP3SDOp(Enum):
V_ADD_CO_CI_U32 = 288
V_SUB_CO_CI_U32 = 289
V_SUBREV_CO_CI_U32 = 290
@@ -1856,7 +1857,7 @@ class VOP3SDOp(ReprEnum):
V_SUB_CO_U32 = 769
V_SUBREV_CO_U32 = 770
class VOPCOp(ReprEnum):
class VOPCOp(Enum):
V_CMP_F_F16_E32 = 0
V_CMP_LT_F16_E32 = 1
V_CMP_EQ_F16_E32 = 2
@@ -2238,7 +2239,7 @@ class VOPCOp(ReprEnum):
V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
class VOPDOp(ReprEnum):
class VOPDOp(Enum):
V_DUAL_FMAC_F32 = 0
V_DUAL_FMAAK_F32 = 1
V_DUAL_FMAMK_F32 = 2
+27 -26
View File
@@ -1,7 +1,8 @@
# autogenerated from AMD ISA XML - do not edit
from extra.assembly.amd.autogen.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401
from enum import Enum
from extra.assembly.amd.autogen.common import Fmt, FMT_BITS, OpType # noqa: F401
class DSOp(ReprEnum):
class DSOp(Enum):
DS_ADD_U32 = 0
DS_SUB_U32 = 1
DS_RSUB_U32 = 2
@@ -126,7 +127,7 @@ class DSOp(ReprEnum):
DS_LOAD_B96 = 254
DS_LOAD_B128 = 255
class HWREG(ReprEnum):
class HWREG(Enum):
HW_REG_WAVE_MODE = 1
HW_REG_WAVE_STATUS = 2
HW_REG_WAVE_STATE_PRIV = 4
@@ -147,7 +148,7 @@ class HWREG(ReprEnum):
HW_REG_SHADER_CYCLES_LO = 29
HW_REG_SHADER_CYCLES_HI = 30
class MSG(ReprEnum):
class MSG(Enum):
MSG_RTN_GET_DOORBELL = 128
MSG_RTN_GET_DDID = 129
MSG_RTN_GET_TMA = 130
@@ -158,7 +159,7 @@ class MSG(ReprEnum):
MSG_RTN_GET_SE_HW_ID = 135
MSG_RTN_ILLEGAL_MSG = 255
class SMEMOp(ReprEnum):
class SMEMOp(Enum):
S_LOAD_B32 = 0
S_LOAD_B64 = 1
S_LOAD_B128 = 2
@@ -188,7 +189,7 @@ class SMEMOp(ReprEnum):
S_BUFFER_PREFETCH_DATA = 39
S_PREFETCH_DATA_PC_REL = 40
class SOP1Op(ReprEnum):
class SOP1Op(Enum):
S_MOV_B32 = 0
S_MOV_B64 = 1
S_CMOV_B32 = 2
@@ -277,7 +278,7 @@ class SOP1Op(ReprEnum):
S_TRUNC_F16 = 109
S_RNDNE_F16 = 110
class SOP2Op(ReprEnum):
class SOP2Op(Enum):
S_ADD_CO_U32 = 0
S_SUB_CO_U32 = 1
S_ADD_CO_I32 = 2
@@ -353,7 +354,7 @@ class SOP2Op(ReprEnum):
S_SUB_NC_U64 = 84
S_MUL_U64 = 85
class SOPCOp(ReprEnum):
class SOPCOp(Enum):
S_CMP_EQ_I32 = 0
S_CMP_LG_I32 = 1
S_CMP_GT_I32 = 2
@@ -401,7 +402,7 @@ class SOPCOp(ReprEnum):
S_CMP_NEQ_F16 = 93
S_CMP_NLT_F16 = 94
class SOPKOp(ReprEnum):
class SOPKOp(Enum):
S_MOVK_I32 = 0
S_VERSION = 1
S_CMOVK_I32 = 2
@@ -412,7 +413,7 @@ class SOPKOp(ReprEnum):
S_SETREG_IMM32_B32 = 19
S_CALL_B64 = 20
class SOPPOp(ReprEnum):
class SOPPOp(Enum):
S_NOP = 0
S_SETKILL = 1
S_SETHALT = 2
@@ -457,7 +458,7 @@ class SOPPOp(ReprEnum):
S_WAIT_LOADCNT_DSCNT = 72
S_WAIT_STORECNT_DSCNT = 73
class VBUFFEROp(ReprEnum):
class VBUFFEROp(Enum):
BUFFER_LOAD_FORMAT_X = 0
BUFFER_LOAD_FORMAT_XY = 1
BUFFER_LOAD_FORMAT_XYZ = 2
@@ -548,14 +549,14 @@ class VBUFFEROp(ReprEnum):
TBUFFER_STORE_D16_FORMAT_XYZ = 142
TBUFFER_STORE_D16_FORMAT_XYZW = 143
class VDSDIROp(ReprEnum):
class VDSDIROp(Enum):
DS_PARAM_LOAD = 0
DS_DIRECT_LOAD = 1
class VEXPORTOp(ReprEnum):
class VEXPORTOp(Enum):
EXPORT = 0
class VFLATOp(ReprEnum):
class VFLATOp(Enum):
FLAT_LOAD_U8 = 16
FLAT_LOAD_I8 = 17
FLAT_LOAD_U16 = 18
@@ -614,7 +615,7 @@ class VFLATOp(ReprEnum):
FLAT_ATOMIC_PK_ADD_F16 = 89
FLAT_ATOMIC_PK_ADD_BF16 = 90
class VGLOBALOp(ReprEnum):
class VGLOBALOp(Enum):
GLOBAL_LOAD_U8 = 16
GLOBAL_LOAD_I8 = 17
GLOBAL_LOAD_U16 = 18
@@ -681,7 +682,7 @@ class VGLOBALOp(ReprEnum):
GLOBAL_ATOMIC_PK_ADD_BF16 = 90
GLOBAL_ATOMIC_ORDERED_ADD_B64 = 115
class VIMAGEOp(ReprEnum):
class VIMAGEOp(Enum):
IMAGE_LOAD = 0
IMAGE_LOAD_MIP = 1
IMAGE_LOAD_PCK = 2
@@ -716,7 +717,7 @@ class VIMAGEOp(ReprEnum):
IMAGE_ATOMIC_PK_ADD_F16 = 134
IMAGE_ATOMIC_PK_ADD_BF16 = 135
class VINTERPOp(ReprEnum):
class VINTERPOp(Enum):
V_INTERP_P10_F32 = 0
V_INTERP_P2_F32 = 1
V_INTERP_P10_F16_F32 = 2
@@ -724,7 +725,7 @@ class VINTERPOp(ReprEnum):
V_INTERP_P10_RTZ_F16_F32 = 4
V_INTERP_P2_RTZ_F16_F32 = 5
class VOP1Op(ReprEnum):
class VOP1Op(Enum):
V_NOP_E32 = 0
V_MOV_B32_E32 = 1
V_READFIRSTLANE_B32_E32 = 2
@@ -906,7 +907,7 @@ class VOP1Op(ReprEnum):
V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E32
V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E32
class VOP2Op(ReprEnum):
class VOP2Op(Enum):
V_CNDMASK_B32_E32 = 1
V_ADD_F64_E32 = 2
V_ADD_F32_E32 = 3
@@ -1006,7 +1007,7 @@ class VOP2Op(ReprEnum):
V_LDEXP_F16 = V_LDEXP_F16_E32
V_PK_FMAC_F16 = V_PK_FMAC_F16_E32
class VOP3Op(ReprEnum):
class VOP3Op(Enum):
V_CMP_LT_F16_E64 = 1
V_CMP_EQ_F16_E64 = 2
V_CMP_LE_F16_E64 = 3
@@ -1731,7 +1732,7 @@ class VOP3Op(ReprEnum):
V_CVT_PK_F32_FP8 = V_CVT_PK_F32_FP8_E64
V_CVT_PK_F32_BF8 = V_CVT_PK_F32_BF8_E64
class VOP3POp(ReprEnum):
class VOP3POp(Enum):
V_PK_MAD_I16 = 0
V_PK_MUL_LO_U16 = 1
V_PK_ADD_I16 = 2
@@ -1789,7 +1790,7 @@ class VOP3POp(ReprEnum):
V_SWMMAC_F32_16X16X32_BF8_FP8 = 89
V_SWMMAC_F32_16X16X32_BF8_BF8 = 90
class VOP3SDOp(ReprEnum):
class VOP3SDOp(Enum):
V_ADD_CO_CI_U32 = 288
V_SUB_CO_CI_U32 = 289
V_SUBREV_CO_CI_U32 = 290
@@ -1801,7 +1802,7 @@ class VOP3SDOp(ReprEnum):
V_SUB_CO_U32 = 769
V_SUBREV_CO_U32 = 770
class VOPCOp(ReprEnum):
class VOPCOp(Enum):
V_CMP_LT_F16_E32 = 1
V_CMP_EQ_F16_E32 = 2
V_CMP_LE_F16_E32 = 3
@@ -2127,7 +2128,7 @@ class VOPCOp(ReprEnum):
V_CMPX_CLASS_F32 = V_CMPX_CLASS_F32_E32
V_CMPX_CLASS_F64 = V_CMPX_CLASS_F64_E32
class VOPDOp(ReprEnum):
class VOPDOp(Enum):
V_DUAL_FMAC_F32 = 0
V_DUAL_FMAAK_F32 = 1
V_DUAL_FMAMK_F32 = 2
@@ -2146,7 +2147,7 @@ class VOPDOp(ReprEnum):
V_DUAL_LSHLREV_B32 = 17
V_DUAL_AND_B32 = 18
class VSAMPLEOp(ReprEnum):
class VSAMPLEOp(Enum):
IMAGE_MSAA_LOAD = 24
IMAGE_SAMPLE = 27
IMAGE_SAMPLE_D = 28
@@ -2206,7 +2207,7 @@ class VSAMPLEOp(ReprEnum):
IMAGE_GATHER4_C_B_CL = 101
IMAGE_GATHER4H = 144
class VSCRATCHOp(ReprEnum):
class VSCRATCHOp(Enum):
SCRATCH_LOAD_U8 = 16
SCRATCH_LOAD_I8 = 17
SCRATCH_LOAD_U16 = 18
+16 -16
View File
@@ -101,11 +101,11 @@ class VFLAT(Inst):
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VGLOBAL(Inst):
encoding = FixedBitField(31, 24, 0b11101110)
encoding = FixedBitField(31, 24, 0b11101100)
op = EnumBitField(21, 14, VGLOBALOp, {VGLOBALOp.GLOBAL_LOAD_U8, VGLOBALOp.GLOBAL_LOAD_I8, VGLOBALOp.GLOBAL_LOAD_U16, VGLOBALOp.GLOBAL_LOAD_I16, VGLOBALOp.GLOBAL_LOAD_B32, VGLOBALOp.GLOBAL_LOAD_B64, VGLOBALOp.GLOBAL_LOAD_B96, VGLOBALOp.GLOBAL_LOAD_B128, VGLOBALOp.GLOBAL_STORE_B8, VGLOBALOp.GLOBAL_STORE_B16, VGLOBALOp.GLOBAL_STORE_B32, VGLOBALOp.GLOBAL_STORE_B64, VGLOBALOp.GLOBAL_STORE_B96, VGLOBALOp.GLOBAL_STORE_B128, VGLOBALOp.GLOBAL_LOAD_D16_U8, VGLOBALOp.GLOBAL_LOAD_D16_I8, VGLOBALOp.GLOBAL_LOAD_D16_B16, VGLOBALOp.GLOBAL_LOAD_D16_HI_U8, VGLOBALOp.GLOBAL_LOAD_D16_HI_I8, VGLOBALOp.GLOBAL_LOAD_D16_HI_B16, VGLOBALOp.GLOBAL_STORE_D16_HI_B8, VGLOBALOp.GLOBAL_STORE_D16_HI_B16, VGLOBALOp.GLOBAL_LOAD_ADDTID_B32, VGLOBALOp.GLOBAL_STORE_ADDTID_B32, VGLOBALOp.GLOBAL_INV, VGLOBALOp.GLOBAL_WB, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B32, VGLOBALOp.GLOBAL_ATOMIC_ADD_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_SUB_CLAMP_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_I32, VGLOBALOp.GLOBAL_ATOMIC_MIN_U32, VGLOBALOp.GLOBAL_ATOMIC_MAX_I32, VGLOBALOp.GLOBAL_ATOMIC_MAX_U32, VGLOBALOp.GLOBAL_ATOMIC_AND_B32, VGLOBALOp.GLOBAL_ATOMIC_OR_B32, VGLOBALOp.GLOBAL_ATOMIC_XOR_B32, VGLOBALOp.GLOBAL_ATOMIC_INC_U32, VGLOBALOp.GLOBAL_ATOMIC_DEC_U32, VGLOBALOp.GLOBAL_ATOMIC_SWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_CMPSWAP_B64, VGLOBALOp.GLOBAL_ATOMIC_ADD_U64, VGLOBALOp.GLOBAL_ATOMIC_SUB_U64, VGLOBALOp.GLOBAL_ATOMIC_MIN_I64, VGLOBALOp.GLOBAL_ATOMIC_MIN_U64, VGLOBALOp.GLOBAL_ATOMIC_MAX_I64, VGLOBALOp.GLOBAL_ATOMIC_MAX_U64, VGLOBALOp.GLOBAL_ATOMIC_AND_B64, VGLOBALOp.GLOBAL_ATOMIC_OR_B64, VGLOBALOp.GLOBAL_ATOMIC_XOR_B64, VGLOBALOp.GLOBAL_ATOMIC_INC_U64, VGLOBALOp.GLOBAL_ATOMIC_DEC_U64, VGLOBALOp.GLOBAL_WBINV, VGLOBALOp.GLOBAL_ATOMIC_COND_SUB_U32, VGLOBALOp.GLOBAL_ATOMIC_MIN_NUM_F32, VGLOBALOp.GLOBAL_ATOMIC_MAX_NUM_F32, VGLOBALOp.GLOBAL_LOAD_BLOCK, VGLOBALOp.GLOBAL_STORE_BLOCK, VGLOBALOp.GLOBAL_ATOMIC_ADD_F32, VGLOBALOp.GLOBAL_LOAD_TR_B128, VGLOBALOp.GLOBAL_LOAD_TR_B64, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_F16, VGLOBALOp.GLOBAL_ATOMIC_PK_ADD_BF16, VGLOBALOp.GLOBAL_ATOMIC_ORDERED_ADD_B64})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
@@ -114,7 +114,20 @@ class VGLOBAL(Inst):
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VSCRATCH(Inst):
encoding = FixedBitField(31, 24, 0b11101100)
op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
saddr = SGPRField(6, 0, default=NULL)
nv = BitField(7, 7)
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = BitField(62, 55)
ioffset = BitField(95, 72)
class VIMAGE(Inst):
@@ -240,19 +253,6 @@ class VSAMPLE(Inst):
vaddr2 = BitField(87, 80)
vaddr3 = BitField(95, 88)
class VSCRATCH(Inst):
encoding = FixedBitField(31, 24, 0b11101101)
op = EnumBitField(21, 14, VSCRATCHOp, {VSCRATCHOp.SCRATCH_LOAD_U8, VSCRATCHOp.SCRATCH_LOAD_I8, VSCRATCHOp.SCRATCH_LOAD_U16, VSCRATCHOp.SCRATCH_LOAD_I16, VSCRATCHOp.SCRATCH_LOAD_B32, VSCRATCHOp.SCRATCH_LOAD_B64, VSCRATCHOp.SCRATCH_LOAD_B96, VSCRATCHOp.SCRATCH_LOAD_B128, VSCRATCHOp.SCRATCH_STORE_B8, VSCRATCHOp.SCRATCH_STORE_B16, VSCRATCHOp.SCRATCH_STORE_B32, VSCRATCHOp.SCRATCH_STORE_B64, VSCRATCHOp.SCRATCH_STORE_B96, VSCRATCHOp.SCRATCH_STORE_B128, VSCRATCHOp.SCRATCH_LOAD_D16_U8, VSCRATCHOp.SCRATCH_LOAD_D16_I8, VSCRATCHOp.SCRATCH_LOAD_D16_B16, VSCRATCHOp.SCRATCH_LOAD_D16_HI_U8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_I8, VSCRATCHOp.SCRATCH_LOAD_D16_HI_B16, VSCRATCHOp.SCRATCH_STORE_D16_HI_B8, VSCRATCHOp.SCRATCH_STORE_D16_HI_B16, VSCRATCHOp.SCRATCH_LOAD_BLOCK, VSCRATCHOp.SCRATCH_STORE_BLOCK})
vdst = VGPRField(39, 32)
vaddr = VGPRField(71, 64)
saddr = SGPRField(6, 0, default=NULL)
nv = BitField(7, 7)
sve = BitField(49, 49)
scope = BitField(51, 50)
th = BitField(54, 52)
vsrc = VGPRField(62, 55)
ioffset = BitField(95, 72)
class SOP1_LIT(SOP1):
op = EnumBitField(15, 8, SOP1Op, {SOP1Op.S_MOV_B32, SOP1Op.S_MOV_B64, SOP1Op.S_CMOV_B32, SOP1Op.S_CMOV_B64, SOP1Op.S_BREV_B32, SOP1Op.S_BREV_B64, SOP1Op.S_CTZ_I32_B32, SOP1Op.S_CTZ_I32_B64, SOP1Op.S_CLZ_I32_U32, SOP1Op.S_CLZ_I32_U64, SOP1Op.S_CLS_I32, SOP1Op.S_CLS_I32_I64, SOP1Op.S_SEXT_I32_I8, SOP1Op.S_SEXT_I32_I16, SOP1Op.S_BITSET0_B32, SOP1Op.S_BITSET0_B64, SOP1Op.S_BITSET1_B32, SOP1Op.S_BITSET1_B64, SOP1Op.S_BITREPLICATE_B64_B32, SOP1Op.S_ABS_I32, SOP1Op.S_BCNT0_I32_B32, SOP1Op.S_BCNT0_I32_B64, SOP1Op.S_BCNT1_I32_B32, SOP1Op.S_BCNT1_I32_B64, SOP1Op.S_QUADMASK_B32, SOP1Op.S_QUADMASK_B64, SOP1Op.S_WQM_B32, SOP1Op.S_WQM_B64, SOP1Op.S_NOT_B32, SOP1Op.S_NOT_B64, SOP1Op.S_AND_SAVEEXEC_B32, SOP1Op.S_AND_SAVEEXEC_B64, SOP1Op.S_OR_SAVEEXEC_B32, SOP1Op.S_OR_SAVEEXEC_B64, SOP1Op.S_XOR_SAVEEXEC_B32, SOP1Op.S_XOR_SAVEEXEC_B64, SOP1Op.S_NAND_SAVEEXEC_B32, SOP1Op.S_NAND_SAVEEXEC_B64, SOP1Op.S_NOR_SAVEEXEC_B32, SOP1Op.S_NOR_SAVEEXEC_B64, SOP1Op.S_XNOR_SAVEEXEC_B32, SOP1Op.S_XNOR_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_SAVEEXEC_B32, SOP1Op.S_AND_NOT0_SAVEEXEC_B64, SOP1Op.S_OR_NOT0_SAVEEXEC_B32, SOP1Op.S_OR_NOT0_SAVEEXEC_B64, SOP1Op.S_AND_NOT1_SAVEEXEC_B32, SOP1Op.S_AND_NOT1_SAVEEXEC_B64, SOP1Op.S_OR_NOT1_SAVEEXEC_B32, SOP1Op.S_OR_NOT1_SAVEEXEC_B64, SOP1Op.S_AND_NOT0_WREXEC_B32, SOP1Op.S_AND_NOT0_WREXEC_B64, SOP1Op.S_AND_NOT1_WREXEC_B32, SOP1Op.S_AND_NOT1_WREXEC_B64, SOP1Op.S_MOVRELS_B32, SOP1Op.S_MOVRELS_B64, SOP1Op.S_MOVRELD_B32, SOP1Op.S_MOVRELD_B64, SOP1Op.S_MOVRELSD_2_B32, SOP1Op.S_GETPC_B64, SOP1Op.S_SETPC_B64, SOP1Op.S_SWAPPC_B64, SOP1Op.S_RFE_B64, SOP1Op.S_SENDMSG_RTN_B32, SOP1Op.S_SENDMSG_RTN_B64, SOP1Op.S_BARRIER_SIGNAL, SOP1Op.S_BARRIER_SIGNAL_ISFIRST, SOP1Op.S_GET_BARRIER_STATE, SOP1Op.S_BARRIER_INIT, SOP1Op.S_BARRIER_JOIN, SOP1Op.S_ALLOC_VGPR, SOP1Op.S_SLEEP_VAR, SOP1Op.S_CEIL_F32, SOP1Op.S_FLOOR_F32, SOP1Op.S_TRUNC_F32, SOP1Op.S_RNDNE_F32, SOP1Op.S_CVT_F32_I32, SOP1Op.S_CVT_F32_U32, SOP1Op.S_CVT_I32_F32, SOP1Op.S_CVT_U32_F32, SOP1Op.S_CVT_F16_F32, SOP1Op.S_CVT_F32_F16, SOP1Op.S_CVT_HI_F32_F16, SOP1Op.S_CEIL_F16, SOP1Op.S_FLOOR_F16, SOP1Op.S_TRUNC_F16, SOP1Op.S_RNDNE_F16})
literal = BitField(63, 32)
@@ -29,16 +29,15 @@ def _matches(data: bytes, cls: type[Inst]) -> bool:
# Import instruction classes for each architecture
from extra.assembly.amd.autogen.rdna3.ins import (VOP1, VOP1_SDST, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3SD, VOP3P, VOPC, VOPD, VINTERP,
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH)
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_LIT as R4_VOP1_LIT,
VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT,
SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT, SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT,
SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT,
VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P,
VOPC as R4_VOPC, VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
SMEM as R4_SMEM, DS as R4_DS, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
from extra.assembly.amd.autogen.cdna.ins import (VOP1 as C_VOP1, VOP1_SDWA as C_VOP1_SDWA, VOP1_DPP16 as C_VOP1_DPP16,
VOP2 as C_VOP2, VOP2_LIT as C_VOP2_LIT, VOP2_SDWA as C_VOP2_SDWA, VOP2_DPP16 as C_VOP2_DPP16,
VOPC as C_VOPC, VOPC_SDWA_SDST as C_VOPC_SDWA_SDST,
VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3P_MFMA as C_VOP3P_MFMA, VOP3PX2 as C_VOP3PX2,
VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3PX2 as C_VOP3PX2,
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPK_LIT as C_SOPK_LIT, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
FLAT as C_FLAT, GLOBAL as C_GLOBAL, SCRATCH as C_SCRATCH, MUBUF as C_MUBUF)
@@ -48,9 +47,8 @@ _FORMATS = {
"rdna3": [VOPD, VOP3P, VINTERP, VOP3SD, VOP3_SDST, VOP3, DS, GLOBAL, SCRATCH, FLAT, SMEM,
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, VOPC, VOP1_SDST, VOP1, VOP1_LIT, VOP2, VOP2_LIT],
"rdna4": [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3SD, R4_VOP3_SDST, R4_VOP3, R4_DS, R4_GLOBAL, R4_SCRATCH, R4_FLAT, R4_SMEM,
R4_SOP1, R4_SOP1_LIT, R4_SOPC, R4_SOPC_LIT, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_VOP1_LIT,
R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
"cdna": [C_VOP3PX2, C_VOP3P_MFMA, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
R4_SOP1, R4_SOPC, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC, R4_VOP1_SDST, R4_VOP1, R4_SOP2, R4_SOP2_LIT, R4_VOP2, R4_VOP2_LIT],
"cdna": [C_VOP3PX2, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_SOPK_LIT, C_VOPC_SDWA_SDST, C_VOPC,
C_VOP1_DPP16, C_VOP1_SDWA, C_VOP1, C_VOP2_DPP16, C_VOP2_SDWA, C_SOP2, C_VOP2, C_VOP2_LIT],
}
+22 -30
View File
@@ -1,7 +1,6 @@
# RDNA3/RDNA4/CDNA disassembler
from __future__ import annotations
import re, struct
from typing import Callable
from extra.assembly.amd.dsl import Inst, Reg
# Special register mappings for disassembly
@@ -85,11 +84,10 @@ from extra.assembly.amd.autogen.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as
VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4)
from extra.assembly.amd.autogen.cdna.ins import FLAT as C_FLAT, HWREG as HWREG_CDNA
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
def _is_r4(inst: Inst) -> bool: return 'rdna4' in inst.__class__.__module__
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
@@ -102,7 +100,7 @@ def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1
def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
def _vreg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("v", b - 256 if b >= 256 else b, n)
def _areg(b: int, n: int = 1) -> str: b = _unwrap(b); return _reg("a", b - 256 if b >= 256 else b, n) # accumulator registers for GFX90a
def _ttmp(b, n: int = 1) -> str | None: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _ttmp(b, n: int = 1) -> str: b = _unwrap(b); return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
v = _unwrap(v)
@@ -227,7 +225,7 @@ NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
def _disasm_sopp(inst: SOPP) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = _is_r4(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
# Ops that have no argument when simm16 == 0
no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
's_endpgm_ordered_ps_done', 's_code_end'}
@@ -259,14 +257,14 @@ def _disasm_sopp(inst: SOPP) -> str:
dep = lambda v: deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} 0x{inst.simm16:x}"
return f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if name in ('s_gl1_inv', 's_dcache_inv', 's_dcache_inv_vol', 's_dcache_wb', 's_dcache_wb_vol', 's_icache_inv'): return name
soe, imm = getattr(inst, 'soe', 0) or getattr(inst, 'soffset_en', 0), getattr(inst, 'imm', 1)
is_rdna4 = _is_r4(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0)
if cdna:
if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
@@ -305,30 +303,27 @@ def _disasm_smem(inst: SMEM) -> str:
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
def _disasm_flat(inst: FLAT) -> str:
name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
name, cdna = inst.op_name.lower(), _is_cdna(inst)
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
# Global/scratch uses 13-bit signed offset
offset = inst.ioffset if r4 else inst.offset
if seg != 'flat':
if cdna:
# CDNA: bit 12 is sign bit but not in offset field
raw = int.from_bytes(inst.to_bytes(), 'little')
off_val = offset | ((raw >> 12) & 1) << 12 # get bit 12
off_val = inst.offset | ((raw >> 12) & 1) << 12 # get bit 12
else:
off_val = offset
off_val = inst.offset
off_val = off_val if off_val < 4096 else off_val - 8192 # sign extend 13-bit
else:
off_val = offset
off_val = inst.offset
# Use get_field_bits: data for stores/atomics, d for loads
regs = inst.canonical_op_regs
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
off_s = f" offset:{off_val}" if off_val else ""
if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"
elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
if seg == 'flat': saddr_s = ""
elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
@@ -337,21 +332,18 @@ def _disasm_flat(inst: FLAT) -> str:
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
# RDNA4: vaddr instead of addr, vsrc instead of data
addr = inst.vaddr if r4 else inst.addr
data = inst.vsrc if r4 else inst.data
# load_lds_* instructions: vaddr, saddr (no vdst, data goes to LDS)
if 'load_lds' in name:
addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
return f"{instr} {addr_s}{saddr_s}{mods}"
if seg == 'flat': addr_w = 2 # flat always uses 64-bit vaddr
elif cdna: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(inst.addr, addr_w)
data_s, vdst_s = reg_fn(inst.data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
glc_or_sc0 = inst.sc0 if cdna else inst.glc
if 'atomic' in name:
glc_or_sc0 = inst.sc0 if cdna else inst.glc
return f"{instr} {vdst_s}, {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{saddr_s if seg != 'flat' else ''}{mods}"
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
@@ -466,15 +458,16 @@ def _disasm_vop3sd(inst: VOP3SD) -> str:
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal
op_enum = R4_VOPDOp if _is_r4(inst) else VOPDOp
nx, ny = op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
is_rdna4 = 'rdna4' in inst.__class__.__module__
op_enum = R4_VOPDOp if is_rdna4 else VOPDOp
vdst_y, nx, ny = (_unwrap(inst.vdsty) << 1) | ((_unwrap(inst.vdstx) & 1) ^ 1), op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
def half(n, vd, s0, vs1):
vd, vs1 = _vi(vd), _vi(vs1)
if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, inst.vdsty, inst.srcy0, inst.vsrcy1)}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, vdst_y, inst.srcy0, inst.vsrcy1)}"
def _disasm_vop3p(inst: VOP3P) -> str:
name = inst.op_name.lower()
@@ -552,7 +545,7 @@ _HWREG_BLACKLIST_CDNA = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_R
'HW_REG_SQ_PERF_SNAPSHOT_PC_LO', 'HW_REG_SQ_PERF_SNAPSHOT_PC_HI', 'HW_REG_XCC_ID'}
def _disasm_sopk(inst: SOPK) -> str:
op, name, cdna = inst.op, inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = _is_r4(inst)
is_rdna4 = 'rdna4' in inst.__class__.__module__
hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
def fmt_hwreg(hid, hoff, hsz):
@@ -576,7 +569,7 @@ def _disasm_vinterp(inst: VINTERP) -> str:
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {_lit(inst, inst.src0, inst.neg & 1)}, {_lit(inst, inst.src1, inst.neg & 2)}, {_lit(inst, inst.src2, inst.neg & 4)}" + (" " + mods if mods else "")
DISASM_HANDLERS: dict[type, Callable[..., str]] = {
DISASM_HANDLERS: dict[type, callable] = {
VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3, VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
@@ -589,7 +582,6 @@ DISASM_HANDLERS: dict[type, Callable[..., str]] = {
R4_VOP2: _disasm_vop2, R4_VOP2_LIT: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOPC_LIT: _disasm_vopc,
R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3_SDST_LIT: _disasm_vop3, R4_VOP3_LIT: _disasm_vop3,
R4_VOP3SD: _disasm_vop3sd, R4_VOP3SD_LIT: _disasm_vop3sd, R4_VOP3P: _disasm_vop3p, R4_VOP3P_LIT: _disasm_vop3p,
R4_FLAT: _disasm_flat, R4_GLOBAL: _disasm_flat, R4_SCRATCH: _disasm_flat,
R4_VOPD: _disasm_vopd, R4_VOPD_LIT: _disasm_vopd, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem, R4_DS: _disasm_ds,
R4_SOP1: _disasm_sop1, R4_SOP1_LIT: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOP2_LIT: _disasm_sop2,
R4_SOPC: _disasm_sopc, R4_SOPC_LIT: _disasm_sopc, R4_SOPK: _disasm_sopk, R4_SOPK_LIT: _disasm_sopk}
@@ -604,7 +596,7 @@ from extra.assembly.amd.autogen.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as
VOP1_SDWA as CDNA_VOP1_SDWA, VOP1_DPP16 as CDNA_VOP1_DPP16,
VOP2 as CDNA_VOP2, VOP2_LIT as CDNA_VOP2_LIT, VOP2_SDWA as CDNA_VOP2_SDWA, VOP2_DPP16 as CDNA_VOP2_DPP16,
VOPC as CDNA_VOPC, VOPC_LIT as CDNA_VOPC_LIT, VOPC_SDWA_SDST as CDNA_VOPC_SDWA_SDST,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3P_MFMA as CDNA_VOP3P_MFMA, VOP3PX2 as CDNA_VOP3PX2,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3PX2 as CDNA_VOP3PX2,
SOP1 as CDNA_SOP1, SOP1_LIT as CDNA_SOP1_LIT, SOP2 as CDNA_SOP2, SOP2_LIT as CDNA_SOP2_LIT,
SOPC as CDNA_SOPC, SOPC_LIT as CDNA_SOPC_LIT, SOPK as CDNA_SOPK, SOPK_LIT as CDNA_SOPK_LIT,
SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
@@ -903,5 +895,5 @@ DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP1_LIT: _disasm_vop1,
CDNA_SOP1: _disasm_sop1, CDNA_SOP1_LIT: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOP2_LIT: _disasm_sop2,
CDNA_SOPC: _disasm_sopc, CDNA_SOPC_LIT: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPK_LIT: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_GLOBAL: _disasm_flat, CDNA_SCRATCH: _disasm_flat,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b, CDNA_VOP3P: _disasm_cdna_vop3p,
CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})
+45 -71
View File
@@ -1,24 +1,27 @@
# dsl.py - clean DSL for AMD assembly
from typing import Any
# ══════════════════════════════════════════════════════════════
# Registers - unified src encoding space (0-511)
# ══════════════════════════════════════════════════════════════
def _reg_size(t: str | None) -> int: return {'b64': 2, 'f64': 2, 'u64': 2, 'i64': 2, 'b128': 4}.get(t, 1)
class Reg:
# Register names vary by arch: RDNA has NULL@124/M0@125, CDNA has M0@124/reserved@125
# RDNA4 has DPP8@233, CDNA has SDWA@249/DPP@250/VCCZ@251/EXECZ@252
_NAMES = {102: "FLAT_SCRATCH_LO", 103: "FLAT_SCRATCH_HI", 104: "XNACK_MASK_LO", 105: "XNACK_MASK_HI",
106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
233: "DPP8", 234: "DPP8FI", 235: "SHARED_BASE", 236: "SHARED_LIMIT", 237: "PRIVATE_BASE", 238: "PRIVATE_LIMIT",
_NAMES = {106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
240: "0.5", 241: "-0.5", 242: "1.0", 243: "-1.0", 244: "2.0", 245: "-2.0", 246: "4.0", 247: "-4.0",
248: "INV_2PI", 249: "SDWA", 250: "DPP", 251: "VCCZ", 252: "EXECZ", 253: "SCC", 254: "SRC_LDS_DIRECT", 255: "LIT"}
248: "INV_2PI", 250: "DPP16", 253: "SCC", 255: "LIT"}
_PAIRS = {106: "VCC", 126: "EXEC"}
def __init__(self, offset: int = 0, sz: int = 512, *, neg: bool = False, abs_: bool = False, hi: bool = False):
self.offset, self.sz = offset, sz
self.neg, self.abs_, self.hi = neg, abs_, hi
# TODO: remove these legacy aliases
@property
def count(self): return self.sz
@property
def idx(self): return self.offset
def __hash__(self): return hash((self.offset, self.sz, self.neg, self.abs_, self.hi))
def __getitem__(self, key):
if isinstance(key, slice):
@@ -75,13 +78,9 @@ EXEC = src[126:127]
# 128: 0, 129-192: integers 1-64, 193-208: integers -1 to -16
# 240-248: float constants (0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0, -4.0, 1/(2*PI))
INV_2PI = src[248]
SDWA = src[249]
DPP = DPP16 = src[250]
VCCZ = src[251]
EXECZ = src[252]
DPP16 = src[250]
SCC = src[253]
SRC_LDS_DIRECT = src[254]
LIT = src[255] # literal constant marker
# 255: literal constant
v = src[256:511] # VGPR0-255
# ══════════════════════════════════════════════════════════════
@@ -94,13 +93,12 @@ class _Bits:
bits = _Bits()
class BitField:
name: str | None
def __init__(self, hi: int, lo: int, default: int = 0):
self.hi, self.lo, self.default, self.name, self.mask = hi, lo, default, None, (1 << (hi - lo + 1)) - 1
def __set_name__(self, owner, name: str): self.name = name
def __eq__(self, other) -> 'FixedBitField': # type: ignore[override]
def __set_name__(self, owner, name): self.name = name
def __eq__(self, other) -> 'FixedBitField':
if isinstance(other, int): return FixedBitField(self.hi, self.lo, other)
raise TypeError(f"BitField.__eq__ expects int, got {type(other).__name__}")
return NotImplemented
def enum(self, enum_cls) -> 'EnumBitField': return EnumBitField(self.hi, self.lo, enum_cls)
def encode(self, val) -> int:
assert isinstance(val, int), f"BitField.encode expects int, got {type(val).__name__}"
@@ -109,14 +107,11 @@ class BitField:
def set(self, raw: int, val) -> int:
if val is None: val = self.default
encoded = self.encode(val)
# Handle signed values: convert negative to 2's complement
if encoded < 0: encoded = encoded & self.mask
if encoded < 0 or encoded > self.mask: raise RuntimeError(f"field '{self.name}': value {encoded} doesn't fit in {self.hi - self.lo + 1} bits")
return (raw & ~(self.mask << self.lo)) | (encoded << self.lo)
def __get__(self, obj, objtype=None):
if obj is None: return self
return self.decode((obj._raw >> self.lo) & self.mask)
def __set__(self, obj, val): obj._raw = self.set(obj._raw, val)
class FixedBitField(BitField):
def set(self, raw: int, val=None) -> int:
@@ -173,10 +168,7 @@ class SrcField(BitField):
# Resize register based on operand info (skip non-resizable special registers)
# VCC/EXEC pairs (106, 126), NULL (124), M0 (125), float constants (240-255)
if reg.offset not in (124, 125) and not 240 <= reg.offset <= 255:
# Map variant field names (vsrc0->src0, vsrc1->src1, etc.) for DPP/SDWA classes
assert self.name is not None
name = self.name[1:] if self.name.startswith('v') and self.name[1:] in obj.op_regs else self.name
if sz := obj.op_regs.get(name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
return reg
class VGPRField(SrcField):
@@ -219,12 +211,7 @@ class VDSTYField(BitField):
if not isinstance(val, Reg): raise TypeError(f"VDSTYField requires Reg, got {type(val).__name__}")
if not (256 <= val.offset < 512): raise ValueError(f"VDSTYField requires VGPR, got offset {val.offset}")
return (val.offset - 256) >> 1
def __get__(self, obj, objtype=None):
if obj is None: return self
raw = (obj._raw >> self.lo) & self.mask
vdstx_bit0 = (obj.vdstx.offset - 256) & 1
vgpr_idx = (raw << 1) | (vdstx_bit0 ^ 1)
return Reg(256 + vgpr_idx, 1)
def decode(self, raw): return raw # raw value, actual vdsty = (raw << 1) | ((vdstx & 1) ^ 1)
# ══════════════════════════════════════════════════════════════
# Operand info from XML
@@ -253,15 +240,6 @@ def _get_variant(cls, suffix: str):
module = sys.modules.get(cls.__module__)
return getattr(module, f"{cls.__name__}{suffix}", None) if module else None
def _canonical_name(name: str) -> str | None:
"""Map operand name to canonical name."""
if name in ('src0', 'vsrc0', 'ssrc0'): return 's0'
if name in ('src1', 'vsrc1', 'ssrc1'): return 's1'
if name == 'src2': return 's2'
if name in ('vdst', 'sdst', 'sdata'): return 'd'
if name in ('data', 'vdata', 'data0', 'vsrc'): return 'data'
return None
class Inst:
_fields: list[tuple[str, BitField]]
_base_size: int
@@ -277,30 +255,31 @@ class Inst:
cls._base_size = (max(f.hi for _, f in cls._fields) + 8) // 8
def __new__(cls, *args, **kwargs):
# Auto-upgrade to variant if needed (only for base classes, not variants)
# Auto-upgrade to _LIT variant if needed (only for base classes, not variants)
if not any(cls.__name__.endswith(sfx) for sfx in ('_LIT', '_DPP16', '_DPP8', '_SDWA', '_SDWA_SDST', '_MFMA')):
args_iter = iter(args)
for name, field in cls._fields:
if isinstance(field, FixedBitField): continue
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
if not isinstance(field, SrcField): continue
if isinstance(val, Reg) and val.offset == 255 and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
if isinstance(val, Reg) and val.offset == 249:
if (sdwa_cls := _get_variant(cls, '_SDWA') or _get_variant(cls, '_SDWA_SDST')): return sdwa_cls(*args, **kwargs)
if isinstance(val, Reg) and val.offset == 250 and (dpp_cls := _get_variant(cls, '_DPP16')): return dpp_cls(*args, **kwargs)
if _needs_literal(val) and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
lit_cls = _get_variant(cls, '_LIT')
if lit_cls is not None:
# Check if any src field needs a literal
# Map positional args to field names to find src values
args_iter = iter(args)
for name, field in cls._fields:
if isinstance(field, FixedBitField): continue
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
if isinstance(field, SrcField) and _needs_literal(val):
return lit_cls(*args, **kwargs)
return object.__new__(cls)
def __init__(self, *args, **kwargs):
self._raw = 0
# Map positional args to field names (skip FixedBitFields)
args_iter = iter(args)
vals: dict[str, Any] = {}
vals = {}
for name, field in self._fields:
if isinstance(field, FixedBitField): vals[name] = None
elif name in kwargs: vals[name] = kwargs[name]
else: vals[name] = next(args_iter, None)
assert not (remaining := list(args_iter)), f"too many positional args: {remaining}"
remaining = list(args_iter)
assert not remaining, f"too many positional args: {remaining}"
# Extract modifiers from Reg objects and merge into neg/abs/opsel
neg_bits, abs_bits, opsel_bits = 0, 0, 0
for name, bit in [('src0', 0), ('src1', 1), ('src2', 2)]:
@@ -325,16 +304,16 @@ class Inst:
# Set all field values
for name, field in self._fields:
self._raw = field.set(self._raw, vals[name])
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC, SDWA/DPP markers)
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC)
for name, expected in self.op_regs.items():
if (val := vals.get(name)) is None: continue
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or 249 <= val.offset <= 255):
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or val.offset == 253):
raise TypeError(f"{name} expects {expected} register(s), got {val.sz}")
@property
def op_name(self) -> str: return getattr(self, 'op').name
def op_name(self) -> str: return self.op.name
@property
def operands(self) -> dict: return OPERANDS.get(getattr(self, 'op'), {}) if hasattr(self, 'op') else {}
def operands(self) -> dict: return OPERANDS.get(self.op, {}) if hasattr(self, 'op') else {}
def _is_cdna(self) -> bool: return 'cdna' in type(self).__module__
@functools.cached_property
@@ -346,9 +325,9 @@ class Inst:
if not self._is_cdna():
name = self.op_name.lower()
if 'cndmask' in name and 'src2' in bits: bits['src2'] = 32
if '_co_ci_' in name and 'src2' in bits: bits['src2'] = 32 # carry-in source
# VOP3SD: sdst is always wavefront-size dependent (carry-out or condition mask)
if 'VOP3SD' in type(self).__name__ and 'sdst' in bits: bits['sdst'] = 32
if '_co_ci_' in name:
if 'src2' in bits: bits['src2'] = 32
if 'sdst' in bits: bits['sdst'] = 32
if 'cmp' in name and 'vdst' in bits: bits['vdst'] = 32
# GLOBAL/FLAT: addr is 32-bit if saddr is valid SGPR, 64-bit if saddr is NULL
# SCRATCH: addr is always 32-bit (offset from scratch base, not absolute address)
@@ -362,8 +341,8 @@ class Inst:
# VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
if 'f8f6f4' in getattr(self, 'op_name', '').lower():
# Use explicit fields if available (VOP3PX2), else extract from VOP3P-MAI bit positions
cbsz = getattr(self, 'cbsz') if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
blgp = getattr(self, 'blgp') if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
cbsz = self.cbsz if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
blgp = self.blgp if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
bits['src0'], bits['src1'] = vgprs.get(cbsz, 8) * 32, vgprs.get(blgp, 8) * 32
return bits
@@ -377,17 +356,12 @@ class Inst:
"""Get bit widths with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
bits = {'d': 32, 's0': 32, 's1': 32, 's2': 32, 'data': 32}
for name, val in self.op_bits.items():
if (cn := _canonical_name(name)): bits[cn] = val
if name in ('src0', 'vsrc0', 'ssrc0'): bits['s0'] = val
elif name in ('src1', 'vsrc1', 'ssrc1'): bits['s1'] = val
elif name == 'src2': bits['s2'] = val
elif name in ('vdst', 'sdst', 'sdata'): bits['d'] = val
elif name in ('data', 'vdata', 'data0'): bits['data'] = val
return bits
@functools.cached_property
def canonical_operands(self) -> dict:
"""Get operands with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
result = {}
for name, val in self.operands.items():
if (cn := _canonical_name(name)): result[cn] = val
return result
@property
def canonical_op_regs(self) -> dict[str, int]:
"""Get register counts with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
+413 -1064
View File
File diff suppressed because it is too large Load Diff
+779 -1009
View File
File diff suppressed because it is too large Load Diff
+72 -144
View File
@@ -26,7 +26,7 @@ class AluSrc(Enum):
VALU_SALU = 3
class InstOp(Enum):
"""SQTT instruction operation types for RDNA3 (gfx1100).
"""SQTT instruction operation types.
Memory ops appear in two ranges depending on which SIMD executes them:
- 0x1x-0x2x range: ops on traced SIMD
@@ -95,21 +95,6 @@ class InstOp(Enum):
SALU_SAVEEXEC = 0x72 # s_*_saveexec_b32/b64
VALU_CMPX = 0x73 # v_cmpx_*
class InstOpL4(Enum):
"""SQTT instruction operation types for RDNA4 (gfx1200). Different encoding from RDNA3."""
# TODO: we need to do discovery of all of these from instructions
SALU = 0x0
SMEM = 0x1
UNK_02 = 0x2
JUMP_NO = 0x4
UNK_06 = 0x6
VMEM = 0x10
UNK_11 = 0x11
VINTERP = 0x12
UNK_14 = 0x14
OTHER_VMEM = 0x5e
UNK_60 = 0x60
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE BASE CLASS
# ═══════════════════════════════════════════════════════════════════════════════
@@ -132,64 +117,9 @@ class PacketType:
return inst
def __repr__(self) -> str:
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_') and k != 'encoding')
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_'))
return f"{self.__class__.__name__}({fields_str})"
# ═══════════════════════════════════════════════════════════════════════════════
# TS PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
class TS_DELTA_S8_W3(PacketType):
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[63:11]
class TS_DELTA_S8_W3_L4(PacketType): # Layout 4: 64->72 bits
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[71:11]
class TS_DELTA_S5_W3(PacketType):
encoding = bits[4:0] == 0b00110
delta = bits[7:5]
_padding = bits[51:8]
class TS_DELTA_S5_W3_L4(PacketType): # Layout 4: 52->56 bits
encoding = bits[4:0] == 0b00110
delta = bits[9:7]
_padding = bits[55:10]
class TS_DELTA_SHORT(PacketType):
encoding = bits[3:0] == 0b1000
delta = bits[7:4]
class TS_DELTA_OR_MARK(PacketType):
encoding = bits[6:0] == 0b0000001
delta = bits[47:12]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
class TS_DELTA_OR_MARK_L4(PacketType): # Layout 4: 48->64 bits
encoding = bits[6:0] == 0b0000001
delta = bits[63:12]
bit7 = bits[7:7]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool((self.bit9 and not self.bit8) or self.bit7)
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[47:7]
class TS_DELTA_S5_W2_L4(PacketType): # Layout 4: 48->40 bits
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
_padding = bits[39:7]
# ═══════════════════════════════════════════════════════════════════════════════
# PACKET TYPE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════════
@@ -225,6 +155,11 @@ class WAVERDY(PacketType): # exclude: 1 << 3
delta = bits[7:5]
mask = bits[23:8]
class TS_DELTA_S8_W3(PacketType):
encoding = bits[6:0] == 0b0100001
delta = bits[10:8]
_padding = bits[63:11]
class WAVEEND(PacketType): # exclude: 1 << 4
encoding = bits[4:0] == 0b10101
delta = bits[7:5]
@@ -246,36 +181,29 @@ class WAVESTART(PacketType): # exclude: 1 << 4
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
class WAVESTART_L4(PacketType): # Layout 4 has wave field at different position
encoding = bits[4:0] == 0b01100
class TS_DELTA_S5_W2(PacketType):
encoding = bits[4:0] == 0b11100
delta = bits[6:5]
flag7 = bits[7:7]
simd = bits[9:8]
cu_lo = bits[12:10]
wave = bits[19:15]
id7 = bits[31:20]
@property
def cu(self) -> int: return self.cu_lo | (self.flag7 << 3)
_padding = bits[47:7]
class WAVEALLOC(PacketType): # exclude: 1 << 10
encoding = bits[4:0] == 0b00101
delta = bits[7:5]
_padding = bits[19:8]
class WAVEALLOC_L4(PacketType): # Layout 4: 20->24 bits
encoding = bits[4:0] == 0b00101
class TS_DELTA_S5_W3(PacketType):
encoding = bits[4:0] == 0b00110
delta = bits[7:5]
_padding = bits[23:8]
_padding = bits[51:8]
class PERF(PacketType): # exclude: 1 << 11
encoding = bits[4:0] == 0b10110
delta = bits[7:5]
arg = bits[27:8]
class PERF_L4(PacketType): # Layout 4: 28->32 bits
encoding = bits[4:0] == 0b10110
delta = bits[9:7]
arg = bits[31:10]
class TS_DELTA_SHORT(PacketType):
encoding = bits[3:0] == 0b1000
delta = bits[7:4]
class NOP(PacketType):
encoding = bits[3:0] == 0b0000
@@ -316,6 +244,14 @@ class SNAPSHOT(PacketType):
delta = bits[9:7]
snap = bits[63:10]
class TS_DELTA_OR_MARK(PacketType):
encoding = bits[6:0] == 0b0000001
delta = bits[47:12]
bit8 = bits[8:8]
bit9 = bits[9:9]
@property
def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)
class LAYOUT_HEADER(PacketType):
encoding = bits[6:0] == 0b0010001
delta = None # type: ignore
@@ -335,54 +271,56 @@ class INST(PacketType):
wave = bits[12:8]
op = bits[19:13].enum(InstOp)
class INST_L4(PacketType): # Layout 4: different delta position and InstOp encoding
encoding = bits[2:0] == 0b010
delta = bits[5:3]
flag1 = bits[6:6]
flag2 = bits[7:7]
wave = bits[12:8]
op = bits[19:13].enum(InstOpL4)
class UTILCTR(PacketType):
encoding = bits[6:0] == 0b0110001
delta = bits[8:7]
ctr = bits[47:9]
# Packet types with rocprof type IDs as keys
PACKET_TYPES_L3: dict[int, type[PacketType]] = {
1: VALUINST, 2: VMEMEXEC, 3: ALUEXEC, 4: IMMEDIATE, 5: IMMEDIATE_MASK, 6: WAVERDY, 7: TS_DELTA_S8_W3, 8: WAVEEND,
9: WAVESTART, 10: TS_DELTA_S5_W2, 11: WAVEALLOC, 12: TS_DELTA_S5_W3, 13: PERF, 14: UTILCTR, 15: TS_DELTA_SHORT,
16: NOP, 17: TS_WAVE_STATE, 18: EVENT, 19: EVENT_BIG, 20: REG, 21: SNAPSHOT, 22: TS_DELTA_OR_MARK, 23: LAYOUT_HEADER, 24: INST,
}
PACKET_TYPES_L4: dict[int, type[PacketType]] = {
**PACKET_TYPES_L3,
7: TS_DELTA_S8_W3_L4, 9: WAVESTART_L4, 10: TS_DELTA_S5_W2_L4, 11: WAVEALLOC_L4,
12: TS_DELTA_S5_W3_L4, 13: PERF_L4, 22: TS_DELTA_OR_MARK_L4, 24: INST_L4,
}
def _build_decode_tables(packet_types: dict[int, type[PacketType]]) -> tuple[dict[int, tuple], bytes]:
# Build state table: byte -> opcode. Sort by mask specificity (more bits first), NOP last
sorted_types = sorted(packet_types.items(), key=lambda x: (-bin(x[1].encoding.mask).count('1'), x[0] == 16))
state_table = bytes(next((op for op, cls in sorted_types if (b & cls.encoding.mask) == cls.encoding.default), 16) for b in range(256))
# Build decode info: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 8)
decode_info = {}
for opcode, pkt_cls in packet_types.items():
delta_field = getattr(pkt_cls, 'delta', None)
special = {22: 1, 15: 2}.get(opcode, 0) # TS_DELTA_OR_MARK=22, TS_DELTA_SHORT=15
decode_info[opcode] = (pkt_cls, pkt_cls._size_nibbles, delta_field.lo if delta_field else 0, delta_field.mask if delta_field else 0, special)
return decode_info, state_table
# All packet types in encoding priority order (more specific masks first, NOP last as fallback)
PACKET_TYPES: list[type[PacketType]] = [
EVENT, EVENT_BIG,
TS_DELTA_S8_W3, TS_WAVE_STATE, SNAPSHOT, TS_DELTA_OR_MARK, LAYOUT_HEADER, UTILCTR,
IMMEDIATE_MASK, WAVERDY, WAVEEND, WAVESTART, TS_DELTA_S5_W2, WAVEALLOC, TS_DELTA_S5_W3, PERF,
VMEMEXEC, ALUEXEC, IMMEDIATE, TS_DELTA_SHORT, REG,
VALUINST, INST,
NOP,
]
_DECODE_INFO_L3, _STATE_TABLE_L3 = _build_decode_tables(PACKET_TYPES_L3)
_DECODE_INFO_L4, _STATE_TABLE_L4 = _build_decode_tables(PACKET_TYPES_L4)
def _build_state_table() -> tuple[bytes, dict[int, type[PacketType]]]:
table = [len(PACKET_TYPES) - 1] * 256 # default to NOP
opcode_to_class: dict[int, type[PacketType]] = {i: cls for i, cls in enumerate(PACKET_TYPES)}
for byte_val in range(256):
for opcode, pkt_cls in enumerate(PACKET_TYPES):
if (byte_val & pkt_cls.encoding.mask) == pkt_cls.encoding.default:
table[byte_val] = opcode
break
return bytes(table), opcode_to_class
STATE_TO_OPCODE, OPCODE_TO_CLASS = _build_state_table()
# Precompute special case opcodes
_TS_DELTA_OR_MARK_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_OR_MARK)
_TS_DELTA_SHORT_OPCODE = next(op for op, cls in OPCODE_TO_CLASS.items() if cls is TS_DELTA_SHORT)
# Combined lookup: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
# special_case: 0=none, 1=TS_DELTA_OR_MARK, 2=TS_DELTA_SHORT
_DECODE_INFO: dict[int, tuple] = {}
for _opcode, _pkt_cls in OPCODE_TO_CLASS.items():
_delta_field = getattr(_pkt_cls, 'delta', None)
_delta_lo = _delta_field.lo if _delta_field else 0
_delta_mask = _delta_field.mask if _delta_field else 0
_special = 1 if _opcode == _TS_DELTA_OR_MARK_OPCODE else (2 if _opcode == _TS_DELTA_SHORT_OPCODE else 0)
_DECODE_INFO[_opcode] = (_pkt_cls, _pkt_cls._size_nibbles, _delta_lo, _delta_mask, _special)
# ═══════════════════════════════════════════════════════════════════════════════
# DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def decode(data: bytes) -> Iterator[PacketType]:
"""Decode raw SQTT blob, yielding packet instances. Auto-detects layout from LAYOUT_HEADER."""
"""Decode raw SQTT blob, yielding packet instances."""
n, reg, pos, nib_off, nib_count, time = len(data), 0, 0, 0, 16, 0
decode_info, state_table = _DECODE_INFO_L3, _STATE_TABLE_L3 # default to layout 3, will update after seeing LAYOUT_HEADER
while pos + ((nib_count + nib_off + 1) >> 1) <= n:
need = nib_count - nib_off
@@ -395,20 +333,13 @@ def decode(data: bytes) -> Iterator[PacketType]:
# 3. if odd, read low nibble
if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
opcode = state_table[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
opcode = STATE_TO_OPCODE[reg & 0xFF]
pkt_cls, nib_count, delta_lo, delta_mask, special = _DECODE_INFO[opcode]
delta = (reg >> delta_lo) & delta_mask
if special == 1: # TS_DELTA_OR_MARK
pkt = pkt_cls.from_raw(reg, 0) # create packet to check is_marker
if pkt.is_marker: delta = 0
if special == 1 and (reg >> 9) & 1 and not (reg >> 8) & 1: delta = 0 # TS_DELTA_OR_MARK marker
elif special == 2: delta += 8 # TS_DELTA_SHORT
time += delta
pkt = pkt_cls.from_raw(reg, time)
# detect layout from first LAYOUT_HEADER and switch decode tables if needed
# NOTE: CDNA uses a completely different 16-bit header format, not nibbles - not supported here
if pkt_cls is LAYOUT_HEADER and pkt.layout == 4:
decode_info, state_table = _DECODE_INFO_L4, _STATE_TABLE_L4
yield pkt
yield pkt_cls.from_raw(reg, time)
# ═══════════════════════════════════════════════════════════════════════════════
# PRINTER
@@ -424,26 +355,23 @@ PACKET_COLORS = {
def format_packet(p) -> str:
from tinygrad.helpers import colored
name = type(p).__name__
if isinstance(p, (INST, INST_L4)):
op_name = p.op.name if isinstance(p.op, (InstOp, InstOpL4)) else f"0x{p.op:02x}"
if isinstance(p, INST):
op_name = p.op.name if isinstance(p.op, InstOp) else f"0x{p.op:02x}"
fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
elif isinstance(p, (WAVESTART, WAVESTART_L4, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
elif isinstance(p, (WAVESTART, WAVEEND)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
elif hasattr(p, '_fields'):
filt = {'delta', 'encoding'} if not isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_L4)) else {'encoding'}
fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
for k in p._fields if not k.startswith('_') and k not in filt)
for k in p._fields if not k.startswith('_') and k not in {'delta', 'encoding'})
else: fields = ""
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name.replace('_L4', ''), 'white'))} {fields}"
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name, 'white'))} {fields}"
def print_packets(packets) -> None:
from tinygrad.helpers import getenv
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK",
"TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"} if not getenv("NOSKIP") else {"NOP"}
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK", "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"}
for p in packets:
if type(p).__name__.replace("_L4", "") not in skip: print(format_packet(p))
if type(p).__name__ not in skip: print(format_packet(p))
if __name__ == "__main__":
import sys, pickle
-161
View File
@@ -1,161 +0,0 @@
"""SQTT (SQ Thread Trace) packet decoder for CDNA/MI300 GPUs.
CDNA uses a completely different 16-bit header format from RDNA's nibble-based encoding.
"""
from __future__ import annotations
from typing import Iterator
from extra.assembly.amd.dsl import bits
from extra.assembly.amd.sqtt import PacketType
# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
class CDNA_DELTA(PacketType):
"""pkt_fmt=0: 16-bit timestamp delta packet"""
encoding = bits[3:0] == 0
delta = bits[11:4] # (data >> 4) & 0xff
unk_0 = bits[12:12] # (data >> 0xc) & 1
unk_1 = bits[15:13] # (data >> 0xd)
class CDNA_TIMESTAMP(PacketType):
"""pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
encoding = bits[3:0] == 1
unk_0 = bits[15:4]
timestamp = bits[63:16] # stored as (data_word >> 0x10) in low 46 bits of local_58
class CDNA_PKT_2(PacketType):
"""pkt_fmt=2: 64-bit packet (case 0x4)"""
encoding = bits[3:0] == 2
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[7:7] # (data >> 7) + 1 & 1
unk_padding = bits[63:8]
class CDNA_WAVESTART(PacketType):
"""pkt_fmt=3: 32-bit WAVESTART packet (case 0x8)"""
encoding = bits[3:0] == 3
unk_0 = bits[5:5] # (data >> 5) & 1
unk_1 = bits[9:6] # (data >> 6) & 0xf
wave = bits[13:10] # (data >> 10) & 0xf
simd = bits[15:14] # (data >> 0xe) & 3
cu = bits[17:16] # (data >> 0x10) & 3
unk_5 = bits[19:18] # (data >> 0x12) & 3
unk_6 = bits[28:22] # (data >> 0x16) & 0x7f
unk_padding = bits[31:29]
class CDNA_PKT_4(PacketType):
"""pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
encoding = bits[3:0] == 4
unk_0 = bits[5:5] # (data_word >> 5) & 1
unk_1 = bits[9:6] # (data_word >> 6) & 0xf
unk_2 = bits[13:10] # (data_word >> 10) & 0xf
unk_3 = bits[15:14] # (data_word >> 0xe)
class CDNA_PKT_5(PacketType):
"""pkt_fmt=5: 48-bit packet (case 0x10)"""
encoding = bits[3:0] == 5
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[7:7] # (data >> 7) + 1 & 1
unk_2 = bits[15:9] # (data >> 9) & 0x7f
unk_padding = bits[47:16]
class CDNA_WAVEEND(PacketType):
"""pkt_fmt=6: 16-bit WAVEEND packet (case 0x14, same as 0x8/0xc)"""
encoding = bits[3:0] == 6
unk_0 = bits[5:5] # (data_word >> 5) & 1
unk_1 = bits[9:6] # (data_word >> 6) & 0xf
wave = bits[13:10] # (data_word >> 10) & 0xf
simd = bits[15:14] # (data_word >> 0xe)
class CDNA_EXEC(PacketType):
"""pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
encoding = bits[3:0] == 10
unk_0 = bits[8:5] # (data_word >> 5) & 0xf
unk_1 = bits[10:9] # (data_word >> 9) & 3
unk_2 = bits[15:11] # (data_word >> 0xb)
class CDNA_PKT_11(PacketType):
"""pkt_fmt=11: 64-bit packet (case 0x28)"""
encoding = bits[3:0] == 11
unk_0 = bits[8:5] # (data_word >> 5) & 0xf
unk_1 = bits[10:9] # (data_word >> 9) & 3
unk_2 = bits[15:15] # (data_word >> 0xf) & 1
unk_padding = bits[63:16]
class CDNA_INST(PacketType):
"""pkt_fmt=13: 32-bit INST packet (case 0x30)"""
encoding = bits[3:0] == 13
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[9:8] # (data >> 8) & 3
unk_2 = bits[11:10] # (data >> 10) & 3
unk_3 = bits[13:12] # (data >> 0xc) & 3
unk_4 = bits[15:14] # (data >> 0xe) & 3
unk_5 = bits[19:18] # (data >> 0x12) & 3
unk_6 = bits[21:20] # (data >> 0x14) & 3
unk_7 = bits[23:22] # (data >> 0x16) & 3
unk_8 = bits[25:24] # (data >> 0x18) & 3
unk_9 = bits[27:26] # (data >> 0x1a) & 3
unk_padding = bits[31:28]
class CDNA_PKT_14(PacketType):
"""pkt_fmt=14: 64-bit packet (case 0x34)"""
encoding = bits[3:0] == 14
unk_0 = bits[5:5] # (data >> 5) & 1
unk_1 = bits[9:6] # (data >> 6) & 0xf
unk_2 = bits[11:10] # (data >> 10) & 3
unk_3 = bits[24:12] # (data >> 0xc) & 0x1fff
unk_4 = bits[37:25] # (data >> 0x19) & 0x1fff
unk_5 = bits[50:38] # (data >> 0x26) & 0x1fff
unk_6 = bits[51:51] # (data >> 0x33) & 1
unk_padding = bits[63:52]
class CDNA_PKT_15(PacketType):
"""pkt_fmt=15: 48-bit packet (case 0x38, same as 0x10)"""
encoding = bits[3:0] == 15
unk_0 = bits[6:5] # (data >> 5) & 3
unk_1 = bits[7:7] # (data >> 7) + 1 & 1
unk_2 = bits[15:9] # (data >> 9) & 0x7f
unk_padding = bits[47:16]
CDNA_PKT_TYPES: dict[int, type[PacketType]] = {
0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4,
5: CDNA_PKT_5, 6: CDNA_WAVEEND, 10: CDNA_EXEC, 11: CDNA_PKT_11, 13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
}
# Validate CDNA packet definitions
for pkt_fmt, pkt_cls in CDNA_PKT_TYPES.items():
assert pkt_cls.encoding.default == pkt_fmt, f"{pkt_cls.__name__} encoding {pkt_cls.encoding.default} != pkt_fmt {pkt_fmt}"
assert CDNA_PKT_SIZES[pkt_fmt] * 2 == pkt_cls._size_nibbles, f"{pkt_cls.__name__} size {pkt_cls._size_nibbles//2} != {CDNA_PKT_SIZES[pkt_fmt]}"
def decode(data: bytes) -> Iterator[PacketType]:
"""Decode CDNA SQTT blob using 16-bit header format."""
pos, time, ts_offset = 0, 0, None
while pos + 2 <= len(data):
header = int.from_bytes(data[pos:pos+2], 'little')
pkt_fmt = header & 0xf
pkt_size = CDNA_PKT_SIZES[pkt_fmt]
if pos + pkt_size > len(data): break
raw = int.from_bytes(data[pos:pos+pkt_size], 'little')
# pkt_fmt=0 has delta in bits[11:4], accumulate it
if pkt_fmt == 0: time += ((raw >> 4) & 0xff) * 4
# pkt_fmt=1 with unk_0=0 is absolute timestamp - use it to anchor time
if pkt_fmt == 1 and ((raw >> 4) & 0xfff) == 0:
abs_ts = raw >> 16
if ts_offset is None: ts_offset = abs_ts - time # first timestamp: save offset
else: time = ((abs_ts - ts_offset) & ~3) - 4 # subsequent: compute time, align to 4, subtract 4
pkt_cls = CDNA_PKT_TYPES[pkt_fmt]
yield pkt_cls.from_raw(raw, time)
pos += pkt_size
if __name__ == "__main__":
import sys, pickle
if len(sys.argv) < 2:
print("Usage: python sqtt_cdna.py <pkl_file>")
sys.exit(1)
with open(sys.argv[1], "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
for i, event in enumerate(sqtt_events):
print(f"\n=== event {i} ===")
for pkt in decode(event.blob):
print(f"{pkt._time:8}: {pkt}")
+1 -1
View File
@@ -6,7 +6,7 @@ from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd.sqtt import decode, print_packets, INST, VALUINST, IMMEDIATE, WAVESTART, WAVEEND, InstOp, PacketType, IMMEDIATE_MASK
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd import decode_inst
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP, s_endpgm
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
+23 -97
View File
@@ -6,9 +6,7 @@ from pathlib import Path
# Set AMD=1 before importing tinygrad
os.environ["AMD"] = "1"
from extra.assembly.amd.emu import run_asm as python_run_asm, decode_program
from extra.assembly.amd import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP, SOPPOp
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
if not REMU_PATH.exists():
@@ -66,55 +64,6 @@ def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size
return sum(times) / len(times)
def profile_instructions(kernel: bytes):
"""Profile individual instruction compile times."""
from extra.assembly.amd.emu import _get_runner, _canonical_runner_cache
from tinygrad.helpers import Context
_get_runner.cache_clear()
_canonical_runner_cache.clear()
results = []
i = 0
while i < len(kernel):
inst = decode_inst(kernel[i:])
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
inst_bytes = bytes(kernel[i:i + inst.size() + 4])
try: inst_str = repr(inst)
except Exception: inst_str = f"<{type(inst).__name__}>"
# Time the full compile (sink + render + compile)
start = time.perf_counter()
with Context(CCACHE=0):
runner, is_new = _get_runner(inst_bytes)
compile_time = time.perf_counter() - start
results.append({
'inst_str': inst_str + ('' if is_new else ' [CACHED]'),
'compile_ms': compile_time * 1000 if is_new else 0,
})
i += inst.size()
return sorted(results, key=lambda x: x['compile_ms'], reverse=True)
def benchmark_python_split(kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
"""Benchmark Python emulator with compile and execution times."""
from extra.assembly.amd.emu import _get_runner, _canonical_runner_cache
from tinygrad.helpers import Context
_get_runner.cache_clear()
_canonical_runner_cache.clear()
decode_program.cache_clear()
# Measure compile time (decode_program builds sinks, renders, and compiles)
compile_start = time.perf_counter()
with Context(CCACHE=0):
program = decode_program(kernel)
compile_time = time.perf_counter() - compile_start
n_compiled = len(_canonical_runner_cache)
# Execution time
exec_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, iterations)
return compile_time, exec_time, len(program), n_compiled
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
try:
@@ -170,36 +119,14 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
print(f" Error getting kernel: {e}")
return None
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "sin", "gelu", "matmul_small"]
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
def main():
import argparse
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
parser.add_argument("--profile", type=str, default=None, help="Profile instructions for a specific kernel (e.g. 'sin')")
parser.add_argument("--top", type=int, default=20, help="Number of top instructions to show in profile")
args = parser.parse_args()
# Profile mode: show individual instruction timing
if args.profile:
kernel_info = get_tinygrad_kernel(args.profile)
if kernel_info is None:
print(f"Failed to get kernel for '{args.profile}'")
return
kernel = kernel_info[0]
print(f"Profiling instructions for '{args.profile}' kernel...")
print("=" * 110)
results = profile_instructions(kernel)
print(f"{'Instruction':<90} {'Compile(ms)':>12}")
print("-" * 110)
for r in results[:args.top]:
inst = r['inst_str'][:87] + "..." if len(r['inst_str']) > 90 else r['inst_str']
print(f"{inst:<90} {r['compile_ms']:>12.3f}")
print("-" * 110)
total = sum(r['compile_ms'] for r in results)
print(f"{'TOTAL':<90} {total:>12.3f}")
return
rust_remu = get_rust_remu()
if rust_remu is None:
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
@@ -222,45 +149,44 @@ def main():
continue
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
# Benchmark Python emulator (must be first to measure compile time before cache is populated)
py_compile, py_exec, n_insts, n_compiled = benchmark_python_split(kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
n_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_compiled} unique) × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
set_valid_mem_ranges(ranges)
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
if py_compile is not None:
py_exec_rate = total_work / py_exec / 1e6
print(f" Compile: {py_compile*1000:8.3f} ms ({n_compiled} unique)")
print(f" Exec: {py_exec*1000:8.3f} ms ({py_exec_rate:7.2f} M ops/s)")
if 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_exec / rust_time if py_exec else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
speedup = py_time / rust_time if py_time else 0
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
results.append((op_name, n_insts, n_compiled, n_workgroups, py_compile, py_exec, rust_time))
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
# Summary table
print("\n" + "=" * 110)
print("\n" + "=" * 90)
print("SUMMARY")
print("=" * 110)
print(f"{'Name':<16} {'Insts':<6} {'Unique':<6} {'WGs':<5} {'Compile (ms)':<14} {'Exec (ms)':<12} {'Rust (ms)':<12} {'Speedup':<10}")
print("-" * 110)
print("=" * 90)
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
print("-" * 90)
for name, n_insts, n_compiled, n_wgs, py_compile, py_exec, rust_time in results:
compile_ms = f"{py_compile*1000:.3f}" if py_compile else "error"
exec_ms = f"{py_exec*1000:.3f}" if py_exec else "error"
for name, n_insts, n_wgs, py_time, rust_time in results:
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
if rust_time:
rust_ms = f"{rust_time*1000:.3f}"
speedup = f"{py_exec/rust_time:.1f}x" if py_exec else "N/A"
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
else:
rust_ms, speedup = "N/A", "N/A"
print(f"{name:<16} {n_insts:<6} {n_compiled:<6} {n_wgs:<5} {compile_ms:<14} {exec_ms:<12} {rust_ms:<12} {speedup:<10}")
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
if __name__ == "__main__":
main()
@@ -0,0 +1,196 @@
# Usability tests for the RDNA3 ASM DSL
# These tests demonstrate how the DSL *should* work for a good user experience
# Currently many of these tests fail - they document desired behavior
import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst, RawImm, SGPR, VGPR
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
disasm = inst.disasm()
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", disasm)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
class TestReprReadability(unittest.TestCase):
"""
Issue: repr() leaks internal RawImm type and omits zero-valued fields.
When you create v_mov_b32_e32(v[0], v[1]), the repr shows:
VOP1(op=1, src0=RawImm(257))
Problems:
1. vdst=v[0] is omitted because 0 is treated as "default"
2. src0 shows RawImm(257) instead of v[1]
3. User sees encoded values (257 = 256 + 1) instead of register names
Expected repr: VOP1(op=1, vdst=v[0], src0=v[1])
"""
def test_repr_shows_registers_not_raw_imm(self):
inst = v_mov_b32_e32(v[0], v[1])
# Should show v[1], not RawImm(257)
self.assertNotIn("RawImm", repr(inst), "repr should not expose RawImm internal type")
self.assertIn("v[1]", repr(inst), "repr should show register name")
def test_repr_includes_zero_dst(self):
inst = v_mov_b32_e32(v[0], v[1])
# v[0] is a valid destination register, should be shown
self.assertIn("vdst", repr(inst), "repr should include vdst even when 0")
def test_repr_roundtrip(self):
# repr should produce something that can be eval'd back
inst = v_mov_b32_e32(v[0], v[1])
# This would require repr to output valid Python, e.g.:
# "VOP1(op=VOP1Op.V_MOV_B32, vdst=v[0], src0=v[1])"
r = repr(inst)
# At minimum, it should be human-readable
self.assertIn("v[", r, "repr should show register syntax")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
class TestVOPDHelperSignature(unittest.TestCase):
"""
Issue: VOPD helper functions have confusing semantics.
v_dual_mul_f32 is defined as:
v_dual_mul_f32 = functools.partial(VOPD, VOPDOp.V_DUAL_MUL_F32)
This binds VOPDOp.V_DUAL_MUL_F32 to the FIRST positional arg of VOPD.__init__,
which is 'opx'. So v_dual_mul_f32 sets the X operation.
But then test_dual_mul in test_handwritten.py does:
v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], ...)
This passes V_DUAL_MUL_F32 as the SECOND positional arg (opy), making both
X and Y operations the same. This is confusing because:
1. The function name suggests it handles the X operation
2. But you still pass an opcode as the first arg (which becomes opy)
Expected: Either make the helper fully specify both ops, or make the
signature clearer about what the positional arg means.
"""
def test_vopd_helper_opy_should_be_required(self):
# Using only keyword args "works" but opy silently defaults to 0
inst = v_dual_mul_f32(vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32)
# Bug: opy defaults to 0 (V_DUAL_FMAC_F32) silently - should require explicit opy
# This test documents the bug - it should fail once fixed
self.assertNotEqual(inst.opy, VOPDOp.V_DUAL_FMAC_F32, "opy should not silently default to FMAC")
def test_vopd_helper_positional_arg_is_opy(self):
# The first positional arg after the partial becomes opy, not a second opx
inst = v_dual_mul_f32(VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
self.assertEqual(inst.opx, VOPDOp.V_DUAL_MUL_F32) # From partial
self.assertEqual(inst.opy, VOPDOp.V_DUAL_MOV_B32) # From first positional arg
class TestFieldAccessPreservesType(unittest.TestCase):
"""
Issue: Field access loses type information.
After creating an instruction, accessing fields returns encoded int values:
inst = v_mov_b32_e32(v[0], v[1])
inst.vdst # returns 0, not VGPR(0)
This makes it impossible to round-trip register types through field access.
"""
def test_vdst_returns_register(self):
inst = v_mov_b32_e32(v[5], v[1])
vdst = inst.vdst
# Should return a VGPR, not an int
self.assertIsInstance(vdst, (VGPR, int), "vdst should return VGPR or at least be usable")
# Ideally: self.assertIsInstance(vdst, VGPR)
def test_src_returns_register_for_vgpr_source(self):
inst = v_mov_b32_e32(v[0], v[1])
# src0 is encoded as 257 (256 + 1 for v1)
# Ideally it should decode back to v[1]
src0_raw = inst._values.get('src0')
# Currently returns RawImm(257), should return VGPR(1) or similar
self.assertNotIsInstance(src0_raw, RawImm, "source should not be RawImm internally")
class TestArgumentDiscoverability(unittest.TestCase):
"""
Issue: No clear signature for positional arguments.
inspect.signature(s_load_b128) shows: (*args, literal=None, **kwargs)
Users have no way to know the argument order without reading source code.
The order is implicitly defined by the class field definition order.
Possible fixes:
1. Add explicit parameter names to functools.partial
2. Generate type stubs with proper signatures
3. Add docstrings listing the expected arguments
"""
def test_signature_has_named_params(self):
import inspect
sig = inspect.signature(s_load_b128)
params = list(sig.parameters.keys())
# Currently: ['args', 'literal', 'kwargs'] (from *args, literal=None, **kwargs)
# Expected: something like ['sdata', 'sbase', 'soffset', 'offset', 'literal']
self.assertIn('sdata', params, "signature should show field names")
class TestSpecialConstants(unittest.TestCase):
"""
Issue: NULL and other constants are IntEnum values that might be confusing.
NULL = SrcEnum.NULL = 124, but users might expect NULL to be a special object
that clearly represents "no register" rather than a magic number.
"""
def test_null_has_clear_repr(self):
# NULL should have a clear string representation
self.assertIn("NULL", str(NULL) or repr(NULL), "NULL should be clearly identifiable")
def test_null_is_distinguishable_from_int(self):
# NULL should be distinguishable from the raw integer 124
self.assertNotEqual(type(NULL), int, "NULL should not be plain int")
if __name__ == "__main__":
unittest.main()
+43 -2
View File
@@ -5,7 +5,6 @@ from dataclasses import dataclass
@dataclass
class KernelInfo:
code: bytes
src: str
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
@@ -35,4 +34,46 @@ TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items
def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
def get_mattr(arch:str) -> str:
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
# ═══════════════════════════════════════════════════════════════════════════════
# EXECUTION CONTEXT (for testing compiled pseudocode)
# ═══════════════════════════════════════════════════════════════════════════════
class ExecContext:
"""Context for running compiled pseudocode in tests."""
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
self.D0, self.D1 = Reg(d0), Reg(0)
self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
self.lane, self.laneId, self.literal = lane, lane, literal
self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
self.VGPR = vgprs if vgprs is not None else {}
self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
def run(self, code: str):
"""Execute compiled code."""
import extra.assembly.amd.pcode as pcode
ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
# Also include underscore-prefixed helpers that compiled pseudocode uses
for k in ['_pack', '_pack32']:
if hasattr(pcode, k): ns[k] = getattr(pcode, k)
ns.update({
'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
'tmp': self.tmp, 'saveexec': self.saveexec,
'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
})
exec(code, ns)
def _sync(ctx_reg, ns_val):
if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
else: ctx_reg._val = int(ns_val) & self._MASK64
for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
+1 -2
View File
@@ -141,7 +141,7 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
rsrc2 = 0x19c | (128 << 15)
scratch_size = 0x10000 # 64KB per lane, matches .amdhsa_private_segment_fixed_size in run_program_hw
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2, scratch_size)
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2)
assert result == 0, f"run_asm failed with {result}"
return parse_output(bytes(out_buf), n_lanes)
@@ -204,7 +204,6 @@ amdhsa.kernels:
prg = AMDProgram(dev, "test", lib)
out_gpu = dev.allocator.alloc(OUT_BYTES)
assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
out_buf = bytearray(OUT_BYTES)
-44
View File
@@ -138,50 +138,6 @@ class TestDS2AddrMore(unittest.TestCase):
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSB96(unittest.TestCase):
"""Tests for DS_STORE_B96 and DS_LOAD_B96 (96-bit / 3 dwords)."""
def test_ds_store_load_b96(self):
"""DS_STORE_B96 stores 3 VGPRs, DS_LOAD_B96 loads them back."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
ds_store_b96(addr=v[10], data0=v[0:2]),
s_waitcnt(lgkmcnt=0),
ds_load_b96(addr=v[10], vdst=v[4:6]),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
def test_ds_store_b96_with_offset(self):
"""DS_STORE_B96 with non-zero offset."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
DS(DSOp.DS_STORE_B96, addr=v[10], data0=v[0:2], offset0=12),
s_waitcnt(lgkmcnt=0),
DS(DSOp.DS_LOAD_B96, addr=v[10], vdst=v[4:6], offset0=12),
s_waitcnt(lgkmcnt=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
class TestDSB128(unittest.TestCase):
"""Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
-152
View File
@@ -523,157 +523,5 @@ class TestD16HiLoads(unittest.TestCase):
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
class TestGlobalOffset(unittest.TestCase):
"""Tests for GLOBAL instructions with different offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load/store wrong data.
"""
def test_global_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
# Load from offset 100 -> should get 0xAAAAAAAA
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[3], saddr=SrcEnum.NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], v[4]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_store_different_offsets(self):
"""Store to two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0x11111111 at offset 300
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=300),
# Store 0x22222222 at offset 400
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b32(addr=v[0:1], data=v[3], saddr=SrcEnum.NULL, offset=400),
s_waitcnt(vmcnt=0),
# Load back to verify
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=300),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[5], saddr=SrcEnum.NULL, offset=400),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x11111111, f"offset 300: expected 0x11111111, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0x22222222, f"offset 400: expected 0x22222222, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_no_saddr(self):
"""Test negative offset without saddr (VGPR pair for address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from vaddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from vaddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt(vmcnt=0),
# vaddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[6], saddr=SrcEnum.NULL, offset=-100),
# vaddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[7], saddr=SrcEnum.NULL, offset=-100),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[4], 0),
v_mov_b32_e32(v[5], 0),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_with_saddr(self):
"""Test negative offset with saddr (SGPR pair for base address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from saddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from saddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
# saddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[6], saddr=s[4:5], offset=-100),
# saddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[7], saddr=s[4:5], offset=-100),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
if __name__ == '__main__':
unittest.main()
-355
View File
@@ -1,355 +0,0 @@
"""Tests for SCRATCH instructions - scratch (private) memory operations.
Includes: scratch_load_*, scratch_store_*
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
class TestScratchStore(unittest.TestCase):
"""Tests for SCRATCH store instructions."""
def test_scratch_store_b32_basic(self):
"""SCRATCH_STORE_B32 stores 32-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
# Store via scratch
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Load back via scratch
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
def test_scratch_store_b64_basic(self):
"""SCRATCH_STORE_B64 stores 64-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
s_mov_b32(s[5], 0xCAFEBABE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
scratch_store_b64(addr=v[0], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b64(addr=v[0], vdst=v[4:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
def test_scratch_store_b8_basic(self):
"""SCRATCH_STORE_B8 stores single byte to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# First store full word
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Store single byte
v_mov_b32_e32(v[2], 0x42),
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Load back
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Only byte 0 should change from 0xEF to 0x42
self.assertEqual(st.vgpr[0][0], 0xDEADBE42)
def test_scratch_store_b16_basic(self):
"""SCRATCH_STORE_B16 stores 16-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADCAFE)
class TestScratchLoad(unittest.TestCase):
"""Tests for SCRATCH load instructions."""
def test_scratch_load_b96(self):
"""SCRATCH_LOAD_B96 loads 96-bit value correctly."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[4]),
s_mov_b32(s[4], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[4]),
scratch_store_b96(addr=v[0], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b96(addr=v[0], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[5]),
v_mov_b32_e32(v[1], v[6]),
v_mov_b32_e32(v[2], v[7]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][2], 0xCCCCCCCC)
def test_scratch_load_b128(self):
"""SCRATCH_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
s_mov_b32(s[4], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[4]),
s_mov_b32(s[4], 0x12345678),
v_mov_b32_e32(v[4], s[4]),
s_mov_b32(s[4], 0x9ABCDEF0),
v_mov_b32_e32(v[5], s[4]),
scratch_store_b128(addr=v[0], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_b128(addr=v[0], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[2], v[8]),
v_mov_b32_e32(v[3], v[9]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][2], 0x12345678)
self.assertEqual(st.vgpr[0][3], 0x9ABCDEF0)
def test_scratch_load_u8(self):
"""SCRATCH_LOAD_U8 loads unsigned byte with zero extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADBEAB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_u8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAB)
def test_scratch_load_i8(self):
"""SCRATCH_LOAD_I8 loads signed byte with sign extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0x80), # -128 as signed byte
v_mov_b32_e32(v[2], s[4]),
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_i8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFF80)
def test_scratch_load_u16(self):
"""SCRATCH_LOAD_U16 loads unsigned 16-bit with zero extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADCAFE),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_u16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xCAFE)
def test_scratch_load_i16(self):
"""SCRATCH_LOAD_I16 loads signed 16-bit with sign extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0x8000), # -32768 as signed 16-bit
v_mov_b32_e32(v[2], s[4]),
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
scratch_load_i16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFF8000)
class TestScratchSVE(unittest.TestCase):
"""Tests for SCRATCH SVE (Scratch VGPR Enable) bit behavior."""
def test_scratch_sve_zero_ignores_vaddr(self):
"""With SVE=0, VADDR should be ignored in address calculation."""
TEST_OFFSET = 256
# Store a marker value at offset 256 (where SVE=0 should go)
# Then set v[0] to a non-zero value (100) and store via scratch with SVE=0
# If SVE=0 is handled correctly, the VADDR (100) should be IGNORED,
# and the store should go to offset 256, not 256+100=356
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt(vmcnt=0),
# Now set v[0] to 100 (non-zero) and store 0xBBBBBBBB with SVE=0
# With SVE=0, v[0] should be IGNORED, so this should overwrite offset 256
v_mov_b32_e32(v[0], 100),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt(vmcnt=0),
# Load back from offset 256 (with v[0]=0) - should get 0xBBBBBBBB
v_mov_b32_e32(v[0], 0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# If SVE=0 works correctly, v[0] should be 0xBBBBBBBB (the second store overwrote the first)
# If SVE=0 is wrong (VADDR used), v[0] would be 0xAAAAAAAA (stores went to different locations)
self.assertEqual(st.vgpr[0][0], 0xBBBBBBBB, "SVE=0 should ignore VADDR, both stores should go to same location")
def test_scratch_sve_one_uses_vaddr(self):
"""With SVE=1, VADDR should be used as offset in address calculation."""
TEST_OFFSET = 256
# Store at offset 256 with v[0]=0, then store at offset 256 with v[0]=100 and SVE=1
# With SVE=1, the second store should go to 256+100=356, not 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt(vmcnt=0),
# Now set v[0] to 100 and store 0xBBBBBBBB with SVE=1
# With SVE=1, v[0] IS used, so this should go to offset 256+100=356
v_mov_b32_e32(v[0], 100),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt(vmcnt=0),
# Load back from offset 256 (with v[0]=0) - should still be 0xAAAAAAAA
v_mov_b32_e32(v[0], 0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# If SVE=1 works correctly, v[0] should be 0xAAAAAAAA (stores went to different locations)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, "SVE=1 should use VADDR, stores should go to different locations")
class TestScratchMultiLane(unittest.TestCase):
"""Tests for SCRATCH operations with multiple lanes."""
def test_scratch_store_load_multi_lane(self):
"""SCRATCH store/load works correctly with multiple lanes (private per-lane memory)."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt(lgkmcnt=0),
# Each lane stores its lane ID
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[2], v[255]), # v[255] has packed workitem IDs, low 10 bits = x
v_and_b32_e32(v[2], 0x3FF, v[2]), # extract lane ID
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
# Load back
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=4)
# Each lane should have loaded its own lane ID
for lane in range(4):
self.assertEqual(st.vgpr[lane][0], lane, f"Lane {lane} should have value {lane}")
if __name__ == '__main__':
unittest.main()
-448
View File
@@ -1,448 +0,0 @@
"""Tests for SMEM instructions - scalar memory operations.
Includes: s_load_b32, s_load_b64, s_load_b128, s_load_b256, s_load_b512
Tests both immediate and register offset addressing modes.
"""
import unittest
from extra.assembly.amd.test.hw.helpers import *
# Use offset into output buffer for test data (output buffer is 2124 bytes)
TEST_OFFSET = 2000
# Cache invalidation sequence for scalar loads after vector stores
# s_wait_idle waits for all outstanding memory operations including cache flushes
CACHE_INV = [s_gl1_inv(), s_dcache_inv(), s_wait_idle()]
class TestSLoadRegisterOffset(unittest.TestCase):
"""Tests for s_load with register offset (soffset field).
Bug: s_load_b32(s[dst], s[base:base+1], s[off]) ignores the register offset
and only uses the immediate offset field. This causes incorrect memory loads
when the offset comes from a register.
"""
def test_s_load_b32_register_offset_basic(self):
"""s_load_b32 with register offset should load from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values to output buffer: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[4], 0xAAAAAAAA),
s_mov_b32(s[5], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Now test s_load with register offset
# Put offset value in s[4]: offset = 4 bytes (1 dword)
s_mov_b32(s[4], 4),
# Load from out_ptr + TEST_OFFSET + s[4] (should load 0xBBBBBBBB)
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB,
f"s_load with reg offset 4 should load 0xBBBBBBBB: s[5]=0x{st.sgpr[5]:08x}")
def test_s_load_b32_register_offset_different_from_immediate(self):
"""s_load_b32 with register offset loads different data than immediate offset 0."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[4], 0xAAAAAAAA),
s_mov_b32(s[5], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with immediate offset 0
s_load_b32(s[5], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Load with register offset 4
s_mov_b32(s[4], 4),
s_load_b32(s[6], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# s[5] has dword at offset 0 (0xAAAAAAAA), s[6] has dword at offset 4 (0xBBBBBBBB)
self.assertEqual(st.sgpr[5], 0xAAAAAAAA)
self.assertEqual(st.sgpr[6], 0xBBBBBBBB)
self.assertNotEqual(st.sgpr[5], st.sgpr[6],
f"s_load with reg offset 4 should load different value than offset 0: "
f"s[5]=0x{st.sgpr[5]:08x}, s[6]=0x{st.sgpr[6]:08x}")
def test_s_load_b32_register_offset_same_as_dst(self):
"""s_load_b32 where soffset register is same as destination.
This is the exact pattern that exposes the bug:
s_load_b32(s[8], s[2:3], s[8])
The offset should be read BEFORE the destination is overwritten.
"""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[6], 0xAAAAAAAA),
s_mov_b32(s[7], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[6]),
v_mov_b32_e32(v[3], s[7]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Set up s[4] = 4 (offset in bytes)
s_mov_b32(s[4], 4),
# Load using s[4] as both offset and destination
# Should load from base + 4, then store result in s[4]
s_load_b32(s[4], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Also load with immediate offset 4 for comparison
s_load_b32(s[5], s[2:3], NULL, offset=TEST_OFFSET+4),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# s[4] and s[5] should have the same value (both loaded from offset 4 = 0xBBBBBBBB)
self.assertEqual(st.sgpr[4], 0xBBBBBBBB)
self.assertEqual(st.sgpr[4], st.sgpr[5],
f"s_load with reg offset s[4]=4 should match immediate offset=4: "
f"s[4]=0x{st.sgpr[4]:08x}, s[5]=0x{st.sgpr[5]:08x}")
def test_s_load_b32_register_offset_zero(self):
"""s_load_b32 with register offset = 0 should be same as immediate offset 0."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test value: 0xDEADBEEF at offset
s_mov_b32(s[7], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[7]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 0
s_mov_b32(s[4], 0),
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Load with immediate offset 0
s_load_b32(s[6], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[5], 0xDEADBEEF)
self.assertEqual(st.sgpr[5], st.sgpr[6],
f"s_load with reg offset 0 should match immediate offset 0: "
f"s[5]=0x{st.sgpr[5]:08x}, s[6]=0x{st.sgpr[6]:08x}")
def test_s_load_b32_register_plus_immediate_offset(self):
"""s_load_b32 with both register and immediate offset should add them."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[8], 0xAAAAAAAA),
s_mov_b32(s[9], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[8]),
v_mov_b32_e32(v[3], s[9]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# reg offset = 4, imm offset = 0 -> total offset = 4
s_mov_b32(s[4], 4),
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# reg offset = 0, imm offset = 4 -> total offset = 4
s_mov_b32(s[6], 0),
s_load_b32(s[7], s[2:3], s[6], offset=TEST_OFFSET+4),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Both should load from offset 4 (0xBBBBBBBB)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB)
self.assertEqual(st.sgpr[7], 0xBBBBBBBB)
self.assertEqual(st.sgpr[5], st.sgpr[7],
f"reg_off=4 + imm_off=0 should equal reg_off=0 + imm_off=4: "
f"s[5]=0x{st.sgpr[5]:08x}, s[7]=0x{st.sgpr[7]:08x}")
class TestSLoadMultiDword(unittest.TestCase):
"""Tests for multi-dword s_load with register offset."""
def test_s_load_b64_register_offset(self):
"""s_load_b64 with register offset should load 2 dwords from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA, 0xBBBBBBBB at offset
s_mov_b32(s[10], 0xAAAAAAAA),
s_mov_b32(s[11], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[10]),
v_mov_b32_e32(v[3], s[11]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 0
s_mov_b32(s[4], 0),
s_load_b64(s[6:7], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Compare with immediate offset
s_load_b64(s[8:9], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], 0xAAAAAAAA)
self.assertEqual(st.sgpr[7], 0xBBBBBBBB)
self.assertEqual(st.sgpr[6], st.sgpr[8])
self.assertEqual(st.sgpr[7], st.sgpr[9])
def test_s_load_b128_register_offset(self):
"""s_load_b128 with register offset should load 4 dwords from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
# Store test values: 0xAAAAAAAA, 0xBBBBBBBB, 0xCCCCCCCC, 0xDDDDDDDD at offset
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[14], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[14], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_mov_b32(s[14], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+8),
s_mov_b32(s[14], 0xDDDDDDDD),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+12),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 0 (s_load_b128 requires 4-aligned dest: s[4], s[8], s[12], ...)
s_mov_b32(s[15], 0),
s_load_b128(s[4:7], s[2:3], s[15], offset=TEST_OFFSET),
s_waitcnt(0),
# Compare with immediate offset
s_load_b128(s[8:11], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB)
self.assertEqual(st.sgpr[6], 0xCCCCCCCC)
self.assertEqual(st.sgpr[7], 0xDDDDDDDD)
self.assertEqual(st.sgpr[4], st.sgpr[8])
self.assertEqual(st.sgpr[5], st.sgpr[9])
class TestSLoadLarge(unittest.TestCase):
"""Tests for large s_load operations (s_load_b256, s_load_b512)."""
def test_s_load_b256_basic(self):
"""s_load_b256 loads 8 consecutive dwords."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 8 test values
s_mov_b32(s[20], 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[20], 0x22222222),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_mov_b32(s[20], 0x33333333),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+8),
s_mov_b32(s[20], 0x44444444),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+12),
s_mov_b32(s[20], 0x55555555),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+16),
s_mov_b32(s[20], 0x66666666),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+20),
s_mov_b32(s[20], 0x77777777),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+24),
s_mov_b32(s[20], 0x88888888),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+28),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load all 8 dwords with s_load_b256
s_load_b256(s[4:11], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111)
self.assertEqual(st.sgpr[5], 0x22222222)
self.assertEqual(st.sgpr[6], 0x33333333)
self.assertEqual(st.sgpr[7], 0x44444444)
self.assertEqual(st.sgpr[8], 0x55555555)
self.assertEqual(st.sgpr[9], 0x66666666)
self.assertEqual(st.sgpr[10], 0x77777777)
self.assertEqual(st.sgpr[11], 0x88888888)
def test_s_load_b512_basic(self):
"""s_load_b512 loads 16 consecutive dwords."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 16 test values (use a pattern: 0x10, 0x20, ..., 0x100)
*[instr for i in range(16) for instr in [
s_mov_b32(s[20], (i + 1) * 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET + i * 4),
]],
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load all 16 dwords with s_load_b512
s_load_b512(s[64:79], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(lgkmcnt=0),
# Copy results to lower regs for verification (since st.sgpr only has 16 regs in test)
s_mov_b32(s[4], s[64]),
s_mov_b32(s[5], s[65]),
s_mov_b32(s[6], s[78]),
s_mov_b32(s[7], s[79]),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111, "first dword")
self.assertEqual(st.sgpr[5], 0x22222222, "second dword")
self.assertEqual(st.sgpr[6], 0xFFFFFFFF & (15 * 0x11111111), "15th dword")
self.assertEqual(st.sgpr[7], 0xFFFFFFFF & (16 * 0x11111111), "16th dword")
def test_s_load_b256_with_register_offset(self):
"""s_load_b256 with register offset should add reg offset to address."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store pattern at TEST_OFFSET+8: skip first 2 dwords
*[instr for i in range(8) for instr in [
s_mov_b32(s[20], (i + 1) * 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET + 8 + i * 4),
]],
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load with register offset 8
s_mov_b32(s[20], 8),
s_load_b256(s[4:11], s[2:3], s[20], offset=TEST_OFFSET),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111, "first dword at offset+8")
self.assertEqual(st.sgpr[5], 0x22222222, "second dword at offset+8")
self.assertEqual(st.sgpr[11], 0x88888888, "last dword at offset+8")
class TestSLoadOffset(unittest.TestCase):
"""Tests for s_load with different immediate offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load wrong data.
"""
def test_s_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# Load from offset 100 -> should get 0xAAAAAAAA
s_load_b32(s[4], s[2:3], NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
s_load_b32(s[5], s[2:3], NULL, offset=200),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.sgpr[4]:08x}")
self.assertEqual(st.sgpr[5], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.sgpr[5]:08x}")
def test_s_load_negative_offset(self):
"""Test negative offset (21-bit signed).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from base+200 -> should get 0xAAAA.
Load with offset -100 from base+300 -> should get 0xBBBB."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt(lgkmcnt=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[8], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[8]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[8], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[8]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt(vmcnt=0),
*CACHE_INV,
# base+200, load with offset -100 -> should get value at 100
s_add_u32(s[6], s[2], 200),
s_addc_u32(s[7], s[3], 0),
s_load_b32(s[4], s[6:7], NULL, offset=-100),
# base+300, load with offset -100 -> should get value at 200
s_add_u32(s[6], s[2], 300),
s_addc_u32(s[7], s[3], 0),
s_load_b32(s[5], s[6:7], NULL, offset=-100),
s_waitcnt(lgkmcnt=0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[6], 0),
s_mov_b32(s[7], 0),
s_mov_b32(s[8], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.sgpr[4]:08x}")
self.assertEqual(st.sgpr[5], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.sgpr[5]:08x}")
if __name__ == '__main__':
unittest.main()
-267
View File
@@ -619,272 +619,5 @@ class Test64BitCompare(unittest.TestCase):
self.assertEqual(st.sgpr[4], 1)
class TestSOPPNop(unittest.TestCase):
"""Tests for S_NOP and other SOPP instructions with expression-based for loops.
S_NOP's pcode uses 'for i in 0U : SIMM16.u16[3 : 0].u32 do' which requires
the parser to handle non-constant loop bounds.
"""
def test_s_nop_basic(self):
"""S_NOP executes without side effects."""
# S_NOP with immediate 0 should just do nothing
instructions = [
s_mov_b32(s[0], 42),
s_nop(0), # nop with simm16=0
s_mov_b32(s[1], 100),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 42)
self.assertEqual(st.sgpr[1], 100)
def test_s_nop_with_count(self):
"""S_NOP with count parameter executes multiple nops."""
# S_NOP with immediate 3 should execute 4 nops (0:3 inclusive)
instructions = [
s_mov_b32(s[0], 1),
s_nop(3), # nop with simm16=3 -> 4 iterations
s_add_u32(s[0], s[0], 1),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 2)
class TestNullRegister(unittest.TestCase):
"""Tests for NULL register (124) behavior - writes should be discarded, reads return 0."""
def test_s_mov_b32_from_null(self):
"""S_MOV_B32 from NULL should read as 0."""
instructions = [
s_mov_b32(s[0], 0xDEADBEEF), # Set s[0] to sentinel
s_mov_b32(s[0], NULL), # Read from NULL - should be 0
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 0)
def test_s_add_u32_with_null_src(self):
"""S_ADD_U32 with NULL as source should use 0."""
instructions = [
s_mov_b32(s[0], 100),
s_add_u32(s[1], s[0], NULL), # 100 + 0 = 100
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[1], 100)
def test_s_mov_b32_to_null(self):
"""S_MOV_B32 to NULL (sdst=124) should discard the write."""
instructions = [
s_mov_b32(s[0], 0xDEADBEEF), # Set s[0] to sentinel
s_mov_b32(NULL, 42), # Write to NULL - should be discarded
# s[0] should still be 0xDEADBEEF since NULL write doesn't affect it
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0], 0xDEADBEEF)
def test_s_add_u32_to_null(self):
"""S_ADD_U32 with sdst=NULL should discard result but still set SCC."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 1),
s_add_u32(NULL, s[0], s[1]), # overflow, write to NULL
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
# SCC should still be set from overflow even though result was discarded
self.assertEqual(st.sgpr[2], 1)
self.assertEqual(st.scc, 1)
def test_s_and_b32_to_null(self):
"""S_AND_B32 with sdst=NULL should discard result but still set SCC."""
instructions = [
s_mov_b32(s[0], 0xFF00FF00),
s_mov_b32(s[1], 0x0F0F0F0F),
s_and_b32(NULL, s[0], s[1]), # result=0x0F000F00, non-zero so SCC=1
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 1) # SCC=1 because result was non-zero
self.assertEqual(st.scc, 1)
def test_s_or_b32_to_null_zero_result(self):
"""S_OR_B32 with sdst=NULL and zero result should set SCC=0."""
instructions = [
s_mov_b32(s[0], 0),
s_mov_b32(s[1], 0),
s_or_b32(NULL, s[0], s[1]), # result=0, so SCC=0
s_cselect_b32(s[2], 1, 0), # capture SCC
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[2], 0) # SCC=0 because result was zero
self.assertEqual(st.scc, 0)
class Test64BitSOP1InlineConstants(unittest.TestCase):
"""Tests for 64-bit SOP1 instructions with inline constants.
Regression tests for bug where rsrc_dyn didn't properly handle 64-bit
inline constants, incorrectly duplicating lo bits to hi instead of
zero/sign-extending.
"""
def test_s_mov_b64_inline_0(self):
"""S_MOV_B64 with inline constant 0."""
instructions = [
s_mov_b64(s[0:1], 0),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_16(self):
"""S_MOV_B64 with inline constant 16 should set lo=16, hi=0."""
instructions = [
s_mov_b64(s[0:1], 16),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_64(self):
"""S_MOV_B64 with inline constant 64 (max positive)."""
instructions = [
s_mov_b64(s[0:1], 64),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 64)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_inline_neg1(self):
"""S_MOV_B64 with inline constant -1 should sign-extend."""
instructions = [
s_mov_b64(s[0:1], -1),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
def test_s_mov_b64_inline_neg16(self):
"""S_MOV_B64 with inline constant -16 should sign-extend."""
instructions = [
s_mov_b64(s[0:1], -16),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFF0)
self.assertEqual(st.vgpr[0][1], 0xFFFFFFFF)
def test_s_mov_b64_float_const_1_0(self):
"""S_MOV_B64 with float inline constant 1.0 - casts F32 to F64."""
instructions = [
s_mov_b64(s[0:1], 1.0), # inline constant 242 (1.0f)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
# Hardware casts F32 to F64: 1.0f64 = 0x3FF0000000000000
self.assertEqual(st.vgpr[0][0], 0x00000000) # lo
self.assertEqual(st.vgpr[0][1], 0x3FF00000) # hi
def test_s_or_b64_inline_constant(self):
"""S_OR_B64 with 64-bit inline constant."""
instructions = [
s_mov_b64(s[0:1], 0),
s_or_b64(s[2:3], s[0:1], 16),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_and_b64_inline_constant(self):
"""S_AND_B64 with 64-bit inline constant."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 0xFFFFFFFF),
s_and_b64(s[2:3], s[0:1], 16),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.vgpr[0][1], 0)
class Test64BitSOPLiterals(unittest.TestCase):
"""Tests for 64-bit SOP instructions with 32-bit literals.
Tests the behavior when a 64-bit SOP instruction uses a 32-bit literal
(offset 255 in instruction encoding). The literal is zero-extended to 64 bits.
"""
def test_s_mov_b64_literal(self):
"""S_MOV_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b64(s[0:1], 0x12345678), # literal > 64, uses literal encoding
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_or_b64_literal(self):
"""S_OR_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b64(s[0:1], 0),
s_or_b64(s[2:3], s[0:1], 0x12345678), # literal
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_and_b64_literal(self):
"""S_AND_B64 with 32-bit literal value - zero-extended to 64 bits."""
instructions = [
s_mov_b32(s[0], 0xFFFFFFFF),
s_mov_b32(s[1], 0xFFFFFFFF),
s_and_b64(s[2:3], s[0:1], 0x12345678), # literal
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x12345678)
self.assertEqual(st.vgpr[0][1], 0)
def test_s_mov_b64_literal_negative(self):
"""S_MOV_B64 with 0xFFFFFFFF literal - zero-extended (not sign-extended)."""
instructions = [
s_mov_b64(s[0:1], 0xFFFFFFFF), # -1 as 32-bit, but zero-extended to 64-bit
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][1], 0) # zero-extended, not sign-extended
def test_s_mov_b64_literal_high_bit(self):
"""S_MOV_B64 with 0x80000000 literal - zero-extended (not sign-extended)."""
instructions = [
s_mov_b64(s[0:1], 0x80000000), # high bit set, but zero-extended
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x80000000)
self.assertEqual(st.vgpr[0][1], 0) # zero-extended, not sign-extended
if __name__ == '__main__':
unittest.main()
-118
View File
@@ -237,32 +237,6 @@ class TestF16Ops(unittest.TestCase):
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
def test_v_max_f16_basic(self):
"""V_MAX_F16 returns the maximum of two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_max_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected 0x4000 (f16 2.0), got 0x{result:04x}")
def test_v_min_f16_basic(self):
"""V_MIN_F16 returns the minimum of two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_min_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected 0x3c00 (f16 1.0), got 0x{result:04x}")
def test_v_fmaak_f16_basic(self):
"""V_FMAAK_F16: d = a * b + K."""
instructions = [
@@ -836,81 +810,6 @@ class TestCarryOps(unittest.TestCase):
self.assertEqual(st.vgpr[0][2], 0) # Overflowed to 0
self.assertEqual(st.vcc, 1) # Carry out
def test_v_add_co_ci_u32_clears_carry(self):
"""V_ADD_CO_CI_U32: VCC must be updated even when no carry is generated.
This tests the case where VCC=1 going in (carry-in consumed) but the addition
does not overflow, so VCC must be cleared to 0.
Regression test for: VCC not being written by v_add_co_ci_u32_e32.
"""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 3) # 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out - VCC must be cleared
def test_v_add_co_ci_u32_multilane_clears_vcc(self):
"""V_ADD_CO_CI_U32 with multiple lanes: VCC bits must be updated per-lane.
When VCC has multiple bits set (one per active lane), and the addition doesn't
overflow for any lane, all VCC bits must be cleared.
Regression test for: VCC not being written by v_add_co_ci_u32_e32 in multi-lane case.
"""
instructions = [
s_mov_b32(VCC_LO, 0b11), # VCC = 0b11 (lanes 0,1 have carry-in)
v_mov_b32_e32(v[0], 1), # S0 = 1 for all lanes
v_mov_b32_e32(v[1], 1), # S1 = 1 for all lanes
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vgpr[0][2], 3) # lane 0: 1 + 1 + 1 = 3
self.assertEqual(st.vgpr[1][2], 3) # lane 1: 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out for any lane - all VCC bits must be cleared
def test_v_add_co_ci_u32_preserves_inactive_vcc_bits(self):
"""V_ADD_CO_CI_U32: VCC carry-out overwrites entire VCC register.
VOP2 carry instructions write ALL VCC bits based on carry-out, clearing
bits for lanes that don't overflow regardless of EXEC mask.
Note: This differs from VOPC which only writes active lane bits.
"""
instructions = [
s_mov_b32(VCC_LO, 0x00010000), # VCC bit 16 set
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 0 = 2 (no carry)
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vgpr[0][2], 2) # lane 0: 1 + 1 + 0 = 2
# VCC should be completely cleared (all lanes have no carry-out)
self.assertEqual(st.vcc, 0)
def test_v_add_co_ci_u32_all_lanes_same_result(self):
"""V_ADD_CO_CI_U32: all active lanes should produce the same result.
When the same constant inputs are used across all lanes, each lane should
compute the same result and write to its own VGPR slot.
Regression test for: VGPR writes not happening for all lanes.
"""
instructions = [
s_mov_b32(VCC_LO, 0), # No carry-in
v_mov_b32_e32(v[0], 3), # inline constant 3
v_mov_b32_e32(v[1], 5), # value 5
v_add_co_ci_u32_e32(v[1], 3, v[1]), # v[1] = 3 + v[1] + 0 = 3 + 5 = 8
]
st = run_program(instructions, n_lanes=4)
# All 4 lanes should have v[1] = 8
for lane in range(4):
self.assertEqual(st.vgpr[lane][1], 8, f"lane {lane} should have v[1]=8")
def test_v_sub_co_ci_u32_no_borrow(self):
"""V_SUB_CO_CI_U32: D0 = S0 - S1 - VCC_IN, when VCC_IN=0."""
instructions = [
@@ -961,23 +860,6 @@ class TestCarryOps(unittest.TestCase):
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.sgpr[10], 0) # No carry out
def test_v_add_co_ci_u32_vop3sd_null_sdst(self):
"""VOP3SD V_ADD_CO_CI_U32 with sdst=NULL: carry output is discarded.
When sdst=NULL (register 124), the carry-out should NOT be written anywhere.
We verify this by checking that VCC (which we set to a sentinel value) is unchanged.
"""
instructions = [
s_mov_b32(VCC_LO, 0xDEADBEEF), # Sentinel value in VCC
s_mov_b32(s[6], 0), # carry-in = 0
# VOP3SD with NULL sdst: carry-out should be discarded
# Uses 0xFFFFFFFF + 1 + 0 = 0 with carry-out=1, but carry should not be written
v_add_co_ci_u32(v[0], NULL, 0xFFFFFFFF, 1, s[6]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0) # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
self.assertEqual(st.vcc, 0xDEADBEEF) # VCC unchanged - carry was discarded
if __name__ == '__main__':
unittest.main()
+1 -145
View File
@@ -1077,7 +1077,7 @@ class TestF64Ops(unittest.TestCase):
s_mov_b32(s[3], one_f64 >> 32),
v_mov_b32_e32(v[2], s[2]),
v_mov_b32_e32(v[3], s[3]),
VOP3SD(VOP3SDOp.V_DIV_SCALE_F64, vdst=v[4:5], sdst=s[10], src0=v[0:1], src1=v[0:1], src2=v[2:3]),
VOP3SD(VOP3SDOp.V_DIV_SCALE_F64, vdst=v[4:5], sdst=s[10:11], src0=v[0:1], src1=v[0:1], src2=v[2:3]),
]
st = run_program(instructions, n_lanes=1)
result = i642f(st.vgpr[0][4] | (st.vgpr[0][5] << 32))
@@ -1359,43 +1359,6 @@ class TestF64ToI64Conversion(unittest.TestCase):
self.assertEqual(result, 5000000000)
class TestB64VOPLiteral(unittest.TestCase):
"""Tests for B64 VOP operations with literal encoding.
B64 operations (like V_LSHLREV_B64) should zero-extend the literal to 64 bits,
NOT put it in the high 32 bits like F64 operations do.
"""
def test_v_lshlrev_b64_literal_shift_amount(self):
"""V_LSHLREV_B64 with literal shift amount (src0 is 32-bit)."""
# Shift 1 left by 100 (0x64) - uses literal encoding for src0
# Shift amount is 100 & 63 = 36, so 1 << 36 = 0x1000000000
instructions = [
s_mov_b32(s[0], 1),
s_mov_b32(s[1], 0),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_lshlrev_b64(v[2:3], 100, v[0:1]), # 100 > 64, uses literal encoding
]
st = run_program(instructions, n_lanes=1)
# lo = 0x00000000, hi = 0x00000010 = 1 << (36-32)
self.assertEqual(st.vgpr[0][2], 0x00000000)
self.assertEqual(st.vgpr[0][3], 0x00000010)
def test_v_lshlrev_b64_literal_value(self):
"""V_LSHLREV_B64 with literal as the 64-bit value being shifted (src1).
B64 literals are zero-extended (not shifted to high bits like F64).
0xDEADBEEF << 4 = 0xDEADBEEF0 = lo=0xEADBEEF0, hi=0x0000000D
"""
instructions = [
v_lshlrev_b64(v[0:1], 4, 0xDEADBEEF), # shift literal left by 4
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xEADBEEF0) # lo
self.assertEqual(st.vgpr[0][1], 0x0000000D) # hi
class TestWMMAMore(unittest.TestCase):
"""More WMMA tests."""
@@ -2848,112 +2811,5 @@ class TestMin3Max3Unsigned(unittest.TestCase):
self.assertEqual(st.vgpr[0][1] & 0xFFFF, 0)
class TestVOP3Clamp(unittest.TestCase):
"""Tests for VOP3 clamp modifier (clmp=1).
The clamp modifier restricts float outputs to [0.0, 1.0] range.
This is used by operations like clip(0, 1) which AMD LLVM compiles to
v_max_f32_e64 with clmp=1.
Regression test for: clip(0, 1) bug where emulator ignored clmp field.
"""
def test_v_max_f32_e64_clamp_positive(self):
"""V_MAX_F32_E64 with clamp: value > 1.0 should be clamped to 1.0."""
instructions = [
v_mov_b32_e32(v[0], 2.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 1.0, places=5)
def test_v_max_f32_e64_clamp_negative(self):
"""V_MAX_F32_E64 with clamp: value < 0.0 should be clamped to 0.0."""
instructions = [
v_mov_b32_e32(v[0], -1.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.0, places=5)
def test_v_max_f32_e64_clamp_in_range(self):
"""V_MAX_F32_E64 with clamp: value in [0,1] should pass through."""
instructions = [
v_mov_b32_e32(v[0], 0.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.5, places=5)
def test_v_max_f32_e64_no_clamp(self):
"""V_MAX_F32_E64 without clamp: value > 1.0 should pass through."""
instructions = [
v_mov_b32_e32(v[0], 2.5),
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=0),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.5, places=5)
def test_v_min_f32_e64_clamp_negative(self):
"""V_MIN_F32_E64 with clamp: value < 0.0 should be clamped to 0.0."""
instructions = [
v_mov_b32_e32(v[0], -2.0),
VOP3(VOP3Op.V_MIN_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.0, places=5)
def test_v_add_f32_e64_clamp(self):
"""V_ADD_F32_E64 with clamp: 0.7 + 0.8 = 1.5 -> 1.0."""
instructions = [
v_mov_b32_e32(v[0], 0.7),
v_mov_b32_e32(v[1], 0.8),
VOP3(VOP3Op.V_ADD_F32_E64, vdst=v[2], src0=v[0], src1=v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 1.0, places=5)
def test_v_mul_f32_e64_clamp_underflow(self):
"""V_MUL_F32_E64 with clamp: 0.5 * -2.0 = -1.0 -> 0.0."""
instructions = [
v_mov_b32_e32(v[0], 0.5),
v_mov_b32_e32(v[1], -2.0),
VOP3(VOP3Op.V_MUL_F32_E64, vdst=v[2], src0=v[0], src1=v[1], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 0.0, places=5)
def test_v_fma_f32_clamp(self):
"""V_FMA_F32 with clamp: 2*2+1 = 5 -> 1.0."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 2.0),
v_mov_b32_e32(v[2], 1.0),
VOP3(VOP3Op.V_FMA_F32, vdst=v[3], src0=v[0], src1=v[1], src2=v[2], clmp=1),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][3]), 1.0, places=5)
def test_v_max_f32_e64_clamp_multilane(self):
"""V_MAX_F32_E64 with clamp: test multiple lanes with different values."""
# lane 0: -0.5 -> 0.0
# lane 1: 0.5 -> 0.5
# lane 2: 1.5 -> 1.0
# lane 3: 2.5 -> 1.0
instructions = [
# Setup different values per lane using lane_id
s_mov_b32(s[0], f2i(0.5)),
v_cvt_f32_i32_e32(v[0], v[255]), # Convert lane_id to float
v_mov_b32_e32(v[2], s[0]), # v2 = 0.5
v_sub_f32_e32(v[0], v[0], v[2]), # Subtract 0.5: lane0=-0.5, lane1=0.5, lane2=1.5, lane3=2.5
VOP3(VOP3Op.V_MAX_F32_E64, vdst=v[1], src0=v[0], src1=v[0], clmp=1),
]
st = run_program(instructions, n_lanes=4)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 0.0, places=5, msg="lane 0: -0.5 should clamp to 0.0")
self.assertAlmostEqual(i2f(st.vgpr[1][1]), 0.5, places=5, msg="lane 1: 0.5 should pass through")
self.assertAlmostEqual(i2f(st.vgpr[2][1]), 1.0, places=5, msg="lane 2: 1.5 should clamp to 1.0")
self.assertAlmostEqual(i2f(st.vgpr[3][1]), 1.0, places=5, msg="lane 3: 2.5 should clamp to 1.0")
if __name__ == '__main__':
unittest.main()
+1 -161
View File
@@ -404,99 +404,8 @@ class TestVOP3P(unittest.TestCase):
self.assertAlmostEqual(hi, 0.0, places=1)
class TestWMMAF16(unittest.TestCase):
"""Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
Note: RDNA3 WMMA F16 uses 8 VGPRs for accumulator/output (same as F32 variant),
but values are packed as f16. This differs from RDNA4 which uses 4 VGPRs.
"""
def test_v_wmma_f16_16x16x16_f16_all_ones(self):
"""V_WMMA_F16_16X16X16_F16 with all ones produces 16.0 in f16."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[16:23] (8 regs)
for i in range(16, 24):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[24:31] (8 regs)
for i in range(24, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[0:7] to zero (8 regs for RDNA3 WMMA F16)
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA F16: D = A @ B + C
instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 in f16, stored in lo 16 bits of each VGPR (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = _f16(result & 0xffff)
self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
def test_v_wmma_f16_16x16x16_f16_with_accumulator(self):
"""V_WMMA_F16_16X16X16_F16 with non-zero accumulator."""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
instructions.append(s_mov_b32(s[1], 0x4500)) # f16 5.0 in lo bits only
# Initialize A matrix in v[16:23] (8 regs)
for i in range(16, 24):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[24:31] (8 regs)
for i in range(24, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[0:7] to 5.0 in lo bits (8 regs for RDNA3 WMMA F16)
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
# WMMA F16: D = A @ B + C
instructions.append(v_wmma_f16_16x16x16_f16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 + 5.0 = 21.0 in f16, stored in lo 16 bits (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = _f16(result & 0xffff)
self.assertAlmostEqual(lo, 21.0, places=0, msg=f"v[{reg}] lane {lane}: expected 21.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
def test_v_wmma_f16_16x16x16_f16_high_registers(self):
"""V_WMMA_F16_16X16X16_F16 with high register indices.
Regression test: WMMA was using static register indices instead of dynamic.
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
"""
from extra.assembly.amd.test.hw.helpers import _f16
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[64:71] (8 regs)
for i in range(64, 72):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[80:87] (8 regs)
for i in range(80, 88):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[96:103] to zero (8 regs for RDNA3 WMMA F16)
for i in range(96, 104):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA F16: D = A @ B + C, result in v[96:103]
instructions.append(v_wmma_f16_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
# Copy results to v[0:7] for checking
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
st = run_program(instructions, n_lanes=32)
# Result should be 16.0 in f16, stored in lo 16 bits (hi bits are 0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
lo = _f16(result & 0xffff)
self.assertAlmostEqual(lo, 16.0, places=1, msg=f"v[{reg}] lane {lane}: expected 16.0, got {lo}")
self.assertEqual(result >> 16, 0, msg=f"v[{reg}] lane {lane}: hi bits should be 0")
class TestWMMA(unittest.TestCase):
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions with F32 output."""
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions."""
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
@@ -531,75 +440,6 @@ class TestWMMA(unittest.TestCase):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_f16_high_registers(self):
"""V_WMMA_F32_16X16X16_F16 with high register indices.
Regression test: WMMA was using static register indices instead of dynamic,
causing incorrect results when registers weren't at the default positions.
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
"""
instructions = []
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
# Initialize A matrix in v[64:71]
for i in range(64, 72):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize B matrix in v[80:87]
for i in range(80, 88):
instructions.append(v_mov_b32_e32(v[i], s[0]))
# Initialize C (accumulator) in v[96:103] to zero
for i in range(96, 104):
instructions.append(v_mov_b32_e32(v[i], 0))
# WMMA: D = A @ B + C, result in v[96:103]
instructions.append(v_wmma_f32_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
# Copy results to v[0:7] for checking
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
class TestWMMABF16(unittest.TestCase):
"""Tests for WMMA BF16 instructions."""
def test_v_wmma_f32_16x16x16_bf16_all_ones(self):
"""V_WMMA_F32_16X16X16_BF16 with all ones produces 16.0."""
instructions = []
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
instructions.append(s_mov_b32(s[0], 0x3f803f80))
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], 0))
instructions.append(v_wmma_f32_16x16x16_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
expected = f2i(16.0)
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
def test_v_wmma_f32_16x16x16_bf16_with_accumulator(self):
"""V_WMMA_F32_16X16X16_BF16 with non-zero accumulator."""
instructions = []
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
instructions.append(s_mov_b32(s[0], 0x3f803f80))
instructions.append(s_mov_b32(s[1], f2i(5.0)))
for i in range(16, 32):
instructions.append(v_mov_b32_e32(v[i], s[0]))
for i in range(8):
instructions.append(v_mov_b32_e32(v[i], s[1]))
instructions.append(v_wmma_f32_16x16x16_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
st = run_program(instructions, n_lanes=32)
expected = f2i(21.0) # 16 + 5
for lane in range(32):
for reg in range(8):
result = st.vgpr[lane][reg]
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
class TestSpecialOps(unittest.TestCase):
"""Tests for special operations (SAD, PERM, DOT2)."""
-105
View File
@@ -731,111 +731,6 @@ class TestVCCBehavior(unittest.TestCase):
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
class TestCmpNge(unittest.TestCase):
"""Tests for V_CMP_NGE (not-greater-or-equal) with NaN semantics.
NGE = !(a >= b). With NaN inputs:
- If either input is NaN, a >= b is false, so !(false) = true
- This differs from a < b which returns false for NaN inputs
"""
def test_v_cmp_nge_f32_normal_values(self):
"""v_cmp_nge_f32: basic comparison with normal floats."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(2.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 2.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= 2.0) should be true")
def test_v_cmp_nge_f32_equal_values(self):
"""v_cmp_nge_f32: equal values should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(1.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_greater_value(self):
"""v_cmp_nge_f32: greater value should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(2.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(2.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_neg_inf(self):
"""v_cmp_nge_f32: -inf compared to normal value."""
neg_inf = 0xff800000 # -inf
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(-inf >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(-inf >= 1.0) should be true")
def test_v_cmp_nge_f32_clears_inactive_vcc_bits(self):
"""v_cmp_nge_f32 with partial EXEC clears inactive VCC bits (hardware behavior)."""
neg_inf = 0xff800000 # -inf
instructions = [
# Set VCC to all 1s first
s_mov_b32(VCC_LO, 0xFFFFFFFF),
# Set EXEC to only lane 0
s_mov_b32(EXEC_LO, 0x00000001),
# v0 = 1.0 for lane 0
v_mov_b32_e32(v[0], f2i(1.0)),
# Compare: !(-inf >= 1.0) = true for lane 0
v_cmp_nge_f32_e32(neg_inf, v[0]),
]
st = run_program(instructions, n_lanes=16)
# Hardware clears inactive lane bits, only active lane results remain
# Lane 0 result = 1 (true), lanes 1-15 = 0 (cleared)
self.assertEqual(st.vcc, 0x00000001, "VCC should only have active lane results")
def test_v_cmp_nge_f32_nan_src0(self):
"""v_cmp_nge_f32: NaN in src0 should return true (NaN >= x is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= 1.0) should be true")
def test_v_cmp_nge_f32_nan_src1(self):
"""v_cmp_nge_f32: NaN in src1 should return true (x >= NaN is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= NaN) should be true")
def test_v_cmp_nge_f32_both_nan(self):
"""v_cmp_nge_f32: both NaN should return true."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= NaN) should be true")
class TestCmpxPartialWavefront(unittest.TestCase):
"""Tests for V_CMPX with partial wavefronts (fewer than 32 active lanes).
+29 -135
View File
@@ -1,14 +1,18 @@
# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
import unittest, ctypes
import unittest, ctypes, os
from dataclasses import dataclass
from pathlib import Path
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
from extra.assembly.amd import decode_inst
# Set environment before any tinygrad imports to use MOCKGPU
# This allows generating AMD GPU kernels without requiring real hardware
os.environ["AMD"] = "1"
os.environ["MOCKGPU"] = "1"
os.environ["PYTHON_REMU"] = "1"
from extra.assembly.amd.emu import WaveState, decode_program, WAVE_SIZE, set_valid_mem_ranges, LDSMem
from extra.assembly.amd.test.helpers import KernelInfo
from extra.assembly.amd.test.bench_emu import REMU_PATH
def set_valid_mem_ranges(ranges): pass # emu2 doesn't need this
def _is_f32_nan(bits: int) -> bool:
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
@@ -87,61 +91,33 @@ class PythonEmulator:
def __init__(self):
self.state: WaveState | None = None
self.program: dict | None = None
self.vmem_buf = None
self.lds_buf = None
self.kernel_buf = None # Keep kernel bytes alive
self.lib_addr = 0 # Base address of kernel code
def create(self, kernel: bytes, n_lanes: int):
import ctypes
from tinygrad.device import Buffer, BufferSpec
from tinygrad.dtype import dtypes
# Store kernel in a ctypes buffer so generic instructions can read from vmem at actual PC address
self.kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
self.lib_addr = ctypes.addressof(self.kernel_buf)
# Remap program dict to use actual addresses (like run_asm does)
program_raw = decode_program(kernel)
self.program = {self.lib_addr + offset: val for offset, val in program_raw.items()}
self.state = WaveState(n_lanes)
self.state.pc = self.lib_addr # Set PC to code base address
self.vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
self.lds_buf = Buffer('CPU', 65536 // 4, dtypes.uint32).ensure_allocated()
self.program = decode_program(kernel)
self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
self.state.exec_mask = (1 << n_lanes) - 1
def step(self) -> int:
import ctypes
assert self.program is not None and self.state is not None
pc = self.state.pc
if pc == 0xFFFFFFFFFFFFFFFF or pc not in self.program: return -1
name, fxn, globals_list, _runner = self.program[pc]
if fxn is None: return 1 # unsupported instruction
buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,
2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}
# Direct ctypes call - bypasses HCQ overhead
fxn(*[ctypes.c_uint64(buf_addrs[g]) for g in globals_list], ctypes.c_int32(0))
return -1 if self.state.pc == 0xFFFFFFFFFFFFFFFF else 0
return self.program[self.state.pc]._dispatch(self.state, self.program[self.state.pc])
def set_sgpr(self, idx: int, val: int):
assert self.state is not None
self.state._write_sgpr(idx, val)
self.state.sgpr[idx] = val & 0xffffffff
def set_vgpr(self, lane: int, idx: int, val: int):
assert self.state is not None
self.state._write_vgpr(idx, lane, val)
self.state.vgpr[lane][idx] = val & 0xffffffff
def get_snapshot(self) -> StateSnapshot:
assert self.state is not None
sgpr = [self.state._read_sgpr(i) for i in range(128)]
vgpr = [[self.state._read_vgpr(reg, lane) for reg in range(256)] for lane in range(WAVE_SIZE)]
# Convert actual PC address to word offset for comparison with Rust emulator
pc_offset = (self.state.pc - self.lib_addr) // 4 if self.state.pc != 0xFFFFFFFFFFFFFFFF else 0xFFFFFFFFFFFFFFFF
return StateSnapshot(pc=pc_offset, scc=self.state._read_sgpr(SCC.offset), vcc=sgpr[VCC_LO.offset],
exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
local_size: tuple[int, int, int], program, max_steps: int, debug: bool, trace_len: int,
kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, int]:
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
max_workgroups: int = 8) -> tuple[bool, str, int]:
"""Run a single kernel through both emulators. Returns (success, message, total_steps)."""
gx, gy, gz = global_size
lx, ly, lz = local_size
total_steps = 0
wg_count = 0
@@ -164,52 +140,28 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
emu.set_sgpr(13, gidx)
emu.set_sgpr(14, gidy)
emu.set_sgpr(15, gidz)
# Initialize v[0] with packed workitem IDs for each lane
for lane in range(n_lanes):
tid = lane
z, y, x = tid // (lx * ly), (tid // lx) % ly, tid % lx
emu.set_vgpr(lane, 0, (z << 20) | (y << 10) | x)
step = 0
trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
prev_sync_after = False # Track if previous instruction had known Rust bugs
try:
while step < max_steps:
rust_before = rust.get_snapshot()
python_before = python.get_snapshot()
inst_info = python.program.get(python.lib_addr + python_before.pc * 4) # Convert word offset to actual address
inst_hex_name = inst_info[0] if inst_info else f"unknown at PC={python_before.pc}"
# Decode the instruction to get mnemonic for sync_after checks
try:
# Format is mnemonic_hexbytes, e.g. v_exp_f32_e32_014b027e -> hex is 014b027e
parts = inst_hex_name.rsplit('_', 1)
inst_bytes_hex = parts[1] if len(parts) == 2 else ""
inst_bytes = bytes.fromhex(inst_bytes_hex) if inst_bytes_hex else b''
decoded = decode_inst(inst_bytes) if inst_bytes else None
inst_mnemonic = repr(decoded).split('(')[0] if decoded else ""
except:
inst_mnemonic = ""
# For generic instructions, use function name for sync_after check
if not inst_mnemonic: inst_mnemonic = inst_hex_name
inst_str = inst_hex_name
inst = program.get(python_before.pc)
inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
trace.append((step, python_before.pc, inst_str, rust_before, python_before))
if len(trace) > trace_len: trace.pop(0)
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
# Instructions with known Rust emulator bugs or precision differences - sync Python to Rust after execution
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
# v_div_scale/v_div_fixup: Rust has different VCC handling
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
# s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
# v_exp_f32/v_log_f32/v_ldexp_f32: precision differences in transcendental functions
# s_delay_alu: Rust handles differently
# v_add_co_ci_u32/v_sub_co_ci_u32/v_subrev_co_ci_u32: Rust preserves inactive VCC bits, but hardware clears all bits
sync_after = any(x in inst_mnemonic.lower() for x in ('v_div_scale', 'v_div_fixup', 'v_cvt_f16_f32', 's_add_i32', 's_sub_i32',
'v_exp_f32', 'v_log_f32', 'v_ldexp_f32', 's_delay_alu',
'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'))
# Skip comparison if previous instruction had known Rust bugs (states were synced but may still differ slightly)
diffs = rust_before.diff(python_before, n_lanes) if not prev_sync_after else []
sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
diffs = rust_before.diff(python_before, n_lanes)
if diffs:
trace_lines = []
for idx, (s, pc, d, rb, pb) in enumerate(trace):
@@ -248,12 +200,7 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
for lane in range(n_lanes):
for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
assert python.state is not None
# Convert Rust's word-based PC to Python's actual address
python.state.pc = python.lib_addr + rust_after.pc * 4
python.state._write_sgpr(SCC.offset, rust_after.scc)
python.state._write_sgpr(VCC_LO.offset, rust_after.vcc)
python.state._write_sgpr(EXEC_LO.offset, rust_after.exec_mask)
prev_sync_after = sync_after
python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
if rust_result == -1:
total_steps += step + 1
@@ -307,7 +254,7 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
ok, msg, steps = run_single_kernel(
kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
kernel.local_size, program, max_steps, debug, trace_len, ki
program, max_steps, debug, trace_len, ki
)
total_steps += steps
if not ok:
@@ -334,8 +281,7 @@ def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list,
set_valid_mem_ranges(ranges)
program = decode_program(kernel)
# Legacy wrapper assumes local_size = (n_lanes, 1, 1)
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), program, max_steps, debug, trace_len)
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]]:
@@ -378,7 +324,6 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int],
buf_sizes.append(b.nbytes)
kernels.append(KernelInfo(
code=bytes(sec.content),
src=lowered.prg.p.src,
global_size=tuple(lowered.prg.p.global_size),
local_size=tuple(lowered.prg.p.local_size),
buf_idxs=buf_idxs,
@@ -441,7 +386,6 @@ class TestTinygradKernels(unittest.TestCase):
from tinygrad import dtypes
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
def test_exp(self): self._test_kernel(lambda T: T.empty(1024).exp())
def test_cross_entropy(self):
import numpy as np
np.random.seed(0)
@@ -453,55 +397,5 @@ class TestTinygradKernels(unittest.TestCase):
from tinygrad import dtypes
self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
def test_sin_large_f32(self):
"""Test sin with large values that trigger Payne-Hanek range reduction."""
# Values around 859240 trigger the Payne-Hanek algorithm
# This tests the integer multiply-high instructions used in range reduction
self._test_kernel(lambda T: T([859240.0, 1000000.0, 100594688.0]).sin())
def test_clip_zero_one(self):
"""Test clip(0, 1) - regression for binary_crossentropy failure."""
import numpy as np
np.random.seed(0)
x_np = np.random.uniform(-2, 2, (32, 10)).astype(np.float32).tolist()
self._test_kernel(lambda T: T(x_np).clip(0, 1))
def test_mod_int64(self):
"""Test int64 modulo, especially edge cases like 1 % -1."""
from tinygrad import dtypes
self._test_kernel(lambda T: T([1, 10, -10, 7], dtype=dtypes.int64) % T([-1, 3, 3, -3], dtype=dtypes.int64))
def test_expand_flatten_sum(self):
"""Test flatten of expanded tensor followed by sum.
Bug: flatten() of an expanded tensor produces wrong results for certain sizes.
Sizes that are multiples of 32 work (32, 48, 64), but sizes like 33, 49, 50 fail.
This breaks masked_select and nonzero operations.
"""
import numpy as np
np.random.seed(0)
x_np = np.random.uniform(-2, 2, (33,)).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).unsqueeze(-1).expand(33, 3).flatten().sum())
@unittest.skip("slow and broken with AMD_LLVM=1")
def test_nonzero(self):
"""Test nonzero operation - counts and gathers indices of non-zero elements."""
import numpy as np
np.random.seed(42)
x_np = np.random.rand(10, 5, 3).astype(np.float32)
self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).nonzero())
@unittest.skip("Precision differences in v_exp/v_log accumulate across kernels, causing memory divergence")
def test_softmax_argmax_fused(self):
"""Test fused softmax+argmax - tracks exp2 precision issue.
The fused kernel recomputes softmax inline and Python emulator's exp2 polynomial
has up to 1 ULP error vs native exp2f, causing accumulated differences.
"""
import torch
torch.manual_seed(0)
x_np = torch.rand(4, 10).numpy()
self._test_kernel(lambda T: T(x_np.tolist()).softmax(1).argmax())
if __name__ == "__main__":
unittest.main()
+1 -30
View File
@@ -34,44 +34,15 @@ def custom_add_one(A:UOp, arch:str) -> UOp:
sink = UOp.sink(A.base, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)), *assemble_insts(insts, name, arch)))
def custom_add_var(A:UOp, B:UOp, arch:str) -> UOp:
A,B = A.flatten(), B.flatten()
assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
threads = UOp.special(A.size, "lidx0")
var = UOp.variable("var", 0, 10)
insts = [
s_load_b128(s[4:7], s[0:1]),
s_load_b32(s[8], s[0:1], offset=0x10), # all threads load the same variable
s_waitcnt(lgkmcnt=0),
v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset, different per thread
global_load_b32(v[1], v[0], saddr=s[6:7]),
s_waitcnt(vmcnt=0),
v_add_nc_u32_e32(v[1], s[8], v[1]),
global_store_b32(addr=v[0], data=v[1], saddr=s[4:5]),
s_endpgm(),
]
sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(name:=f"custom_add_one_{A.size}"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=(*sink.src, sink)),
*assemble_insts(insts, name, arch, kernarg_size=16)))
class TestCustomKernel(unittest.TestCase):
def test_simple(self):
a = Tensor.full((16, 16), 1.).contiguous().realize()
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].renderer.arch))[0]
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_add_one, arch=Device[Device.DEFAULT].arch))[0]
ei = a.schedule()[-1].lower()
self.assertEqual(ei.prg.estimates.ops, a.numel())
self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
ei.run()
self.assertTrue((a.numpy() == 2.).all())
def test_variable(self):
b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
a = Tensor.zeros_like(b).contiguous().realize()
a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_var, arch=Device[Device.DEFAULT].renderer.arch))[0]
ei = a.schedule()[-1].lower()
for i in range(4):
ei.run({"var":i})
self.assertTrue((a.numpy() == 1+i).all())
if __name__ == "__main__":
unittest.main()
-332
View File
@@ -1,332 +0,0 @@
"""Tests for the pcode parser."""
import unittest
from collections import defaultdict
from tinygrad.helpers import DEBUG
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, Ops
from extra.assembly.amd.emu import parse_pcode
from extra.assembly.amd.pcode import parse_expr
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP1Op, VOP2Op, VOP3Op, SOP1Op, SOP2Op, DSOp
def _srcs():
"""Create minimal source variables for pcode parsing."""
u32 = lambda v=0: UOp.const(dtypes.uint32, v)
return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}
class TestBasicParsing(unittest.TestCase):
"""Test basic pcode parsing for common instruction patterns."""
def test_v_add_f32(self):
"""Test parsing V_ADD_F32 pcode."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], _srcs())
self.assertEqual(len(assigns), 1)
dest, _ = assigns[0]
self.assertTrue(dest.startswith('D0'))
def test_v_lshlrev_b32(self):
"""Test parsing V_LSHLREV_B32 pcode."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_LSHLREV_B32_E32], _srcs())
self.assertEqual(len(assigns), 1)
def test_s_cselect_b32(self):
"""Test parsing S_CSELECT_B32 pcode with ternary."""
_, assigns = parse_pcode(PCODE[SOP2Op.S_CSELECT_B32], _srcs())
self.assertEqual(len(assigns), 1)
def test_v_add_co_ci_u32(self):
"""Test parsing V_ADD_CO_CI_U32 with carry."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_CO_CI_U32_E32], _srcs())
self.assertGreaterEqual(len(assigns), 1)
class TestWithSources(unittest.TestCase):
"""Test pcode parsing with actual source operand values."""
def test_v_add_f32_with_sources(self):
"""Test V_ADD_F32 with actual float constants."""
s0 = UOp.const(dtypes.uint32, 0x3f800000) # 1.0f
s1 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be an ADD operation
self.assertEqual(val.op, Ops.ADD)
def test_v_mul_f32_with_sources(self):
"""Test V_MUL_F32 with actual float constants."""
s0 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f
s1 = UOp.const(dtypes.uint32, 0x40400000) # 3.0f
_, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertEqual(val.op, Ops.MUL)
class TestParseExpr(unittest.TestCase):
"""Test the parse_expr function directly."""
def test_integer_literals(self):
"""Test parsing integer literals."""
self.assertEqual(parse_expr('0', {}).arg, 0)
self.assertEqual(parse_expr('42', {}).arg, 42)
self.assertEqual(parse_expr('42U', {}).arg, 42)
def test_negative_integers(self):
"""Test parsing negative integer literals."""
result = parse_expr('-1', {})
self.assertEqual(result.arg, -1)
self.assertEqual(result.dtype, dtypes.int)
def test_float_literals(self):
"""Test parsing float literals."""
result = parse_expr('1.0F', {})
self.assertEqual(result.arg, 1.0)
self.assertEqual(result.dtype, dtypes.float32)
def test_hex_literals(self):
"""Test parsing hex literals."""
result = parse_expr('0xFF', {})
self.assertEqual(result.arg, 255)
def test_variable_lookup(self):
"""Test variable lookup in parse_expr."""
vars = {'x': UOp.const(dtypes.uint32, 42)}
result = parse_expr('x', vars)
self.assertEqual(result.arg, 42)
def test_binary_ops(self):
"""Test parsing binary operations."""
vars = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
# Addition
result = parse_expr('a + b', vars)
self.assertEqual(result.op, Ops.ADD)
# Subtraction with constant folding
result = parse_expr('10 - 5', {})
self.assertEqual(result.op, Ops.CONST)
self.assertEqual(result.arg, 5)
def test_ternary(self):
"""Test parsing ternary expressions."""
vars = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
result = parse_expr('cond ? a : b', vars)
self.assertEqual(result.op, Ops.WHERE)
class TestForLoopParsing(unittest.TestCase):
"""Test for loop parsing (CLZ/CTZ patterns)."""
def test_clz_pcode_exists(self):
"""Verify CLZ pcode is available."""
pcode = PCODE.get(VOP1Op.V_CLZ_I32_U32_E32)
self.assertIsNotNone(pcode)
self.assertIn('for', pcode.lower())
def test_clz_parsing(self):
"""Test CLZ pcode parsing produces correct structure."""
pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF) # All ones - CLZ should be 0
vars, assigns = parse_pcode(pcode, {'S0': S0})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be a nested WHERE structure
self.assertEqual(val.op, Ops.WHERE)
def test_clz_with_zero(self):
"""Test CLZ with input 0 - should return -1."""
pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
S0 = UOp.const(dtypes.uint32, 0)
vars, assigns = parse_pcode(pcode, {'S0': S0})
# Check that the innermost value (default) is -1 (may be wrapped in CAST)
val = assigns[0][1]
# Traverse to innermost WHERE
while val.op == Ops.WHERE:
val = val.src[2] # false branch
# Unwrap CAST if present
while val.op == Ops.CAST:
val = val.src[0]
self.assertEqual(val.arg, -1)
def test_ctz_parsing(self):
"""Test CTZ pcode parsing."""
pcode = PCODE.get(VOP1Op.V_CTZ_I32_B32_E32)
if pcode is None:
self.skipTest("V_CTZ_I32_B32_E32 pcode not available")
S0 = UOp.const(dtypes.uint32, 1) # LSB set - CTZ should be 0
vars, assigns = parse_pcode(pcode, {'S0': S0})
self.assertEqual(len(assigns), 1)
class TestDSPcodePatterns(unittest.TestCase):
"""Test DS instruction pcode patterns."""
def test_ds_load_b32_pcode(self):
"""Test DS_LOAD_B32 pcode is parseable."""
pcode = PCODE.get(DSOp.DS_LOAD_B32)
self.assertIsNotNone(pcode)
self.assertIn('RETURN_DATA', pcode)
self.assertIn('MEM[', pcode)
def test_ds_store_b32_pcode(self):
"""Test DS_STORE_B32 pcode is parseable."""
pcode = PCODE.get(DSOp.DS_STORE_B32)
self.assertIsNotNone(pcode)
self.assertIn('MEM[', pcode)
self.assertIn('DATA', pcode)
def test_mem_read_parsing(self):
"""Test MEM[addr].type read expression parsing."""
# Create a mock LDS buffer
lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
addr = UOp.const(dtypes.uint32, 0)
vars = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
result = parse_expr('MEM[ADDR + OFFSET].b32', vars)
# Should be an INDEX operation into LDS
self.assertIsNotNone(result)
def test_ds_store_2addr_b32_parsing(self):
"""Test DS_STORE_2ADDR_B32 pcode parsing produces MEM writes."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
self.assertIsNotNone(pcode)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Should have 2 MEM write assignments
self.assertEqual(len(assigns), 2)
for dest, val in assigns:
self.assertTrue(dest.startswith('MEM['))
# val should be (addr, write_val) tuple
self.assertIsInstance(val, tuple)
self.assertEqual(len(val), 2)
def test_ds_load_2addr_b32_parsing(self):
"""Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
self.assertIsNotNone(pcode)
lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'_lds': lds,
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Should have 2 RETURN_DATA assignments
self.assertEqual(len(assigns), 2)
self.assertEqual(assigns[0][0], 'RETURN_DATA[31:0]')
self.assertEqual(assigns[1][0], 'RETURN_DATA[63:32]')
def test_ds_store_address_calculation(self):
"""Test DS_STORE_2ADDR_B32 calculates correct addresses (offset * 4)."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 100),
'OFFSET0': UOp.const(dtypes.uint32, 2),
'OFFSET1': UOp.const(dtypes.uint32, 5),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
addr0, _ = assigns[0][1]
addr1, _ = assigns[1][1]
self.assertEqual(addr0.simplify().arg, 108)
self.assertEqual(addr1.simplify().arg, 120)
def test_ds_store_data_values(self):
"""Test DS_STORE_2ADDR_B32 uses correct data values."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
_, val0 = assigns[0][1]
_, val1 = assigns[1][1]
# DATA[31:0] should preserve the value
self.assertEqual(val0.simplify().arg, 0xAAAAAAAA)
self.assertEqual(val1.simplify().arg, 0xBBBBBBBB)
class TestConditionalParsing(unittest.TestCase):
"""Test conditional (if/elsif/else) pcode parsing."""
def test_ternary_in_assignment(self):
"""Test parsing ternary expression (which becomes WHERE)."""
# S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32
pcode = PCODE[SOP2Op.S_CSELECT_B32]
s0 = UOp.const(dtypes.uint32, 10)
s1 = UOp.const(dtypes.uint32, 20)
scc = UOp.const(dtypes.uint32, 1)
vars, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be a WHERE (ternary becomes WHERE)
self.assertEqual(val.op, Ops.WHERE)
class TestAllPcode(unittest.TestCase):
"""Test that all pcode from all architectures can be parsed."""
def _make_srcs(self):
"""Create dummy source variables for pcode parsing."""
u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
lds = UOp(Ops.DEFINE_GLOBAL, dtypes.uint32.ptr(16384), arg=3)
return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
'SIMM16': u32(), 'SIMM32': u32(), 'OFFSET': u32(), 'OFFSET0': u32(), 'OFFSET1': u32(), 'offset1': u32(),
'ADDR': u32(), 'ADDR_BASE': u32(), 'TADDR': u32(), 'DATA': u32(), 'DATA0': u32(), 'DATA1': u32(), 'DATA2': u32(),
'VDATA': u32(), 'VDATA0': u32(), 'VDATA1': u32(), 'VDATA2': u32(), 'VDATA3': u32(),
'OPSEL': u32(), 'OPSEL_HI': u32(), 'NEG': u32(), 'NEG_HI': u32(), 'CLAMP': u32(),
'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'WAVE_STATUS': u32(),
'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
'_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}
def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float):
"""Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs."""
srcs = self._make_srcs()
passed, skipped, errors = 0, 0, defaultdict(list)
for op, pcode in pcode_dict.items():
try:
parse_pcode(pcode, srcs)
passed += 1
except RuntimeError as e: skipped += 1; errors[str(e)].append(op.name)
except Exception as e: self.fail(f"[{arch}] {op.name}: {e}\nPcode: {pcode[:200]}")
total = len(pcode_dict)
pct = 100 * passed / total
print(f"{arch}: {passed}/{total} ({pct:.1f}%) parsed, {skipped} skipped")
if DEBUG >= 2:
for err, ops in sorted(errors.items(), key=lambda x: -len(x[1])):
print(f" {err}: {', '.join(ops[:5])}{'...' if len(ops) > 5 else ''} ({len(ops)})")
self.assertGreaterEqual(pct, min_pct, f"[{arch}] {pct:.1f}% < {min_pct}% threshold")
def test_parse_all_cdna_pcode(self):
from extra.assembly.amd.autogen.cdna.str_pcode import PCODE as CDNA_PCODE
self._parse_all_pcode(CDNA_PCODE, "CDNA", min_pct=60)
def test_parse_all_rdna3_pcode(self):
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE as RDNA3_PCODE
self._parse_all_pcode(RDNA3_PCODE, "RDNA3", min_pct=90)
def test_parse_all_rdna4_pcode(self):
from extra.assembly.amd.autogen.rdna4.str_pcode import PCODE as RDNA4_PCODE
self._parse_all_pcode(RDNA4_PCODE, "RDNA4", min_pct=65)
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -7,7 +7,7 @@ import unittest
from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import VCC_HI, EXEC_LO, NULL
OFF = NULL # OFF is alias for NULL
from extra.assembly.amd import detect_format
from extra.assembly.amd.decode import detect_format
class TestDS(unittest.TestCase):
+2 -17
View File
@@ -6,21 +6,17 @@ from extra.assembly.amd.autogen.rdna3.ins import *
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd.test.test_roundtrip import compile_asm
class IntegrationTestBase(unittest.TestCase):
class TestIntegration(unittest.TestCase):
inst: Inst
arch: str
def tearDown(self):
if not hasattr(self, 'inst'): return
b = self.inst.to_bytes()
st = self.inst.disasm()
# Test that the instruction can be compiled by LLVM and produces the same bytes
desc = f"{st:25s} {self.inst} {b!r}"
self.assertEqual(b, compile_asm(st, arch=self.arch), desc)
self.assertEqual(b, compile_asm(st), desc)
print(desc)
class TestIntegration(IntegrationTestBase):
arch: str = "rdna3"
def test_wmma(self):
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[184:191], v[136:143], v[0:7])
@@ -128,17 +124,6 @@ class TestIntegration(IntegrationTestBase):
int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
self.assertEqual(self.inst, int_inst)
class TestIntegrationCDNA(IntegrationTestBase):
arch = "cdna"
def test_mfma(self):
from extra.assembly.amd.autogen.cdna.ins import v_mfma_f32_16x16x16_f16
self.inst = v_mfma_f32_16x16x16_f16(v[0:3], v[0:1], v[0:1], 0)
def test_mfma_fp8(self):
from extra.assembly.amd.autogen.cdna.ins import v_mfma_f32_16x16x128_f8f6f4
self.inst = v_mfma_f32_16x16x128_f8f6f4(v[0:3], v[0:5], v[0:5], 1, cbsz=2, blgp=2)
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
+1 -24
View File
@@ -11,7 +11,7 @@ Only compute-relevant instruction formats are tested. Graphics-only formats not
import unittest, re, subprocess, functools
from tinygrad.helpers import fetch
from extra.assembly.amd.disasm import disasm
from extra.assembly.amd import decode_inst, detect_format
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_target, get_mattr
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
@@ -125,26 +125,6 @@ def _make_test(f: str, arch: str, test_type: str):
except ValueError: skipped += 1 # skip invalid opcodes not in enum
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "repr":
# Test that eval(repr(inst)) reproduces the instruction
if arch == "rdna3": import extra.assembly.amd.autogen.rdna3.ins as ins
elif arch == "rdna4": import extra.assembly.amd.autogen.rdna4.ins as ins
elif arch == "cdna": import extra.assembly.amd.autogen.cdna.ins as ins
ns = {k: getattr(ins, k) for k in dir(ins) if not k.startswith('_')}
passed, skipped = 0, 0
for _, data in tests:
try:
decoded = detect_format(data, arch).from_bytes(data)
if decoded.to_bytes()[:len(data)] != data: skipped += 1; continue # skip if binary roundtrip fails
r = repr(decoded)
try:
decoded2 = eval(r, ns) # noqa: S307
if decoded == decoded2: passed += 1
else: skipped += 1
except Exception: skipped += 1
except ValueError: skipped += 1
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "disasm":
to_test = []
for _, data in tests:
@@ -169,15 +149,12 @@ class TestLLVM(unittest.TestCase): pass
for f in RDNA_FILES:
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
setattr(TestLLVM, f"test_rdna3_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "repr"))
for f in CDNA_FILES:
setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
setattr(TestLLVM, f"test_cdna_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "repr"))
for f in RDNA4_FILES:
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
setattr(TestLLVM, f"test_rdna4_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "repr"))
if __name__ == "__main__":
unittest.main()
+403
View File
@@ -0,0 +1,403 @@
#!/usr/bin/env python3
"""Tests for the RDNA3 pseudocode DSL."""
import unittest
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
from extra.assembly.amd.test.helpers import ExecContext
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
# Compile pseudocode functions on demand for regression tests
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', PCODE[VOP3SDOp.V_DIV_SCALE_F32])
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', PCODE[VOPCOp.V_CMP_CLASS_F32_E32])
class TestReg(unittest.TestCase):
def test_u32_read(self):
r = Reg(0xDEADBEEF)
self.assertEqual(int(r.u32), 0xDEADBEEF)
def test_u32_write(self):
r = Reg(0)
r.u32 = 0x12345678
self.assertEqual(r._val, 0x12345678)
def test_f32_read(self):
r = Reg(0x40400000) # 3.0f
self.assertAlmostEqual(float(r.f32), 3.0)
def test_f32_write(self):
r = Reg(0)
r.f32 = 3.0
self.assertEqual(r._val, 0x40400000)
def test_i32_signed(self):
r = Reg(0xFFFFFFFF) # -1 as signed
self.assertEqual(int(r.i32), -1)
def test_u64(self):
r = Reg(0xDEADBEEFCAFEBABE)
self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
def test_f64(self):
r = Reg(0x4008000000000000) # 3.0 as f64
self.assertAlmostEqual(float(r.f64), 3.0)
class TestTypedView(unittest.TestCase):
def test_bit_slice(self):
r = Reg(0xDEADBEEF)
# Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
self.assertEqual(r.u32[7:0].u32, 0xEF)
self.assertEqual(r.u32[15:8].u32, 0xBE)
self.assertEqual(r.u32[23:16].u32, 0xAD)
self.assertEqual(r.u32[31:24].u32, 0xDE)
# Also works with int() for arithmetic
self.assertEqual(int(r.u32[7:0]), 0xEF)
def test_single_bit_read(self):
r = Reg(0b11010101)
self.assertEqual(r.u32[0], 1)
self.assertEqual(r.u32[1], 0)
self.assertEqual(r.u32[2], 1)
self.assertEqual(r.u32[3], 0)
def test_single_bit_write(self):
r = Reg(0)
r.u32[5] = 1
r.u32[3] = 1
self.assertEqual(r._val, 0b00101000)
def test_nested_bit_access(self):
# S0.u32[S1.u32[4:0]] - access bit at position from another register
s0 = Reg(0b11010101)
s1 = Reg(3)
bit_pos = s1.u32[4:0] # TypedView, int value = 3
bit_val = s0.u32[int(bit_pos)] # bit 3 of s0 = 0
self.assertEqual(int(bit_pos), 3)
self.assertEqual(bit_val, 0)
def test_arithmetic(self):
r1 = Reg(0x40400000) # 3.0f
r2 = Reg(0x40800000) # 4.0f
result = r1.f32 + r2.f32
self.assertAlmostEqual(result, 7.0)
def test_comparison(self):
r1 = Reg(5)
r2 = Reg(3)
self.assertTrue(r1.u32 > r2.u32)
self.assertFalse(r1.u32 < r2.u32)
self.assertTrue(r1.u32 != r2.u32)
class TestTypedView(unittest.TestCase):
def test_slice_read(self):
r = Reg(0x56781234)
self.assertEqual(r[15:0].u16, 0x1234)
self.assertEqual(r[31:16].u16, 0x5678)
def test_slice_write(self):
r = Reg(0)
r[15:0].u16 = 0x1234
r[31:16].u16 = 0x5678
self.assertEqual(r._val, 0x56781234)
def test_slice_f16(self):
r = Reg(0)
r[15:0].f16 = 3.0
self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
class TestCompiler(unittest.TestCase):
def test_ternary(self):
result = _expr("a > b ? 1 : 0")
self.assertIn("if", result)
self.assertIn("else", result)
def test_type_prefix_strip(self):
self.assertEqual(_expr("1'0U"), "0")
self.assertEqual(_expr("32'1"), "1")
self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
def test_suffix_strip(self):
self.assertEqual(_expr("0ULL"), "0")
self.assertEqual(_expr("1LL"), "1")
self.assertEqual(_expr("5U"), "5")
self.assertEqual(_expr("3.14F"), "3.14")
def test_boolean_ops(self):
self.assertIn("and", _expr("a && b"))
self.assertIn("or", _expr("a || b"))
self.assertIn("!=", _expr("a <> b"))
def test_pack16(self):
result = _expr("{ a, b }")
self.assertIn("_pack", result)
def test_type_cast_strip(self):
self.assertEqual(_expr("64'U(x)"), "(x)")
self.assertEqual(_expr("32'I(y)"), "(y)")
class TestExecContext(unittest.TestCase):
def test_float_add(self):
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
def test_float_mul(self):
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
ctx.run("D0.f32 = S0.f32 * S1.f32")
self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
def test_scc_comparison(self):
ctx = ExecContext(s0=42, s1=42)
ctx.run("SCC = S0.u32 == S1.u32")
self.assertEqual(ctx.SCC._val, 1)
def test_scc_comparison_false(self):
ctx = ExecContext(s0=42, s1=43)
ctx.run("SCC = S0.u32 == S1.u32")
self.assertEqual(ctx.SCC._val, 0)
def test_ternary(self):
code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
ctx = ExecContext(s0=5, s1=3)
ctx.run(code)
self.assertEqual(ctx.D0._val, 1)
def test_pack(self):
code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
ctx = ExecContext(s0=0x1234, s1=0x5678)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0x56781234)
def test_tmp_with_typed_access(self):
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
self.assertEqual(ctx.D0._val, 300)
def test_s_add_u32_pattern(self):
# Real pseudocode pattern from S_ADD_U32
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
# Test overflow case
ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
ctx.run(code)
self.assertEqual(ctx.D0._val, 0) # Wraps to 0
self.assertEqual(ctx.SCC._val, 1) # Carry set
def test_s_add_u32_no_overflow(self):
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
D0.u32 = tmp.u32""")
ctx = ExecContext(s0=100, s1=200)
ctx.run(code)
self.assertEqual(ctx.D0._val, 300)
self.assertEqual(ctx.SCC._val, 0) # No carry
def test_vcc_lane_read(self):
ctx = ExecContext(vcc=0b1010, lane=1)
# Lane 1 is set
self.assertEqual(ctx.VCC.u64[1], 1)
self.assertEqual(ctx.VCC.u64[2], 0)
def test_vcc_lane_write(self):
ctx = ExecContext(vcc=0, lane=0)
ctx.VCC.u64[3] = 1
ctx.VCC.u64[1] = 1
self.assertEqual(ctx.VCC._val, 0b1010)
def test_for_loop(self):
# CTZ pattern - find first set bit
code = _compile_pseudocode("""tmp = -1
for i in 0 : 31 do
if S0.u32[i] == 1 then
tmp = i
endif
endfor
D0.i32 = tmp""")
ctx = ExecContext(s0=0b1000) # Bit 3 is set
ctx.run(code)
self.assertEqual(ctx.D0._val & MASK32, 3)
def test_result_dict(self):
ctx = ExecContext(s0=5, s1=3)
ctx.D0.u32 = 42
ctx.SCC._val = 1
result = ctx.result()
self.assertEqual(result['d0'], 42)
self.assertEqual(result['scc'], 1)
class TestPseudocodeRegressions(unittest.TestCase):
"""Regression tests for pseudocode instruction emulation bugs."""
def test_v_div_scale_f32_vcc_always_returned(self):
"""V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
This caused division to produce wrong results for multiple lanes."""
# Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
s0 = 0x3f800000 # 1.0
s1 = 0x40400000 # 3.0
s2 = 0x3f800000 # 1.0 (numerator)
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
# Must always have VCC in result
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
def test_v_cmp_class_f32_detects_quiet_nan(self):
"""V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
quiet_nan = 0x7fc00000 # quiet NaN: exponent=255, bit22=1
signal_nan = 0x7f800001 # signaling NaN: exponent=255, bit22=0
# Test quiet NaN detection (bit 1 in mask)
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
# Test signaling NaN detection (bit 0 in mask)
s1_signal = 0b0000000001 # bit 0 = signaling NaN
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
# Test that quiet NaN doesn't match signaling NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
# Test that signaling NaN doesn't match quiet NaN mask
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
def testisNAN_with_typed_view(self):
"""isNAN must work with TypedView objects, not just Python floats.
Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
nan_reg = Reg(0x7fc00000) # quiet NaN
normal_reg = Reg(0x3f800000) # 1.0
inf_reg = Reg(0x7f800000) # +inf
self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
class TestBF16(unittest.TestCase):
"""Tests for BF16 (bfloat16) support."""
def test_bf16_conversion(self):
"""Test bf16 <-> f32 conversion."""
# bf16 is just the top 16 bits of f32
# 1.0f = 0x3f800000, bf16 = 0x3f80
self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
self.assertEqual(_ibf16(1.0), 0x3f80)
# 2.0f = 0x40000000, bf16 = 0x4000
self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
self.assertEqual(_ibf16(2.0), 0x4000)
# -1.0f = 0xbf800000, bf16 = 0xbf80
self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
self.assertEqual(_ibf16(-1.0), 0xbf80)
def test_bf16_special_values(self):
"""Test bf16 special values (inf, nan)."""
import math
# +inf: f32 = 0x7f800000, bf16 = 0x7f80
self.assertTrue(math.isinf(_bf16(0x7f80)))
self.assertEqual(_ibf16(float('inf')), 0x7f80)
# -inf: f32 = 0xff800000, bf16 = 0xff80
self.assertTrue(math.isinf(_bf16(0xff80)))
self.assertEqual(_ibf16(float('-inf')), 0xff80)
# NaN: quiet NaN bf16 = 0x7fc0
self.assertTrue(math.isnan(_bf16(0x7fc0)))
self.assertEqual(_ibf16(float('nan')), 0x7fc0)
def test_bf16_register_property(self):
"""Test Reg.bf16 property."""
r = Reg(0)
r.bf16 = 3.0 # 3.0f = 0x40400000, bf16 = 0x4040
self.assertEqual(r._val & 0xffff, 0x4040)
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
def test_bf16_slice_property(self):
"""Test TypedView.bf16 property."""
r = Reg(0x40404040) # Two bf16 3.0 values
self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
class TestBytePermute(unittest.TestCase):
"""Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
def test_byte_select_0_to_7(self):
"""Test selecting bytes 0-7 from 64-bit data."""
# data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
# Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
data = 0x0706050403020100
for i in range(8):
self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
def test_sign_extend_bytes(self):
"""Test sign extension selectors 8-11."""
# sel 8: sign of byte 1 (bits 15:8)
# sel 9: sign of byte 3 (bits 31:24)
# sel 10: sign of byte 5 (bits 47:40)
# sel 11: sign of byte 7 (bits 63:56)
data = 0x8000800080008000 # All relevant bytes have sign bit set
self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
data = 0x7f007f007f007f00 # No sign bits set
self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
def test_constant_zero(self):
"""Test selector 12 returns 0x00."""
self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
def test_constant_ff(self):
"""Test selectors >= 13 return 0xFF."""
for sel in [13, 14, 15, 255]:
self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
class TestSADHelpers(unittest.TestCase):
"""Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
def test_v_sad_u8_basic(self):
"""Test v_sad_u8 with simple values."""
# s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
result = v_sad_u8(0x04030201, 0x04030201, 0)
self.assertEqual(result, 0)
# s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
result = v_sad_u8(0x05040302, 0x04030201, 0)
self.assertEqual(result, 4)
def test_v_sad_u8_with_accumulator(self):
"""Test v_sad_u8 with non-zero accumulator."""
# s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
result = v_sad_u8(0x05040302, 0x04030201, 100)
self.assertEqual(result, 104)
def test_v_sad_u8_large_diff(self):
"""Test v_sad_u8 with maximum byte differences."""
# s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
result = v_sad_u8(0xffffffff, 0x00000000, 0)
self.assertEqual(result, 1020)
def test_v_msad_u8_basic(self):
"""Test v_msad_u8 masks when reference byte is 0."""
# s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
result = v_msad_u8(0x10101010, 0x00000000, 0)
self.assertEqual(result, 0)
# s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
result = v_msad_u8(0x10101010, 0x01010101, 0)
self.assertEqual(result, 60)
def test_v_msad_u8_partial_mask(self):
"""Test v_msad_u8 with partial masking."""
# s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
# diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
result = v_msad_u8(0x10101010, 0x00010001, 0)
self.assertEqual(result, 30)
def test_v_msad_u8_with_accumulator(self):
"""Test v_msad_u8 with non-zero accumulator."""
result = v_msad_u8(0x10101010, 0x01010101, 50)
self.assertEqual(result, 110) # 60 + 50
if __name__ == '__main__':
unittest.main()
+2 -2
View File
@@ -1,7 +1,7 @@
#!/usr/bin/env python3
"""Test PDF pseudocode extraction from generate.py."""
"""Test PDF pseudocode extraction from amdxml.py."""
import unittest
from extra.assembly.amd.generate import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
from extra.assembly.amd.amdxml import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
EXPECTED_PAGES = {"rdna3": 655, "rdna4": 711, "cdna": 610}
+6 -9
View File
@@ -2,7 +2,7 @@
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re, subprocess, os
from extra.assembly.amd.dsl import Inst
from extra.assembly.amd import decode_inst, detect_format
from extra.assembly.amd.decode import decode_inst, detect_format
from extra.assembly.amd.test.helpers import get_llvm_mc, get_llvm_objdump, get_target, get_mattr
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
@@ -83,21 +83,17 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
arch = self.arch
from extra.assembly.amd.test.test_compare_emulators import get_kernels_from_tinygrad
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
from tinygrad.helpers import AMD_LLVM
from tinygrad.runtime.support.compiler_amd import HIPCompiler
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
# rendered source can be C or llvmir
compiler = (AMDLLVMCompiler if AMD_LLVM else HIPCompiler)(get_target(arch))
compiler = HIPCompiler(get_target(arch))
# First pass: decode all instructions and collect info
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
for ki, kernel in enumerate(kernels):
offset = 0
code = next((s.content for s in elf_loader(compiler.compile(kernel.src))[1] if s.name == ".text"))
while offset < len(code):
remaining = code[offset:]
while offset < len(kernel.code):
remaining = kernel.code[offset:]
fmt = detect_format(remaining, arch)
if fmt is None:
decoded_instrs.append((ki, offset, None, None, None, False, "no format"))
@@ -238,6 +234,7 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
@unittest.skip("RDNA4 decode roundtrip not yet supported")
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
@unittest.skip("CDNA decode roundtrip not yet supported")
+19 -34
View File
@@ -5,11 +5,11 @@ from pathlib import Path
from tinygrad.helpers import DEBUG
from tinygrad.runtime.autogen import rocprof
from tinygrad.runtime.support.elf import elf_loader
from extra.assembly.amd import decode_inst
from extra.assembly.amd.decode import decode_inst
from extra.assembly.amd.autogen.rdna3.ins import SOPP
from extra.assembly.amd.autogen.rdna3.enum import SOPPOp
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_L4, WAVEEND, INST, INST_L4, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES_L3, PACKET_TYPES_L4, InstOp, InstOpL4, print_packets)
from extra.assembly.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, VALUINST, IMMEDIATE, IMMEDIATE_MASK,
ALUEXEC, VMEMEXEC, PACKET_TYPES, InstOp, print_packets)
from extra.assembly.amd.test.helpers import TARGET_TO_ARCH
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
@@ -19,7 +19,6 @@ OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LD
InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
InstOp.OTHER_GLOBAL_STORE_VADDR_128}
OTHER_SIMD_OPS_L4 = {InstOpL4.OTHER_VMEM, InstOpL4.UNK_60}
# ═══════════════════════════════════════════════════════════════════════════════
# ROCPROF DECODER
@@ -92,12 +91,11 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
return occupancy_records, wave_insts
class SQTTExamplesTestBase(unittest.TestCase):
target: str
class TestSQTTExamples(unittest.TestCase):
target = "gfx1100"
@classmethod
def setUpClass(cls):
if cls is SQTTExamplesTestBase: raise unittest.SkipTest("base class")
cls.examples = {}
for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
with open(pkl_path, "rb") as f:
@@ -120,20 +118,18 @@ class SQTTExamplesTestBase(unittest.TestCase):
self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
def test_packet_types_valid(self):
all_classes = set(PACKET_TYPES_L3.values()) | set(PACKET_TYPES_L4.values())
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
for pkt in decode(event.blob):
# Use isinstance to handle layout-specific subclasses (e.g., WAVESTART_L4)
self.assertTrue(any(isinstance(pkt, cls) for cls in all_classes), f"unknown packet type {type(pkt)} in {name}")
self.assertIn(type(pkt), PACKET_TYPES, f"unknown packet type {type(pkt)} in {name}")
def test_wave_lifecycle(self):
for name, (events, *_) in self.examples.items():
if "empty" in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_L4))]), 0, f"no WAVESTART in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, WAVESTART)]), 0, f"no WAVESTART in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
def test_time_monotonic(self):
@@ -148,11 +144,17 @@ class SQTTExamplesTestBase(unittest.TestCase):
if "gemm" not in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_L4))]), 0, f"no INST packets in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, INST)]), 0, f"no INST packets in {name}")
expected: dict[str, list[int]] = {} # override in subclasses
expected = {
"profile_empty_run_0": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_empty_run_1": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_gemm_run_0": [2531, 1844, 1864, 1915, 1942, 1848, 3074, 1919, 1939, 1990, 2017, 1923, 19026, 1919, 1939, 1990, 2017, 1929],
"profile_gemm_run_1": [2554, 1844, 1864, 1915, 1942, 1848, 3084, 1919, 1939, 1990, 2017, 1923, 19010, 1919, 1939, 1990, 2017, 1923],
"profile_plus_run_0": [1900, 1908, 1928, 1979, 2006, 1912],
"profile_plus_run_1": [1856, 1908, 1928, 1979, 2006, 1912],
}
def test_packet_counts(self):
if not self.expected: self.skipTest("no expected packet counts for this target")
for name, (events, *_) in self.examples.items():
with self.subTest(example=name):
if not self.expected.get(name): continue
@@ -176,7 +178,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
for event in events:
wave_starts: dict[tuple[int, int, int], int] = {}
for p in decode(event.blob):
if isinstance(p, (WAVESTART, WAVESTART_L4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
if isinstance(p, WAVESTART): wave_starts[(p.wave, p.simd, p.cu)] = p._time
elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
our_waves.append((wave_starts[key], p._time))
self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
@@ -193,32 +195,15 @@ class SQTTExamplesTestBase(unittest.TestCase):
for event in events:
for p in decode(event.blob):
if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
elif isinstance(p, INST_L4) and p.op not in OTHER_SIMD_OPS_L4: our_insts.append(p._time)
elif isinstance(p, VALUINST): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE_MASK):
for _ in range(bin(p.mask).count('1')): our_insts.append(p._time)
self.assertEqual(sorted(our_insts), sorted(roc_insts), f"instruction times mismatch in {name}")
class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
target = "gfx1100"
expected = {
"profile_empty_run_0": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_empty_run_1": [1803, 1908, 1928, 1979, 2006, 1912],
"profile_gemm_run_0": [2531, 1844, 1864, 1915, 1942, 1848, 3074, 1919, 1939, 1990, 2017, 1923, 19026, 1919, 1939, 1990, 2017, 1929],
"profile_gemm_run_1": [2554, 1844, 1864, 1915, 1942, 1848, 3084, 1919, 1939, 1990, 2017, 1923, 19010, 1919, 1939, 1990, 2017, 1923],
"profile_plus_run_0": [1900, 1908, 1928, 1979, 2006, 1912],
"profile_plus_run_1": [1856, 1908, 1928, 1979, 2006, 1912],
}
#class TestSQTTExamplesRDNA4(TestSQTTExamples): target = "gfx1200"
class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
# CDNA/MI300 (gfx950) uses a completely different 16-bit header packet format, not the nibble-based format.
# See decode_tt_header_stream in ghidra/librocprof-trace-decoder.c - it reads 16-bit headers and uses
# pkt_fmt = header & 0xf to look up packet_class (0x10=2bytes, 0x20=4bytes, 0x30=6bytes, 0x40=8bytes).
# This is NOT implemented yet - the nibble decoder produces garbage for CDNA data.
@unittest.skip("CDNA/MI300 uses 16-bit header format, not nibble-based - decoder not implemented")
class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
target = "gfx950"
#class TestSQTTExamplesCDNA(TestSQTTExamples): target = "gfx950"
if __name__ == "__main__":
unittest.main()
-170
View File
@@ -1,170 +0,0 @@
"""Tests comparing sqtt.py PACKET_TYPES_L3/L4 against AMD's rocprof-trace-decoder binary."""
import unittest, struct, ctypes, pickle
from pathlib import Path
ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
EXAMPLES_DIR = Path(__file__).parent.parent.parent.parent / "sqtt/examples"
def _find_segment(perms: str):
"""Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
with open('/proc/self/maps', 'r') as f:
for line in f:
if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
parts = line.split()
return int(parts[0].split('-')[0], 16), int(parts[2], 16)
return None, None
def _read_array(file_offset: int, count: int):
"""Read an array of uint8 at file_offset from the loaded library."""
base, seg_offset = _find_segment('rw-p')
if base is None: return None
return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
def _load_lib():
if not ROCPROF_LIB.exists(): return False
ctypes.CDLL(str(ROCPROF_LIB))
return True
# ═══════════════════════════════════════════════════════════════════════════════
# RDNA EXTRACTION (nibble-based format)
# ═══════════════════════════════════════════════════════════════════════════════
def extract_bit_tables():
"""Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
if not _load_lib(): return None
return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
def extract_delta_fields():
"""Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
if not _load_lib(): return None
ro_base, ro_offset = _find_segment('r--p')
if ro_base is None: return None
def read_table(file_offset, num_entries):
addr = ro_base + (file_offset - ro_offset)
data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
return {type_id: (lo, hi) for j in range(0, len(data), 12)
for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
def extract_packet_encodings():
"""Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
if not _load_lib(): return None
rw_base, rw_offset = _find_segment('rw-p')
if rw_base is None: return None
# Read base encodings from registration vector at 0x2d340
vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
base = {}
if vec_start and vec_end:
for i in range((vec_end - vec_start) // 32):
addr = vec_start + i * 32
type_id = ctypes.c_uint8.from_address(addr).value
pat_start = ctypes.c_void_p.from_address(addr + 8).value
pat_end = ctypes.c_void_p.from_address(addr + 16).value
if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
pat = list((ctypes.c_uint8 * n).from_address(pat_start))
base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base} # L2 has overrides
# ═══════════════════════════════════════════════════════════════════════════════
# CDNA EXTRACTION (16-bit header format)
# ═══════════════════════════════════════════════════════════════════════════════
def extract_cdna_packet_sizes():
"""Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
from extra.assembly.amd.test.test_sqtt_examples import run_rocprof_decoder
if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
with open(pkl_path, "rb") as f: data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
if not sqtt_events or not prg: return None
# Run decoder to trigger hash table initialization
run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
# Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
rw_base, rw_offset = _find_segment('rw-p')
if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
pkt_sizes, node, seen = {}, head, set()
while node and node not in seen and len(pkt_sizes) < 20:
seen.add(node)
key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
node = ctypes.c_void_p.from_address(node).value
return pkt_sizes if len(pkt_sizes) == 16 else None
# ═══════════════════════════════════════════════════════════════════════════════
# TESTS
# ═══════════════════════════════════════════════════════════════════════════════
class TestSQTTMatchesBinary(unittest.TestCase):
def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
def test_encodings_match_layout3(self): self._test_encodings(3)
def test_encodings_match_layout4(self): self._test_encodings(4)
def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
def test_cdna_packet_sizes(self):
"""Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
pkt_sizes = extract_cdna_packet_sizes()
self.assertIsNotNone(pkt_sizes, "failed to extract CDNA packet sizes")
from extra.assembly.amd.sqtt_cdna import CDNA_PKT_SIZES
for pkt_fmt, size in CDNA_PKT_SIZES.items():
with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
def _test_bit_counts(self, layout: int):
if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
with self.subTest(packet=pkt_cls.__name__):
self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])
def _test_encodings(self, layout: int):
if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
with self.subTest(packet=pkt_cls.__name__):
self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
def _test_delta_fields(self, layout: int):
if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
from extra.assembly.amd.sqtt import PACKET_TYPES_L3, PACKET_TYPES_L4
for type_id, pkt_cls in {3: PACKET_TYPES_L3, 4: PACKET_TYPES_L4}[layout].items():
if type_id not in deltas[layout - 2]: continue
delta = getattr(pkt_cls, 'delta', None)
actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
if __name__ == "__main__":
tables = extract_bit_tables()
encodings = extract_packet_encodings()
deltas = extract_delta_fields()
TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
if encodings and tables:
print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12} {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
print("-" * 140)
for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12} {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
cdna = extract_cdna_packet_sizes()
if cdna: print(f"\nCDNA packet sizes: {cdna}")
unittest.main()
+2 -1
View File
@@ -67,11 +67,12 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))
if not wasm:
thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
for name,cl in bufs_to_save.items():
weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
return '\n'.join(headers + cprog)
else:
if bufs_to_save:
+19 -20
View File
@@ -168,14 +168,12 @@ PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR,
# =============================================================================
class Kernel:
def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
def label(self, name): self.labels[name] = self.pos
def __init__(self, arch='gfx1100'):
self.instructions, self.labels, self.branch_targets, self.arch = [], {}, {}, arch
def emit(self, inst, target=None):
self.instructions.append(inst)
inst._target, inst._pos = target, self.pos
self.pos += inst.size()
return inst
def emit(self, inst): self.instructions.append(inst); return inst
def label(self, name): self.labels[name] = len(self.instructions)
def branch_to(self, label): self.branch_targets[len(self.instructions) - 1] = label
def waitcnt(self, lgkm=None, vm=None):
"""Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
@@ -184,15 +182,16 @@ class Kernel:
self.emit(s_waitcnt(simm16=waitcnt))
def to_asm(self):
# Patch branch offsets: simm16 = (target_pos - branch_end_pos) / 4
for inst in self.instructions:
if inst._target is None: continue
offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
inst.simm16 = offset_dwords
# TODO: replace this with direct ELF
body = ['\t' + inst.disasm() for inst in self.instructions]
import re
# Instruction stream with labels
label_at = {pos: name for name, pos in self.labels.items()}
body = []
for i, inst in enumerate(self.instructions):
if i in label_at: body.append(f'.{label_at[i]}:')
asm = inst.disasm()
if i in self.branch_targets:
asm = re.sub(r'(s_cbranch_\w+|s_branch)\s+\S+', rf'\1 .{self.branch_targets[i]}', asm)
body.append('\t' + asm)
# limit wave occupancy by using more LDS
lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
@@ -316,7 +315,7 @@ def build_kernel(arch='gfx1100'):
k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
# S_LOOP_CTR is already 0 from prologue initialization
k.emit(s_branch(), target='LOOP_ENTRY')
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_ENTRY')
# ===========================================================================
# MAIN GEMM LOOP
@@ -327,12 +326,12 @@ def build_kernel(arch='gfx1100'):
k.label('LOOP_INC')
k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
k.emit(s_cbranch_scc1(), target='EPILOGUE')
k.emit(s_cbranch_scc1(simm16=0)); k.branch_to('EPILOGUE')
k.label('LOOP_ENTRY')
k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0)) # s_cselect doesn't modify SCC
k.emit(s_cbranch_scc0(), target='SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
k.emit(s_cbranch_scc0(simm16=0)); k.branch_to('SKIP_PREFETCH') # branch if loop_ctr >= loop_bound
if not NO_GLOBAL:
# Advance prefetch pointers (VGPR)
@@ -403,7 +402,7 @@ def build_kernel(arch='gfx1100'):
offset = i * 64
k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
k.emit(s_branch(), target='LOOP_INC')
k.emit(s_branch(simm16=0)); k.branch_to('LOOP_INC')
# ===========================================================================
# EPILOGUE: Permute and store results
+1 -2
View File
@@ -3,7 +3,6 @@
import pathlib
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.engine.realize import Estimates
from tinygrad.helpers import getenv
fp = pathlib.Path(__file__).parent/"gemm.s"
@@ -45,7 +44,7 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
sz = UOp.variable("SZ", 256, 8192)
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm", estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
sink = UOp.sink(C.base, A.base, B.base, sz, lidx, gidx, arg=KernelInfo(name="gemm"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
C_asm = Tensor.custom_kernel(C_asm, from_torch(A), from_torch(Bt), fxn=custom_asm_gemm)[0]
File diff suppressed because it is too large Load Diff
+62 -70
View File
@@ -5,7 +5,7 @@ os.environ["AMD_AQL"] = "1"
from tinygrad.device import Device
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from extra.assembly.amd.dsl import Reg, Inst, s, v
from tinygrad.runtime.ops_amd import AMDProgram
NUM_WORKGROUPS = 96
WAVE_SIZE = 32
@@ -17,92 +17,84 @@ DIRECTIVE = ".amdhsa_wavefront_size32 1"
assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> bytes:
preamble = s_mov_b32(counter_sreg, n).to_bytes()
insts_bytes = b"".join([inst.to_bytes() for inst in insts])
sub_inst, cmp_inst = s_sub_u32(counter_sreg, counter_sreg, 1), s_cmp_lg_i32(counter_sreg, 0)
loop_sz = len(insts_bytes) + sub_inst.size() + cmp_inst.size()
branch_inst = s_cbranch_scc1(simm16=-((loop_sz // 4) + 1) & 0xFFFF)
return preamble + insts_bytes + sub_inst.to_bytes() + cmp_inst.to_bytes() + branch_inst.to_bytes() + s_endpgm().to_bytes()
def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs):
def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
if accum:
inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1, acc_cd=1, **kwargs)
instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
vgprIndices[1], vgprIndices[2],
vgprIndices[1], vgprIndices[2], extra)
elif dense:
inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1)
instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
vgprIndices[1], vgprIndices[2],
vgprIndices[1], vgprIndices[2])
else:
inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[3]:vgprIndices[4]], v[vgprIndices[5]])
vgprs:set = set()
for n,_ in inst._fields:
if isinstance(val:=getattr(inst, n), Reg) and val.offset >= v.offset: vgprs |= {val.offset+i for i in range(val.sz)}
inst_bytes = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
inst_hex = "\n".join(" .byte " + ",".join(f"0x{b:02x}" for b in inst_bytes[i:i+16]) for i in range(0, len(inst_bytes), 16)) + "\n"
src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", inst_hex).replace("VGPR_COUNT", str(len(vgprs)))
instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
vgprIndices[1], vgprIndices[2],
vgprIndices[3], vgprIndices[4],
vgprIndices[5])
src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
src = src.replace("DIRECTIVE", DIRECTIVE)
lib = COMPILER.compile(src)
fxn = DEV.runtime("matmul", lib)
fxn = AMDProgram(DEV, "matmul", lib)
elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
if __name__=="__main__":
DEV = Device[Device.DEFAULT]
arch = DEV.renderer.arch
DEVICENUM = os.getenv("DEVICENUM", "0")
try:
DEV = Device['AMD:' + DEVICENUM]
except:
raise RuntimeError("Error while initiating AMD device")
COMPILER = HIPCompiler(arch)
if arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
from extra.assembly.amd.autogen.rdna3.ins import *
if arch == 'gfx1103': NUM_WORKGROUPS = 8
if arch == 'gfx1151': NUM_WORKGROUPS = 32
launchBenchmark(v_wmma_bf16_16x16x16_bf16, (7,8,15))
launchBenchmark(v_wmma_f16_16x16x16_f16, (7,8,15))
launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,15))
launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,15))
launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,9))
launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,11))
elif arch in {'gfx1200', 'gfx1201'}:
from extra.assembly.amd.autogen.rdna4.ins import *
# this instruction does not exist in the rdna4 isa, use the co version
s_sub_u32 = s_sub_co_u32
COMPILER = HIPCompiler(DEV.arch)
if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 32
launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
elif DEV.arch == 'gfx1201':
NUM_WORKGROUPS = 64
launchBenchmark(v_wmma_bf16_16x16x16_bf16, (3,4,7))
launchBenchmark(v_wmma_f16_16x16x16_f16, (3,4,7))
launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,11))
launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,11))
launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,8))
launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_fp8_fp8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_fp8_bf8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_bf8_fp8, (7,8,9))
launchBenchmark(v_wmma_f32_16x16x16_bf8_bf8, (7,8,9))
launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,11))
launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,11))
launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,8))
launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_fp8_fp8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_fp8_bf8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_bf8_fp8", (7,8,9))
launchBenchmark("v_wmma_f32_16x16x16_bf8_bf8", (7,8,9))
FLOPS_PER_MATMUL = 16*16*32*2
launchBenchmark(v_wmma_i32_16x16x32_iu4, (7,8,9))
launchBenchmark(v_swmmac_f32_16x16x32_f16, (7,8,11,12,19,20), False)
launchBenchmark(v_swmmac_f32_16x16x32_bf16, (7,8,11,12,19,20), False)
launchBenchmark(v_swmmac_f16_16x16x32_f16, (3,4,7,8,15,16), False)
launchBenchmark(v_swmmac_bf16_16x16x32_bf16, (3,4,7,8,15,16), False)
launchBenchmark(v_swmmac_i32_16x16x32_iu8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_i32_16x16x32_iu4, (7,8,8,9,10,11), False)
launchBenchmark(v_swmmac_f32_16x16x32_fp8_fp8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_f32_16x16x32_fp8_bf8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_f32_16x16x32_bf8_fp8, (7,8,9,10,13,14), False)
launchBenchmark(v_swmmac_f32_16x16x32_bf8_bf8, (7,8,9,10,13,14), False)
launchBenchmark("v_wmma_i32_16X16X32_iu4", (7,8,9))
launchBenchmark("v_swmmac_f32_16x16x32_f16", (7,8,11,12,19,20), False)
launchBenchmark("v_swmmac_f32_16x16x32_bf16", (7,8,11,12,19,20), False)
launchBenchmark("v_swmmac_f16_16x16x32_f16", (3,4,7,8,15,16), False)
launchBenchmark("v_swmmac_bf16_16x16x32_bf16", (3,4,7,8,15,16), False)
launchBenchmark("v_swmmac_i32_16x16x32_iu8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_i32_16x16x32_iu4", (7,8,8,9,10,11), False)
launchBenchmark("v_swmmac_f32_16x16x32_fp8_fp8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_f32_16x16x32_fp8_bf8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
FLOPS_PER_MATMUL = 16*16*64*2
launchBenchmark(v_swmmac_i32_16x16x64_iu4, (7,8,9,10,13,14), False)
elif arch == 'gfx950':
from extra.assembly.amd.autogen.cdna.ins import *
launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
elif DEV.arch == 'gfx950':
DIRECTIVE = ".amdhsa_accum_offset 4"
NUM_WORKGROUPS = 256
WAVE_SIZE = 64
NUM_WAVES = 4
launchBenchmark(v_mfma_f32_16x16x16_f16, (3,0,1), accum=True)
launchBenchmark(v_mfma_f32_16x16x16_bf16, (3,0,1), accum=True)
launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
FLOPS_PER_MATMUL = 16*16*32*2
launchBenchmark(v_mfma_f32_16x16x32_f16, (3,0,3), accum=True)
launchBenchmark(v_mfma_f32_16x16x32_bf16, (3,0,3), accum=True)
launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
FLOPS_PER_MATMUL = 16*16*128*2
launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,7), accum=True) # fp8
launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,5), accum=True, cbsz=2, blgp=2) # fp6
launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,3), accum=True, cbsz=4, blgp=4) # fp4
launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
else:
raise RuntimeError(f"arch {arch} not supported.")
raise RuntimeError(f"arch {DEV.arch} not supported.")
+10 -3
View File
@@ -3,13 +3,20 @@
.p2align 8
.type matmul,@function
matmul:
s_mov_b32 s1, INTERNAL_LOOP
s_mov_b32 s2, 0
inner_loop:
INSTRUCTION
s_sub_u32 s1, s1, 1
s_cmp_lg_i32 s1, s2
s_cbranch_scc1 inner_loop
s_endpgm
.rodata
.p2align 6
.amdhsa_kernel matmul
.amdhsa_next_free_vgpr VGPR_COUNT
.amdhsa_next_free_sgpr 3
.amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
.amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
DIRECTIVE
.end_amdhsa_kernel
@@ -30,4 +37,4 @@ amdhsa.kernels:
.vgpr_count: 32
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
.end_amdgpu_metadata
+23 -111
View File
@@ -1,7 +1,6 @@
# type: ignore
import ctypes, ctypes.util, struct, platform, pathlib, re, time, os, signal
from tinygrad.helpers import from_mv, to_mv, getenv
from tinygrad.runtime.support.c import init_c_struct_t
from tinygrad.helpers import from_mv, to_mv, getenv, init_c_struct_t
from hexdump import hexdump
start = time.perf_counter()
@@ -11,21 +10,18 @@ processor = platform.processor()
IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]
MMAP_SYSCALL = {"aarch64": 0xde, "x86_64":0x09}[processor]
IOCTL_PRINT = getenv("IOCTL_PRINT", getenv("IOCTL", 0))
GRAB_PMA = getenv("GRAB_PMA", 0)
def get_struct(argp, stype):
return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents
def dump_struct(st):
if IOCTL_PRINT == 0: return
if getenv("IOCTL", 0) == 0: return
print("\t", st.__class__.__name__, end=" { ")
for v in type(st)._real_fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
for v in type(st)._fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
print("}")
def format_struct(s):
sdats = []
for field in s._real_fields_:
for field in s._fields_:
dat = getattr(s, field[0])
if isinstance(dat, int): sdats.append(f"{field[0]}:0x{dat:X}")
else: sdats.append(f"{field[0]}:{dat}")
@@ -62,29 +58,6 @@ def install_hook(c_function, python_function):
return orig_func
# *** ioctl lib end ***
# PMA buffer tracking for raw PC sampling data (only when GRAB_PMA is enabled)
pma_mem_handle = 0 # hMemPmaBuffer from ALLOC_PMA_STREAM
pma_buffer_size = 0
pma_buffer_va = 0 # actual mapped VA (found via /proc/self/maps)
pma_get_offset = 0 # current read offset in ring buffer
pma_pending_map = False # flag to check for new mapping on next ioctl
pma_maps_before = set() # mappings before MAP_MEMORY
pma_raw_dumps: list[bytes] = [] # raw PMA buffer dumps
def get_pma_raw_dumps() -> list[bytes]: return pma_raw_dumps
def clear_pma_raw_dumps(): pma_raw_dumps.clear()
def get_proc_maps():
"""Read current process memory mappings as set of (start, end) tuples."""
result = set()
with open("/proc/self/maps", "r") as f:
for line in f:
addr_range = line.split()[0]
start, end = addr_range.split("-")
result.add((int(start, 16), int(end, 16)))
return result
from tinygrad.runtime.autogen import nv_570 as nv_gpu
nvescs = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("NV_ESC")}
nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(nv_gpu, "struct_"+x.replace("_CMD_", "_")+"_PARAMS", None))) for x in dir(nv_gpu) if \
@@ -96,7 +69,6 @@ def get_classes():
"NV20_SUBDEVICE_0"}
for nm,val in nv_gpu.__dict__.items():
if not isinstance(val, int): continue
if nm.endswith("PARAMETERS_MESSAGE_ID"): continue
if 0x3000 < val < 0xffff: res[val] = nm
if nm in known_classes: res[val] = nm
return res
@@ -108,92 +80,37 @@ global_ioctl_id = 0
gpus_user_modes = []
gpus_mmio = []
gpus_fifo = []
offset_load = 0
@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
def ioctl(fd, request, argp):
global global_ioctl_id, gpus_user_modes, gpus_mmio
global pma_mem_handle, pma_buffer_size, pma_buffer_va, pma_get_offset, pma_pending_map, pma_maps_before
global_ioctl_id += 1
# Check for new PMA buffer mapping from previous MAP_MEMORY call (only when GRAB_PMA is enabled)
if GRAB_PMA and pma_pending_map:
pma_pending_map = False
new_maps = get_proc_maps()
for start, end in new_maps - pma_maps_before:
if end - start == pma_buffer_size:
pma_buffer_va = start
if IOCTL_PRINT >= 1: print(f"\t PMA buffer mapped at CPU VA=0x{pma_buffer_va:x}")
break
st = time.perf_counter()
ret = libc.syscall(IOCTL_SYSCALL, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
et = time.perf_counter()-st
fn = os.readlink(f"/proc/self/fd/{fd}")
#print(f"ioctl {request:8x} {fn:20s}")
idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
if IOCTL_PRINT >= 1: print(f"#{global_ioctl_id}: ", end="")
if getenv("IOCTL", 0) >= 1: print(f"#{global_ioctl_id}: ", end="")
if itype == ord(nv_gpu.NV_IOCTL_MAGIC):
if nr == nv_gpu.NV_ESC_RM_CONTROL:
s = get_struct(argp, nv_gpu.NVOS54_PARAMETERS)
if s.cmd in nvcmds:
name, struc = nvcmds[s.cmd]
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
print(f"NV_ESC_RM_CONTROL cmd={name:30s} hClient={s.hClient}, hObject={s.hObject}, flags={s.flags}, params={s.params}, paramsSize={s.paramsSize}, status={s.status}")
if struc is not None: dump_struct(get_struct(s.params, struc))
elif hasattr(nv_gpu, name+"_PARAMS"): dump_struct(get_struct(argp, getattr(nv_gpu, name+"_PARAMS")))
elif name == "NVA06C_CTRL_CMD_GPFIFO_SCHEDULE": dump_struct(get_struct(argp, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS))
elif name == "NV83DE_CTRL_CMD_GET_MAPPINGS": dump_struct(get_struct(s.params, nv_gpu.NV83DE_CTRL_DEBUG_GET_MAPPINGS_PARAMETERS))
elif name == "NVB0CC_CTRL_CMD_SET_HS_CREDITS":
hs_params = get_struct(s.params, nv_gpu.NVB0CC_CTRL_SET_HS_CREDITS_PARAMS)
dump_struct(hs_params)
if IOCTL_PRINT >= 2:
for i in range(hs_params.numEntries):
print(f"\t\t", end="")
dump_struct(hs_params.creditInfo[i])
# PMA buffer tracking (only when GRAB_PMA is enabled)
if GRAB_PMA and name == "NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM":
pma_params = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS)
pma_mem_handle = pma_params.hMemPmaBuffer
pma_buffer_size = pma_params.pmaBufferSize
pma_get_offset = 0 # Reset read offset for new stream
if IOCTL_PRINT >= 1: print(f"\t PMA buffer: hMem=0x{pma_mem_handle:x} size={pma_buffer_size}")
if GRAB_PMA and name == "NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT":
pma_update = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS)
if pma_update.bytesAvailable > 0 and pma_buffer_va and pma_buffer_size > 0:
avail = pma_update.bytesAvailable
read_offset = pma_get_offset
# Handle ring buffer wrap-around
if pma_get_offset + avail <= pma_buffer_size:
pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, avail))
else:
# Wrap around: read to end, then from start
first_part = pma_buffer_size - pma_get_offset
second_part = avail - first_part
pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, first_part)) + bytes(to_mv(pma_buffer_va, second_part))
pma_raw_dumps.append(pma_data)
pma_get_offset = (pma_get_offset + avail) % pma_buffer_size
if IOCTL_PRINT >= 2:
print(f"\t PMA data: {avail} bytes from offset=0x{read_offset:x}, new offset=0x{pma_get_offset:x}")
hexdump(pma_data)
# Dump regOps for EXEC_REG_OPS when IOCTL >= 3
if name == "NVB0CC_CTRL_CMD_EXEC_REG_OPS" and struc is not None and IOCTL_PRINT >= 3:
reg_params = get_struct(s.params, struc)
for i in range(reg_params.regOpCount):
print(f"\t\t", end="")
dump_struct(reg_params.regOps[i])
# val = (op.regValueHi << 32) | op.regValueLo
# print(f"\t regOps[{i:3d}]: op={op.regOp} type={op.regType} status={op.regStatus} offset=0x{op.regOffset:08x} value=0x{val:016x}")
else:
if IOCTL_PRINT >= 1: print("unhandled cmd", hex(s.cmd))
if getenv("IOCTL", 0) >= 1: print("unhandled cmd", hex(s.cmd))
# format_struct(s)
# print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
elif nr == nv_gpu.NV_ESC_RM_ALLOC:
s = get_struct(argp, nv_gpu.NVOS21_PARAMETERS)
if IOCTL_PRINT >= 1: print(f"NV_ESC_RM_ALLOC hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
if getenv("IOCTL", 0) >= 1: print(f"NV_ESC_RM_ALLOC hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
if s.pAllocParms is not None:
if s.hClass == nv_gpu.NV01_DEVICE_0: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV0080_ALLOC_PARAMETERS))
if s.hClass == nv_gpu.FERMI_VASPACE_A: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS))
@@ -201,7 +118,6 @@ def ioctl(fd, request, argp):
if s.hClass == nv_gpu.NV1_MEMORY_USER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
if s.hClass == nv_gpu.NV1_MEMORY_SYSTEM: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
if s.hClass == nv_gpu.GT200_DEBUGGER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV83DE_ALLOC_PARAMETERS))
if s.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE: dump_struct(get_struct(s.pAllocParms, nv_gpu.NVB2CC_ALLOC_PARAMETERS))
if s.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
sx = get_struct(s.pAllocParms, nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS)
dump_struct(sx)
@@ -210,35 +126,31 @@ def ioctl(fd, request, argp):
if s.hClass == nv_gpu.TURING_USERMODE_A: gpus_user_modes.append(s.hObjectNew)
elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
# nv_ioctl_nvos33_parameters_with_fd
s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
print(f"NV_ESC_RM_MAP_MEMORY hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, length={s.length} flags={s.flags} pLinearAddress={s.pLinearAddress}")
# Track PMA buffer mapping - save maps now, check for new mapping on next ioctl (after mmap happens)
if GRAB_PMA and pma_mem_handle and s.hMemory == pma_mem_handle:
pma_maps_before = get_proc_maps()
pma_pending_map = True
elif nr == nv_gpu.NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO:
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
s = get_struct(argp, nv_gpu.NVOS56_PARAMETERS)
print(f"NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, pOldCpuAddress={s.pOldCpuAddress} pNewCpuAddress={s.pNewCpuAddress} status={s.status}")
elif nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
s = get_struct(argp, nv_gpu.nv_ioctl_nvos02_parameters_with_fd)
print(f"NV_ESC_RM_ALLOC_MEMORY fd={s.fd}, hRoot={s.params.hRoot}, hObjectParent={s.params.hObjectParent}, hObjectNew={s.params.hObjectNew}, hClass={s.params.hClass}, flags={s.params.flags}, pMemory={s.params.pMemory}, limit={s.params.limit}, status={s.params.status}")
elif nr == nv_gpu.NV_ESC_ALLOC_OS_EVENT:
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
s = get_struct(argp, nv_gpu.nv_ioctl_alloc_os_event_t)
print(f"NV_ESC_ALLOC_OS_EVENT hClient={s.hClient} hDevice={s.hDevice} fd={s.fd} Status={s.Status}")
elif nr == nv_gpu.NV_ESC_REGISTER_FD:
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
s = get_struct(argp, nv_gpu.nv_ioctl_register_fd_t)
print(f"NV_ESC_REGISTER_FD fd={s.ctl_fd}")
elif nr in nvescs:
if IOCTL_PRINT >= 1: print(nvescs[nr])
if getenv("IOCTL", 0) >= 1: print(nvescs[nr])
else:
if IOCTL_PRINT >= 1: print("unhandled NR", nr)
if getenv("IOCTL", 0) >= 1: print("unhandled NR", nr)
elif fn.endswith("nvidia-uvm"):
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
print(f"{nvuvms.get(request, f'UVM UNKNOWN {request=}')}")
if nvuvms.get(request) is not None: dump_struct(get_struct(argp, getattr(nv_gpu, nvuvms.get(request)+"_PARAMS")))
if nvuvms.get(request) == "UVM_MAP_EXTERNAL_ALLOCATION":
@@ -247,7 +159,7 @@ def ioctl(fd, request, argp):
print("perGpuAttributes[{i}] = ", end="")
dump_struct(st.perGpuAttributes[i])
if IOCTL_PRINT >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
if getenv("IOCTL") >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
return ret
@ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long)
@@ -260,14 +172,14 @@ def _mmap(addr, length, prot, flags, fd, offset):
return ret
install_hook(libc.ioctl, ioctl)
if getenv("IOCTL") >= 4: orig_mmap_mv = install_hook(libc.mmap, _mmap)
if getenv("IOCTL") >= 3: orig_mmap_mv = install_hook(libc.mmap, _mmap)
import collections
old_gpputs = collections.defaultdict(int)
def _dump_gpfifo(mark):
launches = []
print("_dump_gpfifo:", mark)
# print("_dump_gpfifo:", mark)
for start, size in gpus_fifo:
gpfifo_controls = nv_gpu.AmpereAControlGPFifo.from_address(start+size*8)
gpfifo = to_mv(start, size * 8).cast("Q")
@@ -293,7 +205,7 @@ def make_qmd_struct_type():
fields.append((name.replace("NVC6C0_QMDV03_00_", "").lower(), ctypes.c_uint32, data[0]-data[1]+1))
if len(fields) >= 2 and fields[-2][0].endswith('_lower') and fields[-1][0].endswith('_upper') and fields[-1][0][:-6] == fields[-2][0][:-6]:
fields = fields[:-2] + [(fields[-1][0][:-6], ctypes.c_uint64, fields[-1][2] + fields[-2][2])]
return init_c_struct_t(0x40 * 4, tuple(fields))
return init_c_struct_t(tuple(fields))
qmd_struct_t = make_qmd_struct_type()
assert ctypes.sizeof(qmd_struct_t) == 0x40 * 4
@@ -310,7 +222,7 @@ def _dump_qmd(address, packets):
subc = (dat>>13) & 7
mthd = (dat<<2) & 0x7FFF
method_name = nvqcmds.get(mthd, f"unknown method #{mthd}")
if IOCTL_PRINT >= 1:
if getenv("IOCTL", 0) >= 1:
print(f"\t\t{method_name}, {typ=} {size=} {subc=} {mthd=}")
for j in range(size): print(f"\t\t\t{j}: {gpfifo[i+j+1]} | 0x{gpfifo[i+j+1]:x}")
if mthd == 792:
-1
View File
@@ -1 +0,0 @@
examples/
-135
View File
@@ -1,135 +0,0 @@
import pickle, os, sys, functools, numpy as np
from pathlib import Path
os.environ["DEV"] = "CUDA"
os.environ["PROFILE"] = os.environ.get("PROFILE", "2")
from extra.nv_pma.cupti import cu_prof_ext
cu_prof_ext.enable_auto()
from tinygrad import Tensor, Device
if not os.environ.get("IOCTL") or not os.environ.get("GRAB_PMA"):
print("Usage: GRAB_PMA=1 IOCTL=1 IOCTL_PRINT=0 python3 extra/nv_pma/collect.py")
sys.exit(1)
assert Device.DEFAULT == "CUDA", "only works with CUDA"
EXAMPLES_DIR = Path(__file__).parent / "examples"
_collectors: list[tuple[str, callable]] = []
def pcsampling_test(name: str):
def decorator(fn):
@functools.wraps(fn)
def wrapper():
cu_prof_ext.clear_pma_raw_dumps()
cu_prof_ext.clear_cupti_pc_samples()
fn()
Device["CUDA"].synchronize()
dumps = cu_prof_ext.get_pma_raw_dumps()
# from hexdump import hexdump
# hexdump(dumps[0][:0x40])
return {"test_name": name, "pma_raw_dumps": list(cu_prof_ext.get_pma_raw_dumps()), "cupti_pc_samples": list(cu_prof_ext.get_cupti_pc_samples())}
_collectors.append((name, wrapper))
return wrapper
return decorator
# Refs
@pcsampling_test("test_plus")
def test_plus():
a = Tensor([1, 2, 3, 4])
b = Tensor([5, 6, 7, 8])
(a + b).realize()
@pcsampling_test("test_matmul")
def test_matmul():
a = Tensor(np.random.rand(12, 12).astype(np.float32))
b = Tensor(np.random.rand(12, 12).astype(np.float32))
(a @ b).realize()
@pcsampling_test("test_reduce_sum")
def test_reduce_sum():
a = Tensor(np.random.rand(1024).astype(np.float32))
a.sum().realize()
@pcsampling_test("test_reduce_max")
def test_reduce_max():
a = Tensor(np.random.rand(1024).astype(np.float32))
a.max().realize()
@pcsampling_test("test_exp")
def test_exp():
a = Tensor(np.random.rand(256).astype(np.float32))
a.exp().realize()
@pcsampling_test("test_softmax")
def test_softmax():
a = Tensor(np.random.rand(64, 64).astype(np.float32))
a.softmax().realize()
@pcsampling_test("test_conv2d")
def test_conv2d():
x = Tensor(np.random.rand(1, 3, 32, 32).astype(np.float32))
w = Tensor(np.random.rand(8, 3, 3, 3).astype(np.float32))
x.conv2d(w).realize()
@pcsampling_test("test_large_matmul")
def test_large_matmul():
a = Tensor(np.random.rand(128, 128).astype(np.float32))
b = Tensor(np.random.rand(128, 128).astype(np.float32))
(a @ b).realize()
@pcsampling_test("test_elementwise_chain")
def test_elementwise_chain():
a = Tensor(np.random.rand(512).astype(np.float32))
((a + 1) * 2 - 0.5).relu().realize()
@pcsampling_test("test_broadcast")
def test_broadcast():
a = Tensor(np.random.rand(64, 1).astype(np.float32))
b = Tensor(np.random.rand(1, 64).astype(np.float32))
(a + b).realize()
@pcsampling_test("test_plus_big")
def test_plus_big():
a = Tensor(np.random.rand(64, 32).astype(np.float32))
b = Tensor(np.random.rand(64, 32).astype(np.float32))
(a + b).realize()
def save_example(name: str, data: dict):
pma_bytes = sum(len(d) for d in data['pma_raw_dumps'])
cupti_samples = sum(r['samples'] for r in data['cupti_pc_samples'])
print(f" PMA: {len(data['pma_raw_dumps'])} buffers, {pma_bytes} bytes")
print(f" CUPTI: {len(data['cupti_pc_samples'])} records, {cupti_samples} samples")
outfile = EXAMPLES_DIR / f"{name}.pkl"
with open(outfile, "wb") as f:
pickle.dump(data, f)
print(f" Saved to {outfile}")
if __name__ == "__main__":
EXAMPLES_DIR.mkdir(exist_ok=True)
# Run specific tests if provided as arguments, otherwise run all
if len(sys.argv) > 1:
test_names = sys.argv[1:]
collectors = [(name, fn) for name, fn in _collectors if name in test_names]
if not collectors:
print(f"Unknown tests: {test_names}")
print(f"Available: {[name for name, _ in _collectors]}")
sys.exit(1)
else:
collectors = _collectors
for name, collect_fn in collectors:
print(f"\nCollecting {name}...")
try:
data = collect_fn()
save_example(name, data)
except Exception as e:
print(f" ERROR: {e}")
import traceback
traceback.print_exc()
-26
View File
@@ -1,26 +0,0 @@
# CUPTI autogen loader for nv_pma
# To regenerate: REGEN=1 python -c "import extra.nv_pma.cupti"
import importlib, pathlib
from tinygrad.helpers import getenv
root = pathlib.Path(__file__).parents[3]
here = pathlib.Path(__file__).parent
def load(name, dll, files, **kwargs):
if not (f:=here/f"{name}.py").exists() or getenv('REGEN'):
kwargs['args'] = kwargs.get('args', [])
f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(name, dll, files, **kwargs))
return importlib.import_module(f"extra.nv_pma.cupti.{name}")
def __getattr__(nm):
match nm:
case "cupti":
return load("cupti", "'/usr/local/cuda/targets/x86_64-linux/lib/libcupti.so'", [
"/usr/local/cuda/include/cupti_result.h", "/usr/local/cuda/include/cupti_activity.h",
"/usr/local/cuda/include/cupti_callbacks.h", "/usr/local/cuda/include/cupti_events.h",
"/usr/local/cuda/include/cupti_metrics.h", "/usr/local/cuda/include/cupti_driver_cbid.h",
"/usr/local/cuda/include/cupti_runtime_cbid.h", "/usr/local/cuda/include/cupti_profiler_target.h",
"/usr/local/cuda/include/cupti_profiler_host.h", "/usr/local/cuda/include/cupti_pmsampling.h",
"/usr/local/cuda/include/generated_cuda_meta.h", "/usr/local/cuda/include/generated_cuda_runtime_api_meta.h"
], args=["-D__CUDA_API_VERSION_INTERNAL", "-I/usr/local/cuda/include"], parse_macros=False)
case _: raise AttributeError(f"no such autogen: {nm}")
-164
View File
@@ -1,164 +0,0 @@
from __future__ import annotations
import ctypes
from tinygrad.helpers import DEBUG, getenv
from extra.nv_pma.cupti import cupti
def stall_reason_name(reason: int) -> str:
name = cupti.CUpti_ActivityPCSamplingStallReason.get(reason)
return name.replace("CUPTI_ACTIVITY_PC_SAMPLING_STALL_", "").lower() if name else str(reason)
class CUPTIProfiler:
def __init__(self):
self.initialized = False
self.pc_sampling_enabled = False
self.buffers: list[ctypes.Array] = []
self.kernel_stalls: dict[int, dict[int, int]] = {}
self.raw_buffers: list[bytes] = []
self.pc_samples: list[dict] = []
def _check_cupti(self, status, soft=False):
if status != cupti.CUPTI_SUCCESS:
if soft: return False
raise RuntimeError(f"CUPTI Error {status}")
return True
def init(self, ctx, device_id: int = 0, profile_level: int = 2):
if self.initialized: return
# Initialize profiler API
init_params = cupti.CUpti_Profiler_Initialize_Params()
init_params.structSize = 16
cupti.cuptiProfilerInitialize(ctypes.byref(init_params))
# Register buffer callbacks for Activity API
self._buf_req_cb = cupti.CUpti_BuffersCallbackRequestFunc(self._buffer_requested)
self._buf_comp_cb = cupti.CUpti_BuffersCallbackCompleteFunc(self._buffer_completed)
self._check_cupti(cupti.cuptiActivityRegisterCallbacks(self._buf_req_cb, self._buf_comp_cb))
# PROFILE=1: kernel timing, PROFILE=2: PC sampling with stall reasons
if profile_level >= 2:
# PC sampling for stall analysis (requires elevated privileges)
if DEBUG >= 1: print(" CUPTI: PC sampling mode (before)")
pc_status = cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING)
if pc_status == cupti.CUPTI_SUCCESS:
config = cupti.CUpti_ActivityPCSamplingConfig()
config.size, config.samplingPeriod = 16, cupti.CUPTI_ACTIVITY_PC_SAMPLING_PERIOD_MIN
cfg_status = cupti.dll.cuptiActivityConfigurePCSampling(ctx, ctypes.byref(config))
if cfg_status == cupti.CUPTI_SUCCESS:
if DEBUG >= 1: print(" CUPTI: PC sampling mode (before stall analysis)")
cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO)
self.pc_sampling_enabled = True
if DEBUG >= 1: print(" CUPTI: PC sampling mode (stall analysis)")
elif cfg_status == 35:
if DEBUG >= 1: print(" CUPTI: PC sampling needs: echo 'options nvidia NVreg_RestrictProfilingToAdminUsers=0'|sudo tee /etc/modprobe.d/nvidia.conf && sudo reboot")
# Fall back to kernel timing if PC sampling setup failed
if not self.pc_sampling_enabled:
self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
else:
# Kernel activity tracing for timing
self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
self.initialized = True
def _buffer_requested(self, buffer, size, max_num_records):
buf = (ctypes.c_uint8 * 1024 * 1024)() # 1MB buffer
self.buffers.append(buf)
buffer[0] = ctypes.cast(buf, ctypes.POINTER(ctypes.c_uint8))
size[0] = ctypes.sizeof(buf)
max_num_records[0] = 0
def _buffer_completed(self, ctx, stream_id, buffer, size, valid_size):
if valid_size > 0:
record = ctypes.POINTER(cupti.CUpti_Activity)()
while cupti.cuptiActivityGetNextRecord(buffer, valid_size, ctypes.byref(record)) == cupti.CUPTI_SUCCESS:
kind = record.contents.kind
if kind == cupti.CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL:
kernel = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityKernel9)).contents
name = ctypes.string_at(kernel.name).decode() if kernel.name else "unknown"
duration_us = (kernel.end - kernel.start) / 1000.0
grid, block = (kernel.gridX, kernel.gridY, kernel.gridZ), (kernel.blockX, kernel.blockY, kernel.blockZ)
print(f" CUPTI: {name[:40]:40s} | {duration_us:10.2f} us | grid={grid} block={block} | regs={kernel.registersPerThread:3d} smem={kernel.staticSharedMemory + kernel.dynamicSharedMemory:6d}B")
elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING:
pc = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSampling3)).contents
cid = pc.correlationId
if cid not in self.kernel_stalls: self.kernel_stalls[cid] = {}
self.kernel_stalls[cid][pc.stallReason] = self.kernel_stalls[cid].get(pc.stallReason, 0) + pc.samples
self.pc_samples.append({
'correlationId': pc.correlationId, 'pcOffset': pc.pcOffset, 'stallReason': pc.stallReason,
'samples': pc.samples, 'latencySamples': pc.latencySamples, 'functionId': pc.functionId, 'sourceLocatorId': pc.sourceLocatorId
})
if DEBUG >= 3:
print(f" PC {pc.pcOffset:#x} stall={stall_reason_name(pc.stallReason)} samples={pc.samples} latency={pc.latencySamples} func={pc.functionId} src={pc.sourceLocatorId}")
elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO:
info = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSamplingRecordInfo)).contents
cid = info.correlationId
if cid in self.kernel_stalls:
stalls = self.kernel_stalls[cid]
total = sum(stalls.values())
if total > 0:
top = sorted(stalls.items(), key=lambda x: -x[1])[:5]
stall_str = " ".join(f"{stall_reason_name(r)}:{100*c//total}%" for r,c in top if c > 0)
print(f" CUPTI stalls (corr={cid}): {total} samples | {stall_str}")
del self.kernel_stalls[cid]
else: print(f" CUPTI: Unhandled activity kind {kind}")
def flush(self):
if not self.initialized: return
self._check_cupti(cupti.cuptiActivityFlushAll(0))
# Module-level profiler instance
_profiler: CUPTIProfiler | None = None
def get_profiler() -> CUPTIProfiler | None:
return _profiler
def get_cupti_raw_buffers() -> list[bytes]:
return _profiler.raw_buffers if _profiler else []
def clear_cupti_raw_buffers():
if _profiler: _profiler.raw_buffers.clear()
def get_cupti_pc_samples() -> list[dict]:
return _profiler.pc_samples if _profiler else []
def clear_cupti_pc_samples():
if _profiler: _profiler.pc_samples.clear()
# Raw PMA buffer access (from ioctl interception)
def get_pma_raw_dumps() -> list[bytes]:
try:
from extra.nv_gpu_driver.nv_ioctl import get_pma_raw_dumps as _get
return _get()
except ImportError: return []
def clear_pma_raw_dumps():
try:
from extra.nv_gpu_driver.nv_ioctl import clear_pma_raw_dumps as _clear
_clear()
except ImportError: pass
def enable(profile_level:int=2):
global _profiler
if _profiler is not None: return
_profiler = CUPTIProfiler()
# Patch CUDADevice to initialize CUPTI profiler
from tinygrad.runtime.ops_cuda import CUDADevice
_orig_init = CUDADevice.__init__
_orig_sync = CUDADevice.synchronize
def _patched_init(self, device: str):
_orig_init(self, device)
device_id = int(device.split(":")[1]) if ":" in device else 0
_profiler.init(self.context, device_id, profile_level)
def _patched_sync(self):
_orig_sync(self)
if _profiler: _profiler.flush()
CUDADevice.__init__ = _patched_init
CUDADevice.synchronize = _patched_sync
def enable_auto():
if (profile_level:=getenv("PROFILE", 0)) > 0: enable(profile_level)
File diff suppressed because it is too large Load Diff
+2 -3
View File
@@ -1,13 +1,12 @@
#!/bin/bash
set -e
export PAGE_SIZE=1
export PYTHONPATH=.
export LOGOPS=/tmp/ops
export CAPTURE_PROCESS_REPLAY=1
rm "$LOGOPS" 2>/dev/null || true
rm $LOGOPS
test/external/process_replay/reset.py
CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/unit/test_winograd.py test/models/test_real_world.py --durations=20
CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
CL=1 python3 -m pytest test/test_tiny.py
# extract, sort and uniq
+1 -1
View File
@@ -42,7 +42,7 @@ def get_struct(argp, stype):
def format_struct(s):
sdats = []
for field_name, *_ in s._real_fields_:
for field_name, field_type in s._fields_:
if field_name in {"__pad", "PADDING_0"}: continue
dat = getattr(s, field_name)
if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
+2 -3
View File
@@ -11,7 +11,7 @@ from tinygrad import Tensor
from tinygrad.helpers import system, OSX
from tinygrad.runtime.ops_amd import AMDProgram
from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
from tinygrad.device import Device
from tinygrad.device import Device, ProfileDeviceEvent
from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
@@ -23,7 +23,7 @@ def save_sqtt():
dev.profile_events.clear()
sqtt:dict[str, list[WaveExec]] = {}
yield sqtt
events = dev.profile_events
events = dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())]
#rctx = decode(events)
#assert len(rctx.inst_execs) > 0, "empty sqtt output"
@@ -45,7 +45,6 @@ matmul:
.rodata
.p2align 6
.amdhsa_kernel matmul
.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
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
-1
View File
@@ -159,7 +159,6 @@ def main() -> None:
if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
n = 0
for s in trace["steps"]:
if "PKTS" in s["name"]: continue
print(s["name"])
data = viz.get_render(s["query"])
print_data(data)
+87
View File
@@ -0,0 +1,87 @@
import os
os.environ["PROFILE"] = "1"
os.environ["PMC"] = "1"
import unittest
import functools, contextlib
import numpy as np
from tinygrad import Tensor, Context, Device
from tinygrad.dtype import dtypes, AddrSpace
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.runtime.ops_amd import ProfilePMCEvent
from tinygrad.engine.realize import get_runner
from tinygrad.viz.serve import unpack_pmc
from extra.sqtt.roc import print_pmc
def copy_kernel(B, A, stride=1):
n_threads = 32
assert A.size >= n_threads, f"{A.size} is too small, min size {n_threads}"
g = UOp.range(A.size//n_threads, 0, AxisType.GLOBAL)
l = UOp.range(n_threads, 1, AxisType.LOCAL)
i = g * n_threads + l
index = (i * stride) % A.size
return B[index].store(A[index]).sink(arg=KernelInfo(name=f"copy_{A.size}_stride_{stride}", opts_to_apply=()))
def lds_kernel(offset:UOp, size:int, inst:str) -> UOp:
tid = UOp.range(offset.size, 0, AxisType.LOCAL)
dst = UOp.placeholder((size,), dtypes.float32, 1, AddrSpace.REG)
#lds = UOp.placeholder((1024,), dtypes.float32, 2, AddrSpace.LOCAL)
u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_s_waitcnt(0);')
u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_s_barrier();', src=(u,))
u = UOp(Ops.CUSTOM, arg='__builtin_amdgcn_sched_barrier(0);', src=(u,))
u = UOp(Ops.CUSTOM, arg=f'asm volatile("{inst} '+'%0, %1" : "=v"({0}) : "v"({1}));', src=(dst, offset[tid], u))
return UOp.sink(u, arg=KernelInfo(name="test_lds", opts_to_apply=()))
dev = Device[Device.DEFAULT]
@contextlib.contextmanager
def save_pmc():
# clear the old traces
dev.profile_events.clear()
pmc:list[ProfilePMCEvent] = []
yield pmc
for e in dev.profile_events:
if isinstance(e, ProfilePMCEvent): pmc.append(e)
@unittest.skipIf(dev.device != "AMD", "tests PMC counters on AMD")
class TestPMC(unittest.TestCase):
@Context(IGNORE_OOB=0)
def test_copy(self, stride:int=1):
N = 1 << 25 # ~134MB
a = Tensor(np.arange(N, dtype=np.uint32)+1).realize()
b = Tensor(np.zeros(N, dtype=np.uint32)).realize()
b = Tensor.custom_kernel(b, a, fxn=functools.partial(copy_kernel, stride=stride))[0]
with save_pmc() as pmc:
b.realize()
print_pmc(pmc)
np.testing.assert_equal(a.numpy(), b.numpy())
def test_copy_uncoalesced(self): return self.test_copy(stride=17)
# test with two threads issuing ds_reads at different offsets
def test_ds_read(self, size=1, inst='ds_read_b32'):
test_banks = 256
offsets = [Tensor([0, b*4]) for b in range(1, test_banks)]
with Context(DEBUG=0): Tensor.realize(*offsets)
k = Tensor.custom_kernel(offsets[0], fxn=functools.partial(lds_kernel, size=size, inst=inst))[0]
# sample all kernels
with save_pmc() as pmc_events:
runner = get_runner(Device.DEFAULT, k.schedule()[0].ast)
# TODO: llvm eliminates lds definition from the ELF, is there another way to pin lds size?
runner._prg.group_segment_size = 1024
for offset in offsets: runner([offset.uop.buffer])
# find read offsets that created bank conflicts from the pmc counters
found:list[Tensor] = []
for i,e in enumerate(pmc_events):
pmc = unpack_pmc(e)["rows"]
# SQ on gfx9, renamed to SQC after gfx10
val = next(total for name,total,_all_instances in pmc if name in {"SQ_LDS_BANK_CONFLICT", "SQC_LDS_BANK_CONFLICT"})
if val > 0: found.append(offsets[i])
print("Found bank conflicts at offsets:", [s.numpy() for s in found])
def test_ds_read_b64(self): self.test_ds_read(2, 'ds_read_b64')
def test_ds_read_b128(self): self.test_ds_read(4, 'ds_read_b128')
if __name__ == "__main__":
unittest.main()
+143
View File
@@ -0,0 +1,143 @@
import os
os.environ["PYTHONPATH"] = "."
os.environ["SQTT"] = "1"
if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
os.environ["PROFILE"] = "1"
# VIZ=1 to launch server
# os.environ["VIZ"] = "1"
os.environ["AMD_LLVM"] = "0"
import unittest
import sys, contextlib
from tinygrad import Tensor, dtypes
from tinygrad.helpers import getenv
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.device import Device, ProfileDeviceEvent
from extra.sqtt.roc import decode, WaveExec
dev = Device[os.environ["DEV"]]
def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
name = sys._getframe(1).f_code.co_name
def fxn(_):
L = UOp.special(l, "lidx0")
G = UOp.special(g, "gidx0")
op = custom("asm volatile (")
for inst in instrs: op = custom(f' "{inst}\\n\\t"', op)
op = custom(");", op)
return UOp.sink(op, L, G, arg=KernelInfo(name=name))
k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
return k
@contextlib.contextmanager
def save_sqtt():
# clear the old traces
dev.profile_events.clear()
sqtt:dict[str, list[WaveExec]] = {}
yield sqtt
# decode sqtt
if os.environ["DEV"] != "AMD": return
rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
assert len(rctx.inst_execs) > 0, "empty sqtt output"
sqtt.update(rctx.inst_execs)
class TestTiming(unittest.TestCase):
def test_v_add(self):
with save_sqtt() as sqtt:
asm_kernel([f"v_add_f32 v{10+i} v{10+i+1} {10+i}" for i in range(3)]).realize()
wave = list(sqtt.values())[0][:-1]
assert all(s.dur == 1 for s in wave)
assert all(s.stall == 0 for s in wave)
def test_chain_v_add_1l(self):
with save_sqtt() as sqtt:
asm_kernel([
"v_add_f32_e32 v1 v0 v0",
"v_add_f32_e32 v2 v1 v1",
]).realize()
wave = list(sqtt.values())[0][:-1]
assert all(s.dur == 1 for s in wave)
assert all(s.stall == 0 for s in wave)
def test_multi_cycle_inst(self):
def custom_vrcp(A, B):
op = custom("float a = 0.0;")
op = custom("float b = (*(data1_1+0));", op)
#op = custom('asm volatile("v_mul_f32_e32 %2 %2 %1" : "+v"(a) : "v"(b));', op)
op = custom('asm volatile("v_rcp_f32_e32 %2 %1" : "+v"(a) : "v"(b));', op)
op = custom('asm volatile("v_add_f32_e64 %1 %1 1.0" : "+v"(a));', op)
op = custom("*(data0_1+0) = a;", op)
return UOp.sink(op, A, B, arg=KernelInfo(name="custom_vrcp"))
out = Tensor([0.]).realize()
inp = Tensor([-2.0]).realize()
with save_sqtt() as sqtt:
Tensor.custom_kernel(out, inp, fxn=custom_vrcp)[0].realize()
wave = list(sqtt.values())[0][0]
for i in range(len(wave.insts)):
if wave.insts[i].inst.startswith("global_store"):
print(f"store diff {wave.insts[i].time-(wave.insts[i-1].time)}")
self.assertEqual(out.item(), 0.5)
def test_wmma(self):
with save_sqtt() as sqtt:
for tc in dev.renderer.get_tensor_cores(dev.arch):
M, K, N = tc.dims
s = 32
a = Tensor.empty(M*s, K*s, dtype=tc.dtype_in)@Tensor.empty(K*s, N*s, dtype=tc.dtype_in)
a.realize()
print(a)
for p,waves in sqtt.items():
for e in waves[0].insts:
if (e.inst.startswith("v_wmma")):
instruction = e.inst.split(" ")[0]
print(f"{instruction:<29} : {e.dur} cycles")
def test_sleep(self):
n = 1
def sleep_kernel(data0):
assert data0.dtype.base == dtypes.ulong
op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
op = custom("unsigned long long t1 = __builtin_readcyclecounter();", op)
op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
with save_sqtt() as sqtt:
diff_hw_reg.realize()
sleep = next((e for e in sqtt[f"sleep_{n}"][0].insts if e.inst.startswith("s_sleep")))
# cycles = sleep dur + overhead of storing hi/lo REG_SHADER_CYCLES
self.assertGreaterEqual(diff_hw_reg.item(), sleep.dur)
def test_nop(self):
with save_sqtt() as sqtt:
asm_kernel(["s_nop 1"]*10).realize()
wave = list(sqtt.values())[0][0]
for e in wave.insts:
print(f"{e.inst} {e.dur=} {e.stall=}")
def test_wave_sched(self):
num_waves = getenv("NUM_WAVES", 16)
num_wgps = getenv("NUM_WGPS", 2)
num_vgpr = getenv("NUM_VGPR", 256)
with save_sqtt() as sqtt:
# 1 cycle decode, no stall
asm_kernel([f"v_mov_b32_e32 v{i} {i}" for i in range(num_vgpr)], l=32*num_waves, g=num_wgps).realize()
waves = list(sqtt.values())[0]
print(len(waves), "waves decoded")
for w in waves:
print(f"{w.wave_id:<2} {w.simd=} {w.cu=} {w.se=} @ clk {w.begin_time}")
def test_ones(self):
N = getenv("N", 4096)
CNT = getenv("CNT", 2)
with save_sqtt() as sqtt:
for _ in range(CNT):
Tensor.ones(N, N).contiguous().realize()
self.assertEqual(len(sqtt), CNT)
if __name__ == "__main__":
unittest.main()
+98 -27
View File
@@ -1,7 +1,6 @@
import math
from tinygrad import Tensor, dtypes
from tinygrad.helpers import DEBUG
from tinygrad.uop.ops import UOp
from extra.thunder.tiny.tk import WARP_THREADS
@@ -13,10 +12,10 @@ Q_BLOCK_SIZE = 16
KV_BLOCK_SIZE = 16
def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
if not isinstance(ref.device, tuple): return Tensor.empty(*shape, device=ref.device)
shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
axis = ref.uop.axis if axis is None else axis
return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.multi(axis), dtype=ref.dtype, device=ref.device)
return Tensor(Tensor.empty(*shape, device=ref.device).uop.multi(axis), device=ref.device)
def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
return _sharded_empty(ref.shape, ref, axis)
@@ -39,12 +38,10 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
B, N, H, D = xq.shape
H_KV = xk.shape[2]
GROUP_SIZE = H // H_KV
num_devices = len(xq.device) if isinstance(xq.device, tuple) else 1
B_local = B // num_devices
if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
print(f"Flash Attention {B=} {N=} {H=} {D=} {H_KV=} {GROUP_SIZE=}")
def custom_forward(ou:UOp, l_vecu:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp) -> UOp:
with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
with Kernel("fa_custom_forward", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
o, q, k, v, mask, l_vec = GL(ou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker), GL(l_vecu, ker)
@@ -142,7 +139,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
return ker.finish()
def custom_backward_q(dqu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_q", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
with Kernel("fa_custom_backward_q", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dq, do, q, k, v, mask = GL(dqu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
@@ -221,21 +218,21 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
dp_block -= delta_vec_reg
att_block *= dp_block
att_block *= 1.0 / math.sqrt(D)
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_mma = warp.copy(att_block_mma.after(att_block), att_block)
dq_reg = warp.mma_AB(dq_reg, k_reg_col_t, att_block_mma)
dq_reg = ker.endrange()
dq_reg *= 1.0 / math.sqrt(D)
dq_reg_transposed = warp.transpose(dq_reg_transposed, dq_reg)
dq = warp.store(dq, dq_reg_transposed, (batch, q_seq, head, 0), axis=1)
return ker.finish()
def custom_backward_kv(dku:UOp, dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp):
with Kernel("fa_custom_backward_kv", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B_local), NUM_WORKERS * WARP_THREADS) as ker:
def custom_backward_k(dku:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp, delta_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_k", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dk, dv, do, q, k, v, mask = GL(dku, ker), GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
dk, do, q, k, v, mask = GL(dku, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
l_vec, delta_vec = GL(l_vecu, ker), GL(delta_vecu, ker)
head_kv = ker.blockIdx_x
@@ -256,7 +253,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dk_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
@@ -270,7 +266,6 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
delta_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
dk_reg = warp.zero(dk_reg)
dv_reg = warp.zero(dv_reg)
# load kv tile
k_reg = warp.load(k_reg, k, (), (batch, kv_seq, head_kv, 0), axis=1)
@@ -309,32 +304,107 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
att_block -= l_vec_reg
att_block = att_block.exp2()
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dv_reg_ = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
dp_block = warp.zero(dp_block.after(g, q_idx, dv_reg_))
dp_block = warp.zero(dp_block.after(g, q_idx))
dp_block = warp.mma_ABt(dp_block, v_reg, do_reg)
dp_block -= delta_vec_reg
att_block *= dp_block
att_block *= 1.0 / math.sqrt(D)
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dk_reg = warp.mma_AB(dk_reg, att_block_row, q_reg_col)
dk_reg = ker.endrange(2)
dv_reg = dv_reg.after(dk_reg)
dk_reg *= 1.0 / math.sqrt(D)
dk = warp.store(dk, dk_reg, (batch, kv_seq, head_kv, 0), axis=1)
return ker.finish()
def custom_backward_v(dvu:UOp, dou:UOp, qu:UOp, ku:UOp, vu:UOp, masku:UOp, l_vecu:UOp) -> UOp:
with Kernel("fa_custom_backward_v", (H_KV, N // (KV_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
dv, do, q, k, v, mask = GL(dvu, ker), GL(dou, ker), GL(qu, ker), GL(ku, ker), GL(vu, ker), GL(masku, ker)
l_vec = GL(l_vecu, ker)
head_kv = ker.blockIdx_x
batch = ker.blockIdx_z
kv_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
q_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_smem = ker.st((Q_BLOCK_SIZE, D), dtypes.bfloat16)
att_smem = ker.st((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_t = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
k_reg_t = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
mask_reg = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.float32)
mask_reg_transposed = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
dv_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.float32, TileLayout.COL)
do_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
do_reg_col = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
att_block_transposed = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
att_block_row = ker.rt((Q_BLOCK_SIZE, KV_BLOCK_SIZE), dtypes.bfloat16)
l_vec_reg = ker.rv(Q_BLOCK_SIZE, dtypes.float32)
dv_reg = warp.zero(dv_reg)
# load kv tile
k_reg = warp.load(k_reg, k, (), (batch, kv_seq, head_kv, 0), axis=1)
k_reg_t = warp.transpose(k_reg_t, k_reg)
v_reg = warp.load(v_reg, v, (), (batch, kv_seq, head_kv, 0), axis=1)
for q_idx in ker.range(N // Q_BLOCK_SIZE):
for g in ker.range(GROUP_SIZE):
head_q = head_kv * GROUP_SIZE + g
# load q and do
q_smem = warp.load(q_smem, q, (), (batch, q_idx, head_q, 0), axis=1)
do_smem = warp.load(do_smem, do, (), (batch, q_idx, head_q, 0), axis=1)
q_reg = warp.load(q_reg, q_smem)
q_reg_t = warp.transpose(q_reg_t, q_reg)
do_reg = warp.load(do_reg, do_smem)
do_reg_col = warp.load(do_reg_col, do_smem)
# load l_vec
l_vec_reg = warp.load(l_vec_reg, l_vec, (), (batch, head_q, 0, q_idx), axis=2)
l_vec_reg *= 1.0 / math.log(2)
# mma qk^t
att_block = warp.zero(att_block.after(g))
att_block = warp.mma_AtB(att_block, k_reg_t, q_reg_t)
att_block *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
# apply attention mask
mask_reg = warp.load(mask_reg, mask, (), (batch, 0, q_idx, kv_seq), axis=2)
mask_reg_transposed = warp.transpose(mask_reg_transposed, mask_reg)
att_block += mask_reg_transposed
att_block -= l_vec_reg
att_block = att_block.exp2()
att_block_mma = warp.copy(att_block_mma, att_block)
att_block_transposed = warp.transpose(att_block_transposed, att_block_mma)
att_smem = warp.store(att_smem, att_block_transposed)
att_block_row = warp.load(att_block_row, att_smem)
dv_reg = warp.mma_AB(dv_reg, att_block_row, do_reg_col)
dv_reg = ker.endrange(2)
dv_reg = warp.map(dv_reg, lambda x, idx: x + v_reg[*idx].cast(dtypes.float32) * 1e-30)
dk = warp.store(dk, dk_reg, (batch, kv_seq, head_kv, 0), axis=1)
dv = warp.store(dv, dv_reg, (batch, kv_seq, head_kv, 0), axis=1)
return ker.finish(2)
return ker.finish()
single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device
@@ -362,7 +432,8 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
delta_vec = (grad * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
grad_q = Tensor.custom_kernel(grad_q, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_q)[0]
grad_k, grad_v = Tensor.custom_kernel(grad_k, grad_v, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_kv)[:2]
grad_k = Tensor.custom_kernel(grad_k, grad, xq, xk, xv, attn_mask, l_vec, delta_vec, fxn=custom_backward_k)[0]
grad_v = Tensor.custom_kernel(grad_v, grad, xq, xk, xv, attn_mask, l_vec, fxn=custom_backward_v)[0]
return (None, None, grad_q.uop, grad_k.uop, grad_v.uop, None)
attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, attn_mask, fxn=custom_forward, grad_fxn=grad)[:2]
+3 -1
View File
@@ -189,11 +189,13 @@ class Group:
self.ker.push_store(c_store, c)
return c.after(c_store).reshape(c.shape)
map_rid = 400
def map(self, a:ALL_TILES, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):
a = cast(UOp, a)
assert self.warps == 1
rngs_for_shape = tuple(self.ker.raw_range(dim) for dim in a.shape)
rngs_for_shape = tuple(UOp.range(dim, Group.map_rid + i) for i, dim in enumerate(a.shape))
Group.map_rid += len(a.shape)
if op.__code__.co_argcount == 1:
to_store = op(a[*rngs_for_shape]) # type: ignore
+13 -3
View File
@@ -163,6 +163,19 @@ def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None)
def randperm_generator(n, generator=None, out=None):
return out.copy_(wrap(Tensor.randperm(n, generator=generator, device=unwrap(out).device)))
@torch.library.impl("aten::cummax", "privateuseone")
def cummax(self, dim):
values, indices = unwrap(self).cummax(dim)
return (wrap(values), wrap(indices.cast(dtypes.int64)))
@torch.library.impl("aten::cummin", "privateuseone")
def cummin(self, dim):
values, indices = unwrap(self).cummin(dim)
return (wrap(values), wrap(indices.cast(dtypes.int64)))
@torch.library.impl("aten::nonzero", "privateuseone")
def nonzero(self): return wrap(unwrap(self).nonzero())
@torch.library.impl("aten::_linalg_eigh", "privateuseone")
# TODO: move to tinygrad
def _linalg_eigh(self, UPLO: str = 'U'):
@@ -627,9 +640,6 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
"device": _from_torch_device(device) if device else None}.items() if v is not None}),
"aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
"aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
"aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
"aten.nonzero": Tensor.nonzero,
"aten.unfold": Tensor.unfold,
}}
+1 -1
View File
@@ -338,7 +338,7 @@ class TestHCQ(unittest.TestCase):
et = float(sig_en.timestamp - sig_st.timestamp)
print(f"exec kernel time: {et:.2f} us")
assert 0.1 <= et <= (3000000 if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)
assert 0.1 <= et <= (100000 if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)
def test_speed_copy_bandwidth(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
+7 -40
View File
@@ -65,15 +65,6 @@ class TestMainOnnxOps(TestOnnxOps):
outputs = ["y"]
self.helper_test_single_op("Conv", inputs, attributes, outputs, atol=1e-4)
def test_pad_constant_value_zero(self):
from tinygrad.nn.onnx import onnx_ops
Pad = onnx_ops["Pad"]
x = Tensor.arange(4).reshape(1, 1, 2, 2).float()
pads = [0, 0, 1, 1, 0, 0, 1, 1]
out = Pad(x, pads, constant_value=0, value=3)
expected = x.pad((pads[3], pads[7], pads[2], pads[6], pads[1], pads[5], pads[0], pads[4]), value=0)
self.assertEqual(out.tolist(), expected.tolist())
def test_gather(self):
# test const negative indices
inputs = {
@@ -87,7 +78,8 @@ class TestMainOnnxOps(TestOnnxOps):
def test_gather_jit_different_indices(self):
# Gather should not assume indices is const when it can change at runtime
from tinygrad import TinyJit
from tinygrad.nn.onnx import onnx_ops
from tinygrad.nn.onnx import onnx_ops, _cached_to_python_const
_cached_to_python_const.cache_clear()
Gather = onnx_ops["Gather"]
x = Tensor([10, 20, 30, 40, 50])
@@ -101,32 +93,6 @@ class TestMainOnnxOps(TestOnnxOps):
def gather_jit(x, indices): return Gather(x, indices)
self.assertEqual([gather_jit(x, Tensor(idx)).tolist() for idx in indices_list], expected)
def test_gather_jit_const_zero_index(self):
# Gather with const index=0 (falsy in Python) should work with JIT cache
from tinygrad import TinyJit
# Create model: y = Gather(x, 0) + x where 0 is from initializer
# The Add ensures there's a kernel to JIT
x_input = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (5,))
y_output = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, (5,))
idx_init = onnx.numpy_helper.from_array(np.array(0, dtype=np.int64), name="idx")
graph = onnx.helper.make_graph([
onnx.helper.make_node("Gather", ["x", "idx"], ["g"], axis=0),
onnx.helper.make_node("Add", ["g", "x"], ["y"])],
"test_gather_zero", [x_input], [y_output], [idx_init])
model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 13)])
with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmp:
onnx.save(model, tmp.name)
runner = OnnxRunner(tmp.name)
@TinyJit
def run_gather(x): return runner({"x": x})["y"]
# Run multiple times - JIT capture should use cached index=0 correctly
for val in [[10, 20, 30, 40, 50], [100, 200, 300, 400, 500], [1, 2, 3, 4, 5]]:
result = run_gather(Tensor(val, dtype=dtypes.float32))
np.testing.assert_equal(result.numpy(), np.array(val) + val[0])
# NOTE: resize OP is sensitive to numerical errors
def _test_resize_scales(self, scale_values, **kwargs):
for sc in scale_values:
@@ -172,9 +138,9 @@ class TestMainOnnxOps(TestOnnxOps):
self._test_if(np.array([[1, 2, 3], [4, 5, 6]]).astype(np.float32), np.array([[6, 5, 4, 3, 2, 1]]).astype(np.float32))
def test_if_jit_different_shapes(self):
# TODO: If with different output shapes and non-const condition should raise
# When shapes differ, Python selection evaluates condition at graph build time, breaking JIT
from tinygrad import TinyJit
from tinygrad.engine.jit import JitError
# then: x+1 shape (3,), else: x[:2]+1 shape (2,)
x_input = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, (3,))
then_out = onnx.helper.make_tensor_value_info("res", onnx.TensorProto.FLOAT, (3,))
@@ -203,9 +169,10 @@ class TestMainOnnxOps(TestOnnxOps):
def run_if(cond, x): return runner({"cond": cond, "x": x})["res"]
x = Tensor([1.0, 2.0, 3.0])
with self.assertRaises(JitError):
for _ in range(3):
run_if(Tensor([True]), x)
self.assertEqual(run_if(Tensor([True]), x).tolist(), [2, 3, 4]) # x + 1
self.assertEqual(run_if(Tensor([True]), x).tolist(), [2, 3, 4])
self.assertEqual(run_if(Tensor([True]), x).tolist(), [2, 3, 4])
self.assertEqual(run_if(Tensor([False]), x).tolist(), [2, 3, 4]) # wrong! should be [2, 3]
def test_resize_downsample_scales_linear_align_corners(self):
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-131
+1 -1
View File
@@ -9,7 +9,7 @@ from tinygrad.engine.realize import method_cache
from tinygrad.helpers import Profiling
class FakeProgram:
def __init__(self, name:str, prg:bytes, **kwargs): pass
def __init__(self, name:str, prg:bytes): pass
def __call__(self, *bufs, global_size, local_size, vals=(), wait=False): pass
class FakeAllocator(Allocator[Compiled]):
+1
View File
@@ -74,6 +74,7 @@ if __name__ == "__main__":
apply_movement_op.cache_clear()
_apply_reshape.cache_clear()
fold_divmod_general.cache_clear()
UOp.const.cache_clear()
Tensor._device_seeds.clear()
Tensor._device_rng_counters.clear()
+1 -1
View File
@@ -55,7 +55,7 @@ if __name__ == "__main__":
with Context(CORRECT_DIVMOD_FOLDING=1):
simplified_expr = expr.simplify()
solver = z3.Solver(ctx=z3.Context())
solver = z3.Solver()
solver.set(timeout=5000) # some expressions take very long verify, but its very unlikely they actually return sat
z3_expr, z3_simplified_expr, v1, v2, v3 = uops_to_z3(solver, expr, simplified_expr, u1, u2, u3)
check = solver.check(z3_simplified_expr != z3_expr)
+3 -1
View File
@@ -69,7 +69,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
def not_support_multi_device():
# CL and CUDA don't support multi device if in CI
return CI and Device.DEFAULT in ("CL", "CUDA")
return CI and REAL_DEV in ("CL", "CUDA")
def needs_second_gpu(fn):
@functools.wraps(fn)
@@ -79,3 +79,5 @@ def needs_second_gpu(fn):
except Exception as e: self.skipTest(f"second device not available: {e}")
return fn(self, *args, **kwargs)
return wrapper
REAL_DEV = Device.DEFAULT
-11
View File
@@ -1,7 +1,6 @@
import pathlib, re, ctypes, mmap, collections, functools, copy, os
import tinygrad.runtime.autogen.kfd as kfd
import tinygrad.runtime.autogen.am.am as am
import tinygrad.runtime.autogen.amdgpu_drm as amdgpu_drm
from tinygrad.helpers import from_mv
from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props
@@ -34,16 +33,6 @@ class DRMFileDesc(VirtFileDesc):
super().__init__(fd)
self.driver, self.gpu = driver, gpu
def ioctl(self, fd, request, argp):
struct = amdgpu_drm.struct_drm_amdgpu_info.from_address(argp)
if struct.query == amdgpu_drm.AMDGPU_INFO_DEV_INFO:
dev_info = amdgpu_drm.struct_drm_amdgpu_info_device.from_address(struct.return_pointer)
# mock of gfx1100
for se in range(4):
for sa in range(4): dev_info.cu_bitmap[se][sa] = 0xff if (se * 4 + sa) < 12 else 0
return 0
raise NotImplementedError(f"unknown DRM ioctl query {struct.query}")
def mmap(self, start, sz, prot, flags, fd, offset): return libc.mmap(start, sz, prot, flags|mmap.MAP_ANONYMOUS, -1, 0)
class AMDDriver(VirtDriver):
+1 -10
View File
@@ -9,7 +9,6 @@ SDMA_MAX_COPY_SIZE = 0x400000
regCOMPUTE_PGM_LO = 0x1bac + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_PGM_RSRC2 = 0x1bb3 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_TMPRING_SIZE = 0x1bb8 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_USER_DATA_0 = 0x1be0 + amd_gpu.GC_BASE__INST0_SEG0
regCOMPUTE_NUM_THREAD_X = 0x1ba7 + amd_gpu.GC_BASE__INST0_SEG0
regGRBM_GFX_INDEX = 0x2200 + amd_gpu.GC_BASE__INST0_SEG1
@@ -186,18 +185,10 @@ class PM4Executor(AMDQueue):
for st,sz in self.gpu.mapped_ranges:
if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
# Get scratch size from COMPUTE_TMPRING_SIZE register
# For gfx11: WAVESIZE = ceildiv(64 * size_per_thread, 256), so size_per_thread ≈ WAVESIZE * 256 / 64 = WAVESIZE * 4
try: tmpring_size = self.gpu.regs[regCOMPUTE_TMPRING_SIZE]
except KeyError: tmpring_size = 0
wavesize = (tmpring_size >> 12) & 0x3FFF # WAVESIZE field is bits 12:25 for gfx11
scratch_size = wavesize * 4 # This gives the scratch size per thread (lane)
assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
# Pass valid memory ranges, rsrc2, and scratch_size to Python emulator
# Pass valid memory ranges and rsrc2 to Python emulator for bounds checking and SGPR/VGPR layout
if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
if hasattr(remu, 'rsrc2'): remu.rsrc2 = rsrc2
if hasattr(remu, 'scratch_size'): remu.scratch_size = scratch_size
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")
+5 -4
View File
@@ -19,15 +19,16 @@ class PythonRemu:
"""Python RDNA3 emulator wrapper that matches the libremu.so interface."""
valid_mem_ranges: set[tuple[int, int]] = set()
rsrc2: int = 0x19c # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
scratch_size: int = 0 # private_segment_fixed_size from kernel descriptor
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.amd.emu import run_asm
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size)
from extra.assembly.amd.emu import run_asm, set_valid_mem_ranges
# Pad ranges to handle GPU loads that may read past small buffers (e.g. s_load_b128 on 12-byte buffer)
set_valid_mem_ranges({(start, size + 4096) for start, size in self.valid_mem_ranges})
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2)
def _try_dlopen_remu():
# Use Python emulator only if PYTHON_REMU=1
if int(getenv("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"]
+7 -17
View File
@@ -19,7 +19,6 @@ NVAllocation = collections.namedtuple('NVAllocation', ['device', 'size', 'is_sig
NVChannelGroup = collections.namedtuple('NVChannelGroup', ['device'])
NVContextShare = collections.namedtuple('NVContextShare', ['channel_group'])
NVGPFIFO = collections.namedtuple('NVGPFIFO', ['device', 'token'])
NVProfiler = collections.namedtuple('NVProfiler', ['subdevice'])
class NVCtlFileDesc(VirtFileDesc):
def __init__(self, fd, driver):
@@ -141,10 +140,6 @@ class NVDriver(VirtDriver):
struct.hObjectNew = self._alloc_handle()
elif struct.hClass == nv_gpu.GT200_DEBUGGER:
struct.hObjectNew = self._alloc_handle()
elif struct.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE:
assert struct.hObjectParent in self.object_by_handle and isinstance(self.object_by_handle[struct.hObjectParent], NVSubDevice)
struct.hObjectNew = self._alloc_handle()
self.object_by_handle[struct.hObjectNew] = NVProfiler(self.object_by_handle[struct.hObjectParent])
else: raise RuntimeError(f"Unknown {struct.hClass} to rm_alloc")
return 0
@@ -209,14 +204,6 @@ class NVDriver(VirtDriver):
elif struct.cmd == nv_gpu.NV0000_CTRL_CMD_SYSTEM_GET_BUILD_VERSION_V2:
params = nv_gpu.NV0000_CTRL_SYSTEM_GET_BUILD_VERSION_V2_PARAMS.from_address(params_ptr)
params.driverVersionBuffer = b"570.00.00\0"
elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_GR_GET_TPC_MASK:
params = nv_gpu.NV2080_CTRL_GR_GET_TPC_MASK_PARAMS.from_address(params_ptr)
params.tpcMask = 0x1 # one TPC
# Profiler commands - just pass through for mockgpu
elif struct.cmd in (nv_gpu.NVB0CC_CTRL_CMD_POWER_REQUEST_FEATURES, nv_gpu.NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM,
nv_gpu.NVB0CC_CTRL_CMD_RESERVE_HWPM_LEGACY, nv_gpu.NVB0CC_CTRL_CMD_RESERVE_PM_AREA_PC_SAMPLER,
nv_gpu.NVB0CC_CTRL_CMD_BIND_PM_RESOURCES, nv_gpu.NVB0CC_CTRL_CMD_SET_HS_CREDITS,
nv_gpu.NVB0CC_CTRL_CMD_EXEC_REG_OPS, nv_gpu.NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT): pass
else: raise RuntimeError(f"Unknown {struct.cmd} to rm_control")
return 0
@@ -226,11 +213,14 @@ class NVDriver(VirtDriver):
elif nr == nv_gpu.NV_ESC_RM_CONTROL: return self.rm_control(argp)
elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
st:Any = nv_gpu.nv_ioctl_nvos33_parameters_with_fd.from_address(argp)
obj = self.object_by_handle.get(st.params.hMemory)
file = self.opened_fds.get(st.fd)
if isinstance(obj, NVUserMode) and isinstance(file, NVDevFileDesc):
obj = self.object_by_handle[st.params.hMemory]
if isinstance(obj, NVUserMode):
file = self.opened_fds[st.fd]
assert isinstance(file, NVDevFileDesc)
file._mapping_userland = True
elif isinstance(obj, NVAllocation) and obj.is_signal and isinstance(file, NVDevFileDesc):
elif isinstance(obj, NVAllocation) and obj.is_signal:
file = self.opened_fds[st.fd]
assert isinstance(file, NVDevFileDesc)
file._mapping_signal = True
elif nr == nv_gpu.NV_ESC_RM_FREE:
st = nv_gpu.NVOS00_PARAMETERS.from_address(argp)
-29
View File
@@ -24,7 +24,6 @@ binary_operations = [operator.add, operator.sub, operator.mul, operator.lt, oper
integer_binary_operations = binary_operations + [(Tensor.bitwise_xor, np.bitwise_xor), (Tensor.bitwise_and, np.bitwise_and),
(Tensor.bitwise_or, np.bitwise_or), (Tensor.maximum, np.maximum), operator.mod]
integer_unary_operations = [operator.neg]
unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin),
(Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal), (Tensor.cos, np.cos)]
@@ -182,34 +181,6 @@ class TestDTypeALU(unittest.TestCase):
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
@given(ht.uint8, strat.sampled_from(integer_unary_operations))
def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
@unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}")
@given(ht.uint16, strat.sampled_from(integer_unary_operations))
def test_uint16_unary(self, a, op): universal_test_unary(a, dtypes.uint16, op)
@unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}")
@given(ht.uint32, strat.sampled_from(integer_unary_operations))
def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op)
@unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
@given(ht.int8, strat.sampled_from(integer_unary_operations))
def test_int8_unary(self, a, op): universal_test_unary(a, dtypes.int8, op)
@given(ht.int16, strat.sampled_from(integer_unary_operations))
def test_int16_unary(self, a, op): universal_test_unary(a, dtypes.int16, op)
@given(ht.int32, strat.sampled_from(integer_unary_operations))
def test_int32_unary(self, a, op): universal_test_unary(a, dtypes.int32, op)
@unittest.skipUnless(is_dtype_supported(dtypes.int64), f"no int64 on {Device.DEFAULT}")
@given(ht.int64, strat.sampled_from(integer_unary_operations))
def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
@given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)
+4 -3
View File
@@ -4,10 +4,11 @@ from tinygrad import Device, dtypes, Tensor, Context
from tinygrad.device import LRUAllocator, is_dtype_supported
from tinygrad.dtype import ImageDType
from tinygrad.helpers import prod, unwrap
from test.helpers import REAL_DEV
IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageCopy(unittest.TestCase):
def test_image_copyout_1x8(self, img_type=dtypes.imagef):
it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
@@ -41,7 +42,7 @@ class TestImageCopy(unittest.TestCase):
assert (it == it2).sum().item() == prod(sz)
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageDType(unittest.TestCase):
def test_image_pitch(self):
def __validate(imgdt, expected_pitch):
@@ -198,7 +199,7 @@ class TestImageDType(unittest.TestCase):
self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
self.assertEqual(len(sched), 9)
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
@unittest.skipUnless(REAL_DEV in IMAGE_SUPPORTED_DEVICES, "Images not supported")
class TestImageRealization(unittest.TestCase):
def test_image_dtype_expand(self):
data = Tensor.randn(9*32*4).realize()
+5 -4
View File
@@ -3,7 +3,7 @@ import unittest, functools
import numpy as np
from hypothesis import given, settings, strategies as strat
from test.helpers import assert_jit_cache_len, not_support_multi_device, needs_second_gpu
from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV, needs_second_gpu
from tinygrad.tensor import Tensor
from tinygrad.engine.jit import TinyJit, JitError, GraphRunner, MultiGraphRunner, graph_class
from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
@@ -23,7 +23,7 @@ def _simple_test(add, extract=lambda x: x, N=10):
class TestJit(unittest.TestCase):
@settings(deadline=2e4)
@unittest.skipUnless(Device.DEFAULT in ["CPU"], f"no support on {Device.DEFAULT}")
@unittest.skipUnless(REAL_DEV in ["CPU"], f"no support on {REAL_DEV}")
@given(strat.sampled_from([Tensor.exp2, Tensor.log2, Tensor.sin]))
def test_approx_jit_timeout(self, op):
with Context(TRANSCENDENTAL=2):
@@ -260,6 +260,7 @@ class TestJit(unittest.TestCase):
assert len(res3) == 5, "All values should be different, rand works in jit."
assert res3 != res2, "Jit rand is diff with diff seeds"
@unittest.expectedFailure # TODO: fix
def test_jit_v_nojit_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
@@ -289,7 +290,7 @@ class TestJit(unittest.TestCase):
with_jit.add(o1.numpy()[0][0])
with_jit.add(o2.numpy()[0][0])
assert len(with_jit) == 10, "All values should be different."
assert with_jit != without_jit, "TODO: fix. jit and non-jit should produce the same random values with the same seed"
assert with_jit == without_jit, "Jit rand produced different values from no jit."
def test_jit_multiple_random_regen(self):
def f(a, b):
@@ -783,7 +784,7 @@ class TestJitGraphSplit(unittest.TestCase):
def test_jit_multidev_xfer(self):
if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU is not a valid default device for this test (zero-copies)")
if Device.DEFAULT == "METAL": raise unittest.SkipTest("Metal is flaky, with multidevice (same as metal llama 4gpu?)")
if Device.DEFAULT == "METAL" or REAL_DEV == "METAL": raise unittest.SkipTest("Metal is flaky, with multidevice (same as metal llama 4gpu?)")
try: Device[f"{Device.DEFAULT}:1"]
except Exception: raise unittest.SkipTest("no multidevice")

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