forked from tinygrad/tinygrad
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
48dd1d6543 | ||
|
|
a06023cd36 | ||
|
|
ac232bceb5 | ||
|
|
f19fbadce4 | ||
|
|
8eb762d6fa |
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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!")
|
||||
@@ -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()
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
-32
@@ -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
|
||||
+7
-11
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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],
|
||||
}
|
||||
@@ -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
@@ -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
File diff suppressed because it is too large
Load Diff
+779
-1009
File diff suppressed because it is too large
Load Diff
+72
-144
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)."""
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)."""
|
||||
|
||||
@@ -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).
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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):
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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}
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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.")
|
||||
|
||||
@@ -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
@@ -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 +0,0 @@
|
||||
examples/
|
||||
@@ -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()
|
||||
@@ -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}")
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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.
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
}}
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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]):
|
||||
|
||||
Vendored
+1
@@ -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()
|
||||
|
||||
|
||||
Vendored
+1
-1
@@ -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
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,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
@@ -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
Reference in New Issue
Block a user