forked from tinygrad/tinygrad
assembly/amd: replace pcode with ucode (#14002)
* a bunch of todos for my boy claude
* uops have types
* lil cleanups
* simpler ucode
* isNAN
* calls
* move more
* cleanup pcode_parse
* cvt functions
* fix parser bugs
* no void
* minmax
* more pcode parse
* pretty print
* transform
* comments
* move to transform
* assign/declare
* simpler norm
* single PM
* just Uops
* simpler
* more typed
* all rewrite
* less verbose
* work
* spec
* transform
* work
* simpler spec
* less spec
* bitcast
* simpler
* simp ucode
* work
* more in pcode_transform
* remove junk
* more functions
* bug
* no void assign
* load/store
* wave
* fixes
* move denorm
* move more functions
* tests
* cat is shape None
* uop syntax
* move a few more
* program_spec
* cat stuff
* assign fix clear
* unused
* nans
* fp bits
* works with simplify
* remove junk
* special
* meh
* more
* more
* update test pcode parse
* improve parser
* parse some for loops
* merge master
* dead files
* tests pass
* emu2
* better emu2
* test_plus works
* uselessly write more instructions
* use pcode
* something
* something
* bench_emu
* progress
* ds works
* work
* work
* more passing
* run compare
* bench_emu
* more pcode
* a few more
* bugfixes
* bugfix
* test fixes
* tests pass without USE_HW
* all hw tests pass
* add more hw tests
* new hw tests
* bit
* less handcode
* parse more
* consolidate pcode
* fixes
* rsrc
* lane pcode
* cleanups
* simpler
* emu bugs
* one cmp test fails
* fix decode and upd name
* fix name and test harness
* _ftz_f32
* fix denorm
* fix VOPD and use load
* fix carry bug
* no load where / just invalid
* clean
* simpler
* merge sops
* refactoring
* simplifications
* bugfixes
* new tests
* f16 sin fix
* assertion and hw tests
* cvt functions
* one more failure
* bugfixes
* bugfix + regression
* more tests
* fmac
* no manual unrolling
* ordering
* LLVM backend is a lot faster
* compile inst
* more bugs
* f16
* bugfix
* fix regression
* one clang call
* 1M inst
* scratch works
* do scratch correctly
* cleanup
* regression
* cmp
* fmamk fixes
* merge
* fix vcmpx
* unify memory
* remove unused code
* ignore oob for test
* cleanups
* fix mbs
* unify cmp
* test
* minor cleanups
* bump timeout
* fix tests
* revert the CMPLE stuff
* remove opt
* less diff
* simpler
* revert
* support multiple backends
* memset is a lot faster
* split out in bench emu
* improve timing
* timing
* cache that
* cache that
* simpler and faster
* tokenize
* binop table
* simpler
* move to parser
* tok for lambda
* refactor
* expr_parser
* delete emu2_pcode
* import cleanup
* lil
* if parse
* work
* simpler
* no v
* trig preop is faster
* durations for tests
* fix cmp bug
* sdst
* remove scartch_size hack
* null behavior
* _MXCSRContext
* bugfixes
* DEBUG >= 3
* test smem crashes my gpu
* debug
* test
* test smem
* profiler
* full inst
* bugfix
* rtag(1)
* pc is 64-bit and word
* pc is real code now
* dynamic
* more dynamic
* fix oob access
* fix crash, more dyn
* all dyn
* really all dyn
* correct null mask
* lit + format
* 21s on the tests
* 13s on the tests
* canonical name
* simm16
* more dyn
* 14s
* proper saddr dedup
* dyn
* debug 5
* better 5
* revert dynamic stuff
* that can be dyn
* negative offsets
* dyn wmma
* f16 wmma support / ops / dtype / dtype_alu
* symbolic changes not needed
* ConstFloat
* more uop.const
* __eq__
* uop tests
* fix f16
* bf16 tensor cores
* whitespace
* remove cast roundtrip
* Revert "remove cast roundtrip"
This reverts commit c5bb0381c3.
* just the fix
* remove dead paths
* llvm runs
This commit is contained in:
George Hotz
@@ -659,7 +659,11 @@ jobs:
|
||||
testamdasm:
|
||||
name: AMD ASM IDE
|
||||
runs-on: ubuntu-24.04
|
||||
timeout-minutes: 10
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
AMD: 1
|
||||
PYTHON_REMU: 1
|
||||
MOCKGPU: 1
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
uses: actions/checkout@v4
|
||||
@@ -685,16 +689,16 @@ jobs:
|
||||
- name: Install rocprof-trace-decoder
|
||||
run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
|
||||
- name: Run RDNA3 emulator tests
|
||||
run: python -m pytest -n=auto extra/assembly/amd/ --durations 20
|
||||
run: AMD_LLVM=0 python -m pytest -n=auto extra/assembly/amd/ --durations 20
|
||||
- name: Run RDNA3 emulator tests (AMD_LLVM=1)
|
||||
run: AMD_LLVM=1 python -m pytest -n=auto extra/assembly/amd/ --durations 20
|
||||
- name: Run RDNA3 dtype tests
|
||||
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
|
||||
run: AMD_LLVM=0 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
|
||||
- name: Run RDNA3 dtype tests (AMD_LLVM=1)
|
||||
run: AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py
|
||||
run: AMD_LLVM=1 pytest -n=auto test/test_dtype_alu.py test/test_dtype.py --durations 20
|
||||
# TODO: run all once emulator is faster
|
||||
- name: Run RDNA3 ops tests
|
||||
run: SKIP_SLOW_TEST=1 AMD=1 PYTHON_REMU=1 MOCKGPU=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril"
|
||||
run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
|
||||
|
||||
testnvidia:
|
||||
strategy:
|
||||
|
||||
@@ -109,6 +109,8 @@ 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):
|
||||
|
||||
@@ -1,479 +0,0 @@
|
||||
# RDNA3 emulator - executes compiled pseudocode from AMD ISA PDF
|
||||
# mypy: ignore-errors
|
||||
from __future__ import annotations
|
||||
import ctypes, functools
|
||||
from enum import IntEnum
|
||||
from tinygrad.runtime.autogen import hsa
|
||||
from extra.assembly.amd.dsl import Inst, NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, v, s
|
||||
from extra.assembly.amd.pcode import _f32, _i32, _sext, _f16, _i16, _f64, _i64
|
||||
from extra.assembly.amd.decode import decode_inst
|
||||
from extra.assembly.amd.pcode import compile_pseudocode
|
||||
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
|
||||
from extra.assembly.amd.autogen.rdna3.ins import (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, GLOBAL, SCRATCH, VOPD,
|
||||
SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, SCRATCHOp, VOPDOp)
|
||||
|
||||
# Constants and helpers defined locally (not imported from dsl.py)
|
||||
MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
|
||||
FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
|
||||
|
||||
class SGPRArray:
|
||||
"""SGPR array indexed by Reg or int."""
|
||||
__slots__ = ('_data',)
|
||||
def __init__(self, size: int): self._data = [0] * size
|
||||
def __getitem__(self, key): return self._data[getattr(key, 'offset', key)]
|
||||
def __setitem__(self, key, val): self._data[getattr(key, 'offset', key)] = val
|
||||
def __len__(self): return len(self._data)
|
||||
def __iter__(self): return iter(self._data)
|
||||
|
||||
class VGPRLane:
|
||||
"""Single lane of VGPRs indexed by Reg (offset 256-511) or int (0-255)."""
|
||||
__slots__ = ('_data',)
|
||||
def __init__(self, size: int): self._data = [0] * size
|
||||
def __getitem__(self, key):
|
||||
i = getattr(key, 'offset', key)
|
||||
return self._data[i - 256 if i >= 256 else i]
|
||||
def __setitem__(self, key, val):
|
||||
i = getattr(key, 'offset', key)
|
||||
self._data[i - 256 if i >= 256 else i] = val
|
||||
def __len__(self): return len(self._data)
|
||||
def __iter__(self): return iter(self._data)
|
||||
|
||||
WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
|
||||
|
||||
# Inline constants for src operands 128-254. Build tables for f32, f16, and f64 formats.
|
||||
_FLOAT_CONSTS = {v: k for k, v in FLOAT_ENC.items()} | {248: 0.15915494309189535} # INV_2PI
|
||||
def _build_inline_consts(mask, to_bits):
|
||||
tbl = list(range(65)) + [((-i) & mask) for i in range(1, 17)] + [0] * (127 - 81)
|
||||
for k, v in _FLOAT_CONSTS.items(): tbl[k - 128] = to_bits(v)
|
||||
return tbl
|
||||
_INLINE_CONSTS = _build_inline_consts(MASK32, _i32)
|
||||
_INLINE_CONSTS_F16 = _build_inline_consts(0xffff, _i16)
|
||||
_INLINE_CONSTS_F64 = _build_inline_consts(MASK64, _i64)
|
||||
|
||||
# Helper: extract/write 16-bit half from/to 32-bit value
|
||||
def _src16(raw: int, is_hi: bool) -> int: return ((raw >> 16) & 0xffff) if is_hi else (raw & 0xffff)
|
||||
def _dst16(cur: int, val: int, is_hi: bool) -> int: return (cur & 0x0000ffff) | ((val & 0xffff) << 16) if is_hi else (cur & 0xffff0000) | (val & 0xffff)
|
||||
def _vgpr_hi(src) -> bool: return src.offset >= 256 and ((src.offset - 256) & 0x80) != 0
|
||||
def _vgpr_masked(src): return v[(src.offset - 256) & 0x7f] if src.offset >= 256 else src
|
||||
|
||||
# VOP3 source modifier: apply abs/neg to value
|
||||
def _mod_src(val: int, idx: int, neg: int, abs_: int, is64: bool = False) -> int:
|
||||
to_f, to_i = (_f64, _i64) if is64 else (_f32, _i32)
|
||||
if (abs_ >> idx) & 1: val = to_i(abs(to_f(val)))
|
||||
if (neg >> idx) & 1: val = to_i(-to_f(val))
|
||||
return val
|
||||
|
||||
# Read source operand with VOP3 modifiers
|
||||
def _read_src(st, inst, src, idx: int, lane: int, neg: int, abs_: int, opsel: int) -> int:
|
||||
if src is None: return 0
|
||||
src_off = src.offset
|
||||
src_bits = inst.canonical_op_bits[f's{idx}']
|
||||
literal, is_src_64, is_src_16 = inst._literal, src_bits == 64, src_bits == 16
|
||||
if is_src_64: return _mod_src(st.rsrc64(src, lane, literal), idx, neg, abs_, is64=True)
|
||||
if isinstance(inst, VOP3P):
|
||||
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
|
||||
if 'FMA_MIX' in inst.op_name:
|
||||
raw = st.rsrc(src, lane, literal)
|
||||
sign_bit = (15 if not (opsel & (1 << idx)) else 31) if (opsel_hi >> idx) & 1 else 31
|
||||
if inst.neg_hi & (1 << idx): raw &= ~(1 << sign_bit)
|
||||
if neg & (1 << idx): raw ^= (1 << sign_bit)
|
||||
return raw
|
||||
raw = st.rsrc_f16(src, lane, literal)
|
||||
hi = _src16(raw, opsel_hi & (1 << idx)) ^ (0x8000 if inst.neg_hi & (1 << idx) else 0)
|
||||
lo = _src16(raw, opsel & (1 << idx)) ^ (0x8000 if neg & (1 << idx) else 0)
|
||||
return (hi << 16) | lo
|
||||
if is_src_16 and isinstance(inst, VOP3):
|
||||
raw = st.rsrc_f16(src, lane, literal) if 128 <= src_off < 255 else st.rsrc(src, lane, literal)
|
||||
val = _src16(raw, bool(opsel & (1 << idx)))
|
||||
if abs_ & (1 << idx): val &= 0x7fff
|
||||
if neg & (1 << idx): val ^= 0x8000
|
||||
return val
|
||||
if is_src_16 and isinstance(inst, (VOP1, VOP2, VOPC)):
|
||||
if src_off >= 256: return _src16(_mod_src(st.rsrc(_vgpr_masked(src), lane, literal), idx, neg, abs_), _vgpr_hi(src))
|
||||
return _mod_src(st.rsrc_f16(src, lane, literal), idx, neg, abs_) & 0xffff
|
||||
return _mod_src(st.rsrc(src, lane, literal), idx, neg, abs_)
|
||||
|
||||
# Helper: get number of dwords from memory op name
|
||||
def _op_ndwords(name: str) -> int:
|
||||
if '_B128' in name: return 4
|
||||
if '_B96' in name: return 3
|
||||
if any(s in name for s in ('_B64', '_U64', '_I64', '_F64')): return 2
|
||||
return 1
|
||||
|
||||
# Helper: build multi-dword int from consecutive VGPRs
|
||||
def _vgpr_read(V: VGPRLane, reg, ndwords: int) -> int:
|
||||
return sum(V[reg + i] << (32 * i) for i in range(ndwords))
|
||||
|
||||
# Helper: write multi-dword value to consecutive VGPRs
|
||||
def _vgpr_write(V: VGPRLane, reg, val: int, ndwords: int):
|
||||
for i in range(ndwords): V[reg + i] = (val >> (32 * i)) & MASK32
|
||||
|
||||
# Memory access
|
||||
_valid_mem_ranges: list[tuple[int, int]] = []
|
||||
def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: _valid_mem_ranges.clear(); _valid_mem_ranges.extend(ranges)
|
||||
def _mem_valid(addr: int, size: int) -> bool:
|
||||
return not _valid_mem_ranges or any(s <= addr and addr + size <= s + z for s, z in _valid_mem_ranges)
|
||||
def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint64 if size == 8 else ctypes.c_uint32).from_address(addr)
|
||||
def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
|
||||
def mem_write(addr: int, size: int, val: int) -> None:
|
||||
if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
|
||||
|
||||
def _make_mem_accessor(read_fn, write_fn):
|
||||
"""Create a memory accessor class with the given read/write functions."""
|
||||
class _MemAccessor:
|
||||
__slots__ = ('_addr',)
|
||||
def __init__(self, addr: int): self._addr = int(addr)
|
||||
u8 = property(lambda s: read_fn(s._addr, 1), lambda s, v: write_fn(s._addr, 1, int(v)))
|
||||
u16 = property(lambda s: read_fn(s._addr, 2), lambda s, v: write_fn(s._addr, 2, int(v)))
|
||||
u32 = property(lambda s: read_fn(s._addr, 4), lambda s, v: write_fn(s._addr, 4, int(v)))
|
||||
u64 = property(lambda s: read_fn(s._addr, 8), lambda s, v: write_fn(s._addr, 8, int(v)))
|
||||
i8 = property(lambda s: _sext(read_fn(s._addr, 1), 8), lambda s, v: write_fn(s._addr, 1, int(v)))
|
||||
i16 = property(lambda s: _sext(read_fn(s._addr, 2), 16), lambda s, v: write_fn(s._addr, 2, int(v)))
|
||||
i32 = property(lambda s: _sext(read_fn(s._addr, 4), 32), lambda s, v: write_fn(s._addr, 4, int(v)))
|
||||
i64 = property(lambda s: _sext(read_fn(s._addr, 8), 64), lambda s, v: write_fn(s._addr, 8, int(v)))
|
||||
b8, b16, b32, b64 = u8, u16, u32, u64
|
||||
return _MemAccessor
|
||||
|
||||
_GlobalMemAccessor = _make_mem_accessor(mem_read, mem_write)
|
||||
|
||||
class _GlobalMem:
|
||||
"""Global memory wrapper that supports MEM[addr].u32 style access."""
|
||||
def __getitem__(self, addr) -> _GlobalMemAccessor: return _GlobalMemAccessor(addr)
|
||||
GlobalMem = _GlobalMem()
|
||||
|
||||
class LDSMem:
|
||||
"""LDS memory wrapper that supports MEM[addr].u32 style access."""
|
||||
__slots__ = ('_lds',)
|
||||
def __init__(self, lds: bytearray): self._lds = lds
|
||||
def _read(self, addr: int, size: int) -> int:
|
||||
addr = addr & 0xffff
|
||||
return int.from_bytes(self._lds[addr:addr+size], 'little') if addr + size <= len(self._lds) else 0
|
||||
def _write(self, addr: int, size: int, val: int):
|
||||
addr = addr & 0xffff
|
||||
if addr + size <= len(self._lds): self._lds[addr:addr+size] = (int(val) & ((1 << (size*8)) - 1)).to_bytes(size, 'little')
|
||||
def __getitem__(self, addr): return _make_mem_accessor(self._read, self._write)(addr)
|
||||
|
||||
# SMEM dst register count (for writing result back to SGPRs)
|
||||
SMEM_DST_COUNT = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
|
||||
|
||||
# VOPD op -> VOP3 op mapping (VOPD is dual-issue of VOP1/VOP2 ops, use VOP3 enums for pseudocode lookup)
|
||||
_VOPD_TO_VOP = {
|
||||
VOPDOp.V_DUAL_FMAC_F32: VOP3Op.V_FMAC_F32_E64, VOPDOp.V_DUAL_FMAAK_F32: VOP2Op.V_FMAAK_F32_E32, VOPDOp.V_DUAL_FMAMK_F32: VOP2Op.V_FMAMK_F32_E32,
|
||||
VOPDOp.V_DUAL_MUL_F32: VOP3Op.V_MUL_F32_E64, VOPDOp.V_DUAL_ADD_F32: VOP3Op.V_ADD_F32_E64, VOPDOp.V_DUAL_SUB_F32: VOP3Op.V_SUB_F32_E64,
|
||||
VOPDOp.V_DUAL_SUBREV_F32: VOP3Op.V_SUBREV_F32_E64, VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: VOP3Op.V_MUL_DX9_ZERO_F32_E64,
|
||||
VOPDOp.V_DUAL_MOV_B32: VOP3Op.V_MOV_B32_E64, VOPDOp.V_DUAL_CNDMASK_B32: VOP3Op.V_CNDMASK_B32_E64,
|
||||
VOPDOp.V_DUAL_MAX_F32: VOP3Op.V_MAX_F32_E64, VOPDOp.V_DUAL_MIN_F32: VOP3Op.V_MIN_F32_E64,
|
||||
VOPDOp.V_DUAL_ADD_NC_U32: VOP3Op.V_ADD_NC_U32_E64, VOPDOp.V_DUAL_LSHLREV_B32: VOP3Op.V_LSHLREV_B32_E64, VOPDOp.V_DUAL_AND_B32: VOP3Op.V_AND_B32_E64,
|
||||
}
|
||||
|
||||
|
||||
class WaveState:
|
||||
__slots__ = ('sgpr', 'vgpr', 'scc', 'pc', '_pend_sgpr', 'lds', 'n_lanes')
|
||||
def __init__(self, lds: LDSMem | None = None, n_lanes: int = WAVE_SIZE):
|
||||
self.sgpr, self.vgpr = SGPRArray(SGPR_COUNT), [VGPRLane(VGPR_COUNT) for _ in range(WAVE_SIZE)]
|
||||
self.sgpr[EXEC_LO], self.scc, self.pc, self._pend_sgpr, self.lds, self.n_lanes = 0xffffffff, 0, 0, {}, lds, n_lanes
|
||||
|
||||
@property
|
||||
def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
|
||||
@vcc.setter
|
||||
def vcc(self, v: int): self.sgpr[VCC_LO], self.sgpr[VCC_HI] = v & MASK32, (v >> 32) & MASK32
|
||||
@property
|
||||
def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
|
||||
@exec_mask.setter
|
||||
def exec_mask(self, v: int): self.sgpr[EXEC_LO], self.sgpr[EXEC_HI] = v & MASK32, (v >> 32) & MASK32
|
||||
|
||||
def rsgpr(self, reg) -> int:
|
||||
if reg == NULL: return 0
|
||||
if reg == SCC: return self.scc
|
||||
return self.sgpr[reg]
|
||||
def wsgpr(self, reg, v: int):
|
||||
if reg != NULL: self.sgpr[reg] = v & MASK32
|
||||
def rsgpr64(self, reg) -> int:
|
||||
off = reg.offset
|
||||
return self.sgpr._data[off] | (self.sgpr._data[off + 1] << 32)
|
||||
def wsgpr64(self, reg, v: int):
|
||||
off = reg.offset
|
||||
self.sgpr._data[off] = v & MASK32; self.sgpr._data[off + 1] = (v >> 32) & MASK32
|
||||
|
||||
def _rsrc_base(self, reg, lane: int, consts, literal: int):
|
||||
off = reg.offset
|
||||
if off < SGPR_COUNT: return self.sgpr._data[off]
|
||||
if off == SCC.offset: return self.scc
|
||||
if off < 255: return consts[off - 128]
|
||||
if off == 255: return literal
|
||||
return self.vgpr[lane]._data[off - 256] if off <= 511 else 0
|
||||
def rsrc(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS, literal)
|
||||
def rsrc_f16(self, reg, lane: int, literal: int = 0) -> int: return self._rsrc_base(reg, lane, _INLINE_CONSTS_F16, literal)
|
||||
def rsrc64(self, reg, lane: int, literal: int = 0) -> int:
|
||||
off = reg.offset
|
||||
if 128 <= off < 255: return _INLINE_CONSTS_F64[off - 128]
|
||||
if off == 255: return literal << 32 # 32-bit literal forms upper 32 bits of 64-bit value
|
||||
return self.rsrc(reg, lane, literal) | ((self.rsrc(reg + 1, lane, literal) if off < VCC_LO.offset or 256 <= off <= 511 else 0) << 32)
|
||||
|
||||
def pend_sgpr_lane(self, reg, lane: int, val: int):
|
||||
if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
|
||||
if val: self._pend_sgpr[reg] |= (1 << lane)
|
||||
def commit_pends(self):
|
||||
for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
|
||||
self._pend_sgpr.clear()
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# EXECUTION - All ops use pseudocode from PDF
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_scalar(st: WaveState, inst: Inst):
|
||||
"""Execute scalar instruction. Returns 0 to continue execution."""
|
||||
# Get op enum and lookup compiled function
|
||||
if isinstance(inst, SMEM): ssrc0, sdst = None, None
|
||||
elif isinstance(inst, SOP1): ssrc0, sdst = inst.ssrc0, inst.sdst
|
||||
elif isinstance(inst, SOP2): ssrc0, sdst = inst.ssrc0, inst.sdst
|
||||
elif isinstance(inst, SOPC): ssrc0, sdst = inst.ssrc0, None
|
||||
elif isinstance(inst, SOPK): ssrc0, sdst = inst.sdst, inst.sdst # sdst is both src and dst
|
||||
elif isinstance(inst, SOPP): ssrc0, sdst = None, None
|
||||
else: raise NotImplementedError(f"Unknown scalar type {type(inst)}")
|
||||
|
||||
# SMEM: memory loads
|
||||
if isinstance(inst, SMEM):
|
||||
addr = st.rsgpr64(inst.sbase) + _sext(inst.offset, 21)
|
||||
if inst.soffset != NULL: addr += st.rsrc(inst.soffset, 0, inst._literal)
|
||||
result = inst._fn(GlobalMem, addr & MASK64)
|
||||
if 'SDATA' in result:
|
||||
sdata = result['SDATA']
|
||||
for i in range(SMEM_DST_COUNT.get(inst.op, 1)): st.wsgpr(inst.sdata + i, (sdata >> (i * 32)) & MASK32)
|
||||
st.pc += inst._words
|
||||
return 0
|
||||
|
||||
# Build context - use canonical_op_bits to determine operand sizes
|
||||
literal = inst._literal
|
||||
s0 = st.rsrc64(ssrc0, 0, literal) if inst.canonical_op_bits['s0'] == 64 else (st.rsrc(ssrc0, 0, literal) if not isinstance(inst, (SOPK, SOPP)) else (st.rsgpr(inst.sdst) if isinstance(inst, SOPK) else 0))
|
||||
s1 = st.rsrc64(inst.ssrc1, 0, literal) if inst.canonical_op_bits['s1'] == 64 else (st.rsrc(inst.ssrc1, 0, literal) if isinstance(inst, (SOP2, SOPC)) else inst.simm16 if isinstance(inst, SOPK) else 0)
|
||||
d0 = st.rsgpr64(sdst) if inst.canonical_op_bits['d'] == 64 and sdst is not None else (st.rsgpr(sdst) if sdst is not None else 0)
|
||||
literal = inst.simm16 if isinstance(inst, (SOPK, SOPP)) else inst._literal
|
||||
|
||||
# Call compiled function with int parameters
|
||||
result = inst._fn(s0, s1, 0, d0, st.scc, st.vcc & MASK32, 0, st.exec_mask & MASK32, literal, None, pc=st.pc * 4)
|
||||
|
||||
# Apply results (already int values)
|
||||
if sdst is not None and 'D0' in result:
|
||||
(st.wsgpr64 if inst.canonical_op_bits['d'] == 64 else st.wsgpr)(sdst, result['D0'])
|
||||
if 'SCC' in result: st.scc = result['SCC'] & 1
|
||||
if 'EXEC' in result: st.exec_mask = result['EXEC']
|
||||
if 'PC' in result:
|
||||
# Convert absolute byte address to word offset
|
||||
pc_val = result['PC']
|
||||
new_pc = pc_val if pc_val < 0x8000000000000000 else pc_val - 0x10000000000000000
|
||||
st.pc = new_pc // 4
|
||||
else:
|
||||
st.pc += inst._words
|
||||
return 0
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# VECTOR INSTRUCTIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_vopd(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
|
||||
"""VOPD: dual-issue, execute two ops simultaneously (read all inputs before writes)."""
|
||||
literal = inst._literal
|
||||
sx0, sx1, dx = st.rsrc(inst.srcx0, lane, literal), V[inst.vsrcx1], V[inst.vdstx]
|
||||
sy0, sy1, dy = st.rsrc(inst.srcy0, lane, literal), V[inst.vsrcy1], V[inst.vdsty]
|
||||
V[inst.vdstx] = inst._fnx(sx0, sx1, 0, dx, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
|
||||
V[inst.vdsty] = inst._fny(sy0, sy1, 0, dy, st.scc, st.vcc, lane, st.exec_mask, literal, None)['D0']
|
||||
|
||||
def exec_flat(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
|
||||
"""FLAT/GLOBAL/SCRATCH memory ops."""
|
||||
ndwords = _op_ndwords(inst.op_name)
|
||||
addr = V[inst.addr] | (V[inst.addr + 1] << 32)
|
||||
ADDR = (st.rsgpr64(inst.saddr) + V[inst.addr] + _sext(inst.offset, 13)) & MASK64 if inst.saddr != NULL else (addr + _sext(inst.offset, 13)) & MASK64
|
||||
vdata_src = inst.vdst if 'LOAD' in inst.op_name else inst.data
|
||||
result = inst._fn(GlobalMem, ADDR, _vgpr_read(V, vdata_src, ndwords), V[inst.vdst])
|
||||
if 'VDATA' in result: _vgpr_write(V, inst.vdst, result['VDATA'], ndwords)
|
||||
if 'RETURN_DATA' in result: _vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords)
|
||||
|
||||
def exec_ds(st: WaveState, inst, V: VGPRLane, lane: int) -> None:
|
||||
"""DS (LDS) memory ops."""
|
||||
ndwords = _op_ndwords(inst.op_name)
|
||||
data0, data1 = _vgpr_read(V, inst.data0, ndwords), _vgpr_read(V, inst.data1, ndwords) if inst.data1 is not None else 0
|
||||
result = inst._fn(st.lds, V[inst.addr], data0, data1, inst.offset0, inst.offset1)
|
||||
if 'RETURN_DATA' in result and ('_RTN' in inst.op_name or '_LOAD' in inst.op_name):
|
||||
_vgpr_write(V, inst.vdst, result['RETURN_DATA'], ndwords * 2 if '_2ADDR_' in inst.op_name else ndwords)
|
||||
|
||||
def exec_vop(st: WaveState, inst: Inst, V: VGPRLane, lane: int) -> None:
|
||||
"""VOP1/VOP2/VOP3/VOP3SD/VOP3P/VOPC: standard ALU ops."""
|
||||
is_dst_16 = inst.canonical_op_bits['d'] == 16
|
||||
if isinstance(inst, VOP3P):
|
||||
src0, src1, src2, vdst, dst_hi = inst.src0, inst.src1, inst.src2, inst.vdst, False
|
||||
neg, abs_, opsel = inst.neg, 0, inst.opsel
|
||||
elif isinstance(inst, VOP1):
|
||||
src0, src1, src2, vdst = inst.src0, None, None, inst.vdst
|
||||
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
|
||||
if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
|
||||
elif isinstance(inst, VOP2):
|
||||
src0, src1, src2, vdst = inst.src0, inst.vsrc1, None, inst.vdst
|
||||
neg, abs_, opsel, dst_hi = 0, 0, 0, (inst.vdst.offset & 0x80) != 0 and is_dst_16
|
||||
if is_dst_16: vdst = v[inst.vdst.offset & 0x7f]
|
||||
elif isinstance(inst, (VOP3, VOP3SD)):
|
||||
src0, src1, src2, vdst = inst.src0, inst.src1, (None if isinstance(inst, VOP3) and inst.op.value < 256 else inst.src2), inst.vdst
|
||||
neg, abs_, opsel, dst_hi = (inst.neg, inst.abs, inst.opsel, False) if isinstance(inst, VOP3) else (0, 0, 0, False)
|
||||
elif isinstance(inst, VOPC):
|
||||
src0, src1, src2, vdst, neg, abs_, opsel, dst_hi = inst.src0, inst.vsrc1, None, VCC_LO, 0, 0, 0, False
|
||||
else:
|
||||
raise NotImplementedError(f"exec_vop: unhandled instruction type {type(inst).__name__}")
|
||||
|
||||
s0 = _read_src(st, inst, src0, 0, lane, neg, abs_, opsel)
|
||||
s1 = _read_src(st, inst, src1, 1, lane, neg, abs_, opsel)
|
||||
s2 = _read_src(st, inst, src2, 2, lane, neg, abs_, opsel)
|
||||
if isinstance(inst, VOP2) and is_dst_16: d0 = _src16(V[vdst], dst_hi)
|
||||
elif inst.canonical_op_bits['d'] == 64: d0 = V[vdst] | (V[vdst + 1] << 32)
|
||||
else: d0 = V[vdst]
|
||||
|
||||
if isinstance(inst, VOP3SD) and 'CO_CI' in inst.op_name: vcc_for_fn = st.rsgpr64(inst.src2)
|
||||
elif isinstance(inst, VOP3) and inst.op in (VOP3Op.V_CNDMASK_B32_E64, VOP3Op.V_CNDMASK_B16) and src2 is not None and src2.offset < 256: vcc_for_fn = st.rsgpr64(src2)
|
||||
else: vcc_for_fn = st.vcc
|
||||
src0_off = src0.offset if src0 is not None else 0
|
||||
src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
|
||||
vdst_off = vdst.offset
|
||||
extra_kwargs = {'opsel': opsel, 'opsel_hi': inst.opsel_hi | (inst.opsel_hi2 << 2)} if isinstance(inst, VOP3P) and 'FMA_MIX' in inst.op_name else {}
|
||||
result = inst._fn(s0, s1, s2, d0, st.scc, vcc_for_fn, lane, st.exec_mask, inst._literal, st.vgpr, src0_idx, vdst_off, **extra_kwargs)
|
||||
|
||||
# Check if this is a VOPC instruction (either standalone VOPC or VOP3 with VOPC opcode)
|
||||
is_vopc = isinstance(inst.op, VOPCOp) or (isinstance(inst, VOP3) and inst.op.value < 256)
|
||||
if 'VCC' in result:
|
||||
if isinstance(inst, VOP3SD): st.pend_sgpr_lane(inst.sdst, lane, (result['VCC'] >> lane) & 1)
|
||||
elif isinstance(inst, VOP2) and 'CO_CI' in inst.op_name: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
|
||||
elif is_vopc: st.pend_sgpr_lane(vdst, lane, (result['VCC'] >> lane) & 1) # vdst is VCC_LO for VOPC
|
||||
else: st.pend_sgpr_lane(VCC_LO, lane, (result['VCC'] >> lane) & 1)
|
||||
if 'EXEC' in result:
|
||||
st.pend_sgpr_lane(EXEC_LO, lane, (result['EXEC'] >> lane) & 1)
|
||||
elif is_vopc:
|
||||
st.pend_sgpr_lane(vdst, lane, (result['D0'] >> lane) & 1)
|
||||
if not is_vopc:
|
||||
d0_val = result['D0']
|
||||
if inst.canonical_op_bits['d'] == 64: V[vdst], V[vdst + 1] = d0_val & MASK32, (d0_val >> 32) & MASK32
|
||||
elif not isinstance(inst, VOP3P) and is_dst_16: V[vdst] = _dst16(V[vdst], d0_val, bool(opsel & 8) if isinstance(inst, VOP3) else dst_hi)
|
||||
else: V[vdst] = d0_val & MASK32
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# WMMA (Wave Matrix Multiply-Accumulate)
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_wmma(st: WaveState, inst, op: VOP3POp) -> None:
|
||||
"""Execute WMMA instruction - 16x16x16 matrix multiply across the wave."""
|
||||
src0, src1, src2, vdst = inst.src0.offset, inst.src1.offset, inst.src2.offset, inst.vdst.offset
|
||||
# Read 16x16 f16 matrix from 16 lanes × 8 VGPRs (2 f16 per VGPR)
|
||||
def read_f16_mat(src):
|
||||
return [f for l in range(16) for r in range(8) for v in [st.vgpr[l][src-256+r] if src >= 256 else st.rsgpr(src+r)] for f in [_f16(v&0xffff), _f16((v>>16)&0xffff)]]
|
||||
mat_a, mat_b = read_f16_mat(src0), read_f16_mat(src1)
|
||||
# Read matrix C (16x16 f32) from lanes 0-31, VGPRs src2 to src2+7
|
||||
mat_c = [_f32(st.vgpr[i % 32][src2 - 256 + i // 32] if src2 >= 256 else st.rsgpr(src2 + i // 32)) for i in range(256)]
|
||||
# Compute D = A × B + C (16x16 matrix multiply)
|
||||
mat_d = [sum(mat_a[row*16+k] * mat_b[col*16+k] for k in range(16)) + mat_c[row*16+col] for row in range(16) for col in range(16)]
|
||||
# Write result - f16 packed or f32
|
||||
if op == VOP3POp.V_WMMA_F16_16X16X16_F16:
|
||||
for i in range(0, 256, 2):
|
||||
st.vgpr[(i//2) % 32][vdst - 256 + (i//2)//32] = ((_i16(mat_d[i+1]) & 0xffff) << 16) | (_i16(mat_d[i]) & 0xffff)
|
||||
else:
|
||||
for i in range(256): st.vgpr[i % 32][vdst - 256 + i//32] = _i32(mat_d[i])
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PROGRAM DECODE
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
# Wave-level dispatch functions: (st, inst) -> return_code (0 = continue, -1 = end, -2 = barrier)
|
||||
def dispatch_endpgm(st, inst): return -1
|
||||
def dispatch_barrier(st, inst): st.pc += inst._words; return -2
|
||||
def dispatch_nop(st, inst): st.pc += inst._words; return 0
|
||||
def dispatch_wmma(st, inst): exec_wmma(st, inst, inst.op); st.pc += inst._words; return 0
|
||||
def dispatch_writelane(st, inst): st.vgpr[st.rsrc(inst.src1, 0, inst._literal) & 0x1f][inst.vdst.offset - 256] = st.rsrc(inst.src0, 0, inst._literal) & MASK32; st.pc += inst._words; return 0
|
||||
def dispatch_readlane(st, inst):
|
||||
src0_off = inst.src0.offset
|
||||
src0_idx = (src0_off - 256) if src0_off >= 256 else src0_off
|
||||
s1 = st.rsrc(inst.src1, 0, inst._literal) if getattr(inst, 'src1', None) is not None else 0
|
||||
result = inst._fn(0, s1, 0, 0, st.scc, st.vcc, 0, st.exec_mask, inst._literal, st.vgpr, src0_idx, inst.vdst.offset)
|
||||
st.wsgpr(inst.vdst.offset, result['D0'])
|
||||
st.pc += inst._words; return 0
|
||||
|
||||
# Per-lane dispatch wrapper: wraps per-lane exec functions into wave-level dispatch
|
||||
@functools.cache
|
||||
def dispatch_lane(exec_fn):
|
||||
def dispatch(st, inst):
|
||||
exec_mask, vgpr, n_lanes = st.exec_mask, st.vgpr, st.n_lanes
|
||||
for lane in range(n_lanes):
|
||||
if exec_mask >> lane & 1: exec_fn(st, inst, vgpr[lane], lane)
|
||||
st.commit_pends()
|
||||
st.pc += inst._words
|
||||
return 0
|
||||
return dispatch
|
||||
|
||||
def decode_program(data: bytes) -> dict[int, Inst]:
|
||||
result: dict[int, Inst] = {}
|
||||
i = 0
|
||||
while i < len(data):
|
||||
inst = decode_inst(data[i:])
|
||||
inst._words = inst.size() // 4
|
||||
|
||||
# Determine dispatch function and pcode function
|
||||
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
|
||||
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: inst._dispatch = dispatch_endpgm
|
||||
elif isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER: inst._dispatch = dispatch_barrier
|
||||
elif isinstance(inst, SOPP) and inst.op in (SOPPOp.S_CLAUSE, SOPPOp.S_WAITCNT, SOPPOp.S_WAITCNT_DEPCTR, SOPPOp.S_SENDMSG, SOPPOp.S_SET_INST_PREFETCH_DISTANCE, SOPPOp.S_DELAY_ALU): inst._dispatch = dispatch_nop
|
||||
elif isinstance(inst, (SOP1, SOP2, SOPC, SOPK, SOPP, SMEM)): inst._dispatch = exec_scalar
|
||||
elif isinstance(inst, VOP1) and inst.op == VOP1Op.V_NOP_E32: inst._dispatch = dispatch_nop
|
||||
elif isinstance(inst, VOP3P) and 'WMMA' in inst.op_name: inst._dispatch = dispatch_wmma
|
||||
elif isinstance(inst, VOP3) and inst.op == VOP3Op.V_WRITELANE_B32: inst._dispatch = dispatch_writelane
|
||||
elif isinstance(inst, (VOP1, VOP3)) and inst.op in (VOP1Op.V_READFIRSTLANE_B32_E32, VOP3Op.V_READFIRSTLANE_B32, VOP3Op.V_READLANE_B32): inst._dispatch = dispatch_readlane
|
||||
elif isinstance(inst, VOPD): inst._dispatch = dispatch_lane(exec_vopd)
|
||||
elif isinstance(inst, (FLAT, GLOBAL, SCRATCH)): inst._dispatch = dispatch_lane(exec_flat)
|
||||
elif isinstance(inst, DS): inst._dispatch = dispatch_lane(exec_ds)
|
||||
else: inst._dispatch = dispatch_lane(exec_vop)
|
||||
|
||||
# Compile pcode for instructions that use it (not VOPD which has _fnx/_fny, not special dispatches)
|
||||
# VOPD needs separate functions for X and Y ops
|
||||
if isinstance(inst, VOPD):
|
||||
def _compile_vopd_op(op): return compile_pseudocode(type(op).__name__, op.name, PCODE[op])
|
||||
inst._fnx, inst._fny = _compile_vopd_op(_VOPD_TO_VOP[inst.opx]), _compile_vopd_op(_VOPD_TO_VOP[inst.opy])
|
||||
elif inst._dispatch not in (dispatch_endpgm, dispatch_barrier, dispatch_nop, dispatch_wmma, dispatch_writelane):
|
||||
assert type(inst.op) != int, f"inst op of {inst} is int"
|
||||
inst._fn = compile_pseudocode(type(inst.op).__name__, inst.op.name, PCODE[inst.op])
|
||||
result[i // 4] = inst
|
||||
i += inst._words * 4
|
||||
return result
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# MAIN EXECUTION LOOP
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def exec_wave(program: dict[int, Inst], st: WaveState) -> int:
|
||||
while (inst := program.get(st.pc)) and (result := inst._dispatch(st, inst)) == 0: pass
|
||||
return result
|
||||
|
||||
def exec_workgroup(program: dict[int, Inst], workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, rsrc2: int) -> None:
|
||||
lx, ly, lz = local_size
|
||||
total_threads = lx * ly * lz
|
||||
# GRANULATED_LDS_SIZE is in 512-byte units (see ops_amd.py: lds_size = ((group_segment_size + 511) // 512))
|
||||
lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
|
||||
lds = LDSMem(bytearray(lds_size)) if lds_size else None
|
||||
waves: list[WaveState] = []
|
||||
for wave_start in range(0, total_threads, WAVE_SIZE):
|
||||
n_lanes = min(WAVE_SIZE, total_threads - wave_start)
|
||||
st = WaveState(lds, n_lanes)
|
||||
st.exec_mask = (1 << n_lanes) - 1
|
||||
st.wsgpr64(s[0:1], args_ptr) # s[0:1] = kernel arguments pointer
|
||||
# COMPUTE_PGM_RSRC2: USER_SGPR_COUNT is where workgroup IDs start, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z control which are passed
|
||||
sgpr_idx = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
|
||||
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X: st.sgpr[sgpr_idx] = workgroup_id[0]; sgpr_idx += 1
|
||||
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y: st.sgpr[sgpr_idx] = workgroup_id[1]; sgpr_idx += 1
|
||||
if rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z: st.sgpr[sgpr_idx] = workgroup_id[2]
|
||||
# VGPR0 = packed workitem IDs: (Z << 20) | (Y << 10) | X
|
||||
for tid in range(wave_start, wave_start + n_lanes):
|
||||
st.vgpr[tid - wave_start][0] = ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
|
||||
waves.append(st)
|
||||
while waves:
|
||||
waves = [st for st in waves if exec_wave(program, st) != -1]
|
||||
|
||||
def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c) -> int:
|
||||
program = decode_program((ctypes.c_char * lib_sz).from_address(lib).raw)
|
||||
for gidz in range(gz):
|
||||
for gidy in range(gy):
|
||||
for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, rsrc2)
|
||||
return 0
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,822 +0,0 @@
|
||||
# DSL for RDNA3 pseudocode - makes pseudocode expressions work directly as Python
|
||||
import struct, math, re, functools
|
||||
|
||||
MASK32, MASK64 = 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF
|
||||
|
||||
# Float/int bit conversion functions
|
||||
_struct_f, _struct_I = struct.Struct("<f"), struct.Struct("<I")
|
||||
_struct_e, _struct_H = struct.Struct("<e"), struct.Struct("<H")
|
||||
_struct_d, _struct_Q = struct.Struct("<d"), struct.Struct("<Q")
|
||||
def _f32(i):
|
||||
i = i & MASK32
|
||||
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
|
||||
# Denormal: exponent=0 (bits 23-30) and mantissa!=0 (bits 0-22)
|
||||
if (i & 0x7f800000) == 0 and (i & 0x007fffff) != 0: return 0.0
|
||||
return _struct_f.unpack(_struct_I.pack(i))[0]
|
||||
def _i32(f):
|
||||
if isinstance(f, int): f = float(f)
|
||||
if math.isnan(f): return 0xffc00000 if math.copysign(1.0, f) < 0 else 0x7fc00000
|
||||
if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
|
||||
try:
|
||||
bits = _struct_I.unpack(_struct_f.pack(f))[0]
|
||||
# RDNA3 default mode: flush f32 denormals to zero (FTZ)
|
||||
if (bits & 0x7f800000) == 0 and (bits & 0x007fffff) != 0: return 0x80000000 if bits & 0x80000000 else 0
|
||||
return bits
|
||||
except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
|
||||
def _sext(v, b): return v - (1 << b) if v & (1 << (b - 1)) else v
|
||||
def _f16(i): return _struct_e.unpack(_struct_H.pack(i & 0xffff))[0]
|
||||
def _i16(f):
|
||||
if math.isnan(f): return 0x7e00
|
||||
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
|
||||
try: return _struct_H.unpack(_struct_e.pack(f))[0]
|
||||
except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
|
||||
def _f64(i): return _struct_d.unpack(_struct_Q.pack(i & MASK64))[0]
|
||||
def _i64(f):
|
||||
if math.isnan(f): return 0x7ff8000000000000
|
||||
if math.isinf(f): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
|
||||
try: return _struct_Q.unpack(_struct_d.pack(f))[0]
|
||||
except (OverflowError, struct.error): return 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# INTERNAL HELPERS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _div(a, b):
|
||||
try: return a / b
|
||||
except ZeroDivisionError:
|
||||
if a == 0.0 or math.isnan(a): return float("nan")
|
||||
return math.copysign(float("inf"), a * b) if b == 0.0 else float("inf") if a > 0 else float("-inf")
|
||||
def _check_nan_type(x, quiet_bit_expected, default):
|
||||
try:
|
||||
if not math.isnan(float(x)): return False
|
||||
if hasattr(x, '_reg') and hasattr(x, '_bits'):
|
||||
bits = x._reg._val & ((1 << x._bits) - 1)
|
||||
exp_bits, quiet_pos, mant_mask = {16: (0x1f, 9, 0x3ff), 32: (0xff, 22, 0x7fffff), 64: (0x7ff, 51, 0xfffffffffffff)}.get(x._bits, (0,0,0))
|
||||
exp_shift = {16: 10, 32: 23, 64: 52}.get(x._bits, 0)
|
||||
if exp_bits and ((bits >> exp_shift) & exp_bits) == exp_bits and (bits & mant_mask) != 0:
|
||||
return ((bits >> quiet_pos) & 1) == quiet_bit_expected
|
||||
return default
|
||||
except (TypeError, ValueError): return False
|
||||
def _gt_neg_zero(a, b): return (a > b) or (a == 0 and b == 0 and not math.copysign(1, a) < 0 and math.copysign(1, b) < 0)
|
||||
def _lt_neg_zero(a, b): return (a < b) or (a == 0 and b == 0 and math.copysign(1, a) < 0 and not math.copysign(1, b) < 0)
|
||||
def _fpop(fn):
|
||||
def wrapper(x):
|
||||
x = float(x)
|
||||
if math.isnan(x) or math.isinf(x): return x
|
||||
result = float(fn(x))
|
||||
return math.copysign(0.0, x) if result == 0.0 else result
|
||||
return wrapper
|
||||
def _f_to_int(f, lo, hi): f = float(f); return 0 if math.isnan(f) else (hi if f >= hi else lo if f <= lo else int(f))
|
||||
def _f16_to_f32_bits(bits): return struct.unpack("<e", struct.pack("<H", int(bits) & 0xffff))[0]
|
||||
def _brev(v, bits): return int(bin(v & ((1 << bits) - 1))[2:].zfill(bits)[::-1], 2)
|
||||
def _ctz(v, bits):
|
||||
v, n = int(v) & ((1 << bits) - 1), 0
|
||||
if v == 0: return bits
|
||||
while (v & 1) == 0: v >>= 1; n += 1
|
||||
return n
|
||||
|
||||
def _bf16(i):
|
||||
"""Convert bf16 bits to float. BF16 is just the top 16 bits of f32."""
|
||||
return struct.unpack("<f", struct.pack("<I", (i & 0xffff) << 16))[0]
|
||||
def _ibf16(f):
|
||||
"""Convert float to bf16 bits (truncate to top 16 bits of f32)."""
|
||||
if math.isnan(f): return 0x7fc0 # bf16 quiet NaN
|
||||
if math.isinf(f): return 0x7f80 if f > 0 else 0xff80 # bf16 ±infinity
|
||||
try: return (struct.unpack("<I", struct.pack("<f", float(f)))[0] >> 16) & 0xffff
|
||||
except (OverflowError, struct.error): return 0x7f80 if f > 0 else 0xff80
|
||||
def _trig(fn, x):
|
||||
# V_SIN/COS_F32: hardware does frac on input cycles before computing
|
||||
if math.isinf(x) or math.isnan(x): return float("nan")
|
||||
frac_cycles = fract(x / (2 * math.pi))
|
||||
result = fn(frac_cycles * 2 * math.pi)
|
||||
# Hardware returns exactly 0 for cos(π/2), sin(π), etc. due to lookup table
|
||||
# Round very small results (below f32 precision) to exactly 0
|
||||
if abs(result) < 1e-7: return 0.0
|
||||
return result
|
||||
|
||||
class _SafeFloat(float):
|
||||
"""Float subclass that uses _div for division to handle 0/inf correctly."""
|
||||
def __truediv__(self, o): return _div(float(self), float(o))
|
||||
def __rtruediv__(self, o): return _div(float(o), float(self))
|
||||
|
||||
class _Inf:
|
||||
f16 = f32 = f64 = float('inf')
|
||||
def __neg__(self): return _NegInf()
|
||||
def __pos__(self): return self
|
||||
def __float__(self): return float('inf')
|
||||
def __eq__(self, other): return float(other) == float('inf') if not isinstance(other, _NegInf) else False
|
||||
def __req__(self, other): return self.__eq__(other)
|
||||
class _NegInf:
|
||||
f16 = f32 = f64 = float('-inf')
|
||||
def __neg__(self): return _Inf()
|
||||
def __pos__(self): return self
|
||||
def __float__(self): return float('-inf')
|
||||
def __eq__(self, other): return float(other) == float('-inf') if not isinstance(other, _Inf) else False
|
||||
def __req__(self, other): return self.__eq__(other)
|
||||
|
||||
class _RoundMode:
|
||||
NEAREST_EVEN = 0
|
||||
|
||||
class _WaveMode:
|
||||
IEEE = False
|
||||
|
||||
class _DenormChecker:
|
||||
"""Comparator for denormalized floats. x == DENORM.f32 checks if x is denormalized."""
|
||||
def __init__(self, bits): self._bits = bits
|
||||
def _check(self, other):
|
||||
f = float(other)
|
||||
if math.isinf(f) or math.isnan(f) or f == 0.0: return False
|
||||
if self._bits == 64:
|
||||
bits = struct.unpack("<Q", struct.pack("<d", f))[0]
|
||||
return (bits >> 52) & 0x7ff == 0
|
||||
bits = struct.unpack("<I", struct.pack("<f", f))[0]
|
||||
return (bits >> 23) & 0xff == 0
|
||||
def __eq__(self, other): return self._check(other)
|
||||
def __req__(self, other): return self._check(other)
|
||||
def __ne__(self, other): return not self._check(other)
|
||||
|
||||
class _Denorm:
|
||||
f32 = _DenormChecker(32)
|
||||
f64 = _DenormChecker(64)
|
||||
|
||||
_pack = lambda hi, lo: ((int(hi) & 0xffff) << 16) | (int(lo) & 0xffff)
|
||||
_pack32 = lambda hi, lo: ((int(hi) & 0xffffffff) << 32) | (int(lo) & 0xffffffff)
|
||||
|
||||
class TypedView:
|
||||
"""View into a Reg with typed access. Used for both full-width (Reg.u32) and slices (Reg[31:16])."""
|
||||
__slots__ = ('_reg', '_high', '_low', '_signed', '_float', '_bf16', '_reversed')
|
||||
def __init__(self, reg, high, low=0, signed=False, is_float=False, is_bf16=False):
|
||||
# Handle reversed slices like [0:31] which means bit-reverse
|
||||
if high < low: high, low, reversed = low, high, True
|
||||
else: reversed = False
|
||||
self._reg, self._high, self._low, self._reversed = reg, high, low, reversed
|
||||
self._signed, self._float, self._bf16 = signed, is_float, is_bf16
|
||||
|
||||
def _nbits(self): return self._high - self._low + 1
|
||||
def _mask(self): return (1 << self._nbits()) - 1
|
||||
def _get(self):
|
||||
v = (self._reg._val >> self._low) & self._mask()
|
||||
return _brev(v, self._nbits()) if self._reversed else v
|
||||
def _set(self, v):
|
||||
v = int(v)
|
||||
if self._reversed: v = _brev(v, self._nbits())
|
||||
self._reg._val = (self._reg._val & ~(self._mask() << self._low)) | ((v & self._mask()) << self._low)
|
||||
|
||||
@property
|
||||
def _val(self): return self._get()
|
||||
@property
|
||||
def _bits(self): return self._nbits()
|
||||
|
||||
# Type accessors for slices (e.g., D0[31:16].f16)
|
||||
u8 = property(lambda s: s._get() & 0xff)
|
||||
u16 = property(lambda s: s._get() & 0xffff, lambda s, v: s._set(v))
|
||||
u32 = property(lambda s: s._get() & MASK32, lambda s, v: s._set(v))
|
||||
i16 = property(lambda s: _sext(s._get() & 0xffff, 16), lambda s, v: s._set(v))
|
||||
i32 = property(lambda s: _sext(s._get() & MASK32, 32), lambda s, v: s._set(v))
|
||||
f16 = property(lambda s: _f16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _i16(float(v))))
|
||||
f32 = property(lambda s: _f32(s._get()), lambda s, v: s._set(_i32(float(v))))
|
||||
bf16 = property(lambda s: _bf16(s._get()), lambda s, v: s._set(v if isinstance(v, int) else _ibf16(float(v))))
|
||||
b16, b32 = u16, u32
|
||||
|
||||
# Chained type access (e.g., jump_addr.i64 when jump_addr is already TypedView)
|
||||
@property
|
||||
def i64(s): return s if s._nbits() == 64 and s._signed else int(s)
|
||||
@property
|
||||
def u64(s): return s if s._nbits() == 64 and not s._signed else int(s) & MASK64
|
||||
|
||||
def __getitem__(self, key):
|
||||
if isinstance(key, slice):
|
||||
high, low = int(key.start), int(key.stop)
|
||||
return TypedView(self._reg, high, low)
|
||||
return (self._get() >> int(key)) & 1
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if isinstance(key, slice):
|
||||
high, low = int(key.start), int(key.stop)
|
||||
if high < low: high, low, value = low, high, _brev(int(value), low - high + 1)
|
||||
mask = (1 << (high - low + 1)) - 1
|
||||
self._reg._val = (self._reg._val & ~(mask << low)) | ((int(value) & mask) << low)
|
||||
elif value: self._reg._val |= (1 << int(key))
|
||||
else: self._reg._val &= ~(1 << int(key))
|
||||
|
||||
def __int__(self): return _sext(self._get(), self._nbits()) if self._signed else self._get()
|
||||
def __index__(self): return int(self)
|
||||
def __trunc__(self): return int(float(self)) if self._float else int(self)
|
||||
def __float__(self):
|
||||
if self._float:
|
||||
if self._bf16: return _bf16(self._get())
|
||||
bits = self._nbits()
|
||||
return _f16(self._get()) if bits == 16 else _f32(self._get()) if bits == 32 else _f64(self._get())
|
||||
return float(int(self))
|
||||
def __bool__(s): return bool(int(s))
|
||||
|
||||
# Arithmetic - floats use float(), ints use int()
|
||||
def __add__(s, o): return float(s) + float(o) if s._float else int(s) + int(o)
|
||||
def __radd__(s, o): return float(o) + float(s) if s._float else int(o) + int(s)
|
||||
def __sub__(s, o): return float(s) - float(o) if s._float else int(s) - int(o)
|
||||
def __rsub__(s, o): return float(o) - float(s) if s._float else int(o) - int(s)
|
||||
def __mul__(s, o): return float(s) * float(o) if s._float else int(s) * int(o)
|
||||
def __rmul__(s, o): return float(o) * float(s) if s._float else int(o) * int(s)
|
||||
def __truediv__(s, o): return _div(float(s), float(o)) if s._float else _div(int(s), int(o))
|
||||
def __rtruediv__(s, o): return _div(float(o), float(s)) if s._float else _div(int(o), int(s))
|
||||
def __pow__(s, o): return float(s) ** float(o) if s._float else int(s) ** int(o)
|
||||
def __rpow__(s, o): return float(o) ** float(s) if s._float else int(o) ** int(s)
|
||||
def __neg__(s): return -float(s) if s._float else -int(s)
|
||||
def __abs__(s): return abs(float(s)) if s._float else abs(int(s))
|
||||
|
||||
# Bitwise - GPU shifts mask the shift amount to valid range
|
||||
def __and__(s, o): return int(s) & int(o)
|
||||
def __or__(s, o): return int(s) | int(o)
|
||||
def __xor__(s, o): return int(s) ^ int(o)
|
||||
def __invert__(s): return ~int(s)
|
||||
def __lshift__(s, o): n = int(o); return int(s) << n if 0 <= n < 64 or s._nbits() > 64 else 0
|
||||
def __rshift__(s, o): n = int(o); return int(s) >> n if 0 <= n < 64 or s._nbits() > 64 else 0
|
||||
def __rand__(s, o): return int(o) & int(s)
|
||||
def __ror__(s, o): return int(o) | int(s)
|
||||
def __rxor__(s, o): return int(o) ^ int(s)
|
||||
def __rlshift__(s, o): n = int(s); return int(o) << n if 0 <= n < 64 else 0
|
||||
def __rrshift__(s, o): n = int(s); return int(o) >> n if 0 <= n < 64 else 0
|
||||
|
||||
# Comparison - handle _DenormChecker specially
|
||||
def __eq__(s, o):
|
||||
if isinstance(o, _DenormChecker): return o._check(s)
|
||||
return float(s) == float(o) if s._float else int(s) == int(o)
|
||||
def __ne__(s, o):
|
||||
if isinstance(o, _DenormChecker): return not o._check(s)
|
||||
return float(s) != float(o) if s._float else int(s) != int(o)
|
||||
def __lt__(s, o): return float(s) < float(o) if s._float else int(s) < int(o)
|
||||
def __le__(s, o): return float(s) <= float(o) if s._float else int(s) <= int(o)
|
||||
def __gt__(s, o): return float(s) > float(o) if s._float else int(s) > int(o)
|
||||
def __ge__(s, o): return float(s) >= float(o) if s._float else int(s) >= int(o)
|
||||
|
||||
class Reg:
|
||||
"""GPU register: D0.f32 = S0.f32 + S1.f32 just works. Supports up to 128 bits for DS_LOAD_B128."""
|
||||
__slots__ = ('_val',)
|
||||
def __init__(self, val=0): self._val = int(val)
|
||||
|
||||
# Typed views - TypedView(reg, high, signed, is_float, is_bf16)
|
||||
u64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
|
||||
i64 = property(lambda s: TypedView(s, 63, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK64))
|
||||
b64 = property(lambda s: TypedView(s, 63), lambda s, v: setattr(s, '_val', int(v) & MASK64))
|
||||
f64 = property(lambda s: TypedView(s, 63, is_float=True), lambda s, v: setattr(s, '_val', v if isinstance(v, int) else _i64(float(v))))
|
||||
u32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
|
||||
i32 = property(lambda s: TypedView(s, 31, signed=True), lambda s, v: setattr(s, '_val', int(v) & MASK32))
|
||||
b32 = property(lambda s: TypedView(s, 31), lambda s, v: setattr(s, '_val', int(v) & MASK32))
|
||||
f32 = property(lambda s: TypedView(s, 31, is_float=True), lambda s, v: setattr(s, '_val', _i32(float(v))))
|
||||
u24 = property(lambda s: TypedView(s, 23))
|
||||
i24 = property(lambda s: TypedView(s, 23, signed=True))
|
||||
u16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
|
||||
i16 = property(lambda s: TypedView(s, 15, signed=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
|
||||
b16 = property(lambda s: TypedView(s, 15), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | (int(v) & 0xffff)))
|
||||
f16 = property(lambda s: TypedView(s, 15, is_float=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _i16(float(v))) & 0xffff)))
|
||||
bf16 = property(lambda s: TypedView(s, 15, is_float=True, is_bf16=True), lambda s, v: setattr(s, '_val', (s._val & 0xffff0000) | ((v if isinstance(v, int) else _ibf16(float(v))) & 0xffff)))
|
||||
u8 = property(lambda s: TypedView(s, 7))
|
||||
i8 = property(lambda s: TypedView(s, 7, signed=True))
|
||||
u3 = property(lambda s: TypedView(s, 2)) # 3-bit for opsel fields
|
||||
u1 = property(lambda s: TypedView(s, 0)) # single bit
|
||||
|
||||
def __getitem__(s, key):
|
||||
if isinstance(key, slice): return TypedView(s, int(key.start), int(key.stop))
|
||||
return (s._val >> int(key)) & 1
|
||||
|
||||
def __setitem__(s, key, value):
|
||||
if isinstance(key, slice):
|
||||
high, low = int(key.start), int(key.stop)
|
||||
if high < low: high, low = low, high
|
||||
mask = (1 << (high - low + 1)) - 1
|
||||
s._val = (s._val & ~(mask << low)) | ((int(value) & mask) << low)
|
||||
elif value: s._val |= (1 << int(key))
|
||||
else: s._val &= ~(1 << int(key))
|
||||
|
||||
def __int__(s): return s._val
|
||||
def __index__(s): return s._val
|
||||
def __bool__(s): return bool(s._val)
|
||||
|
||||
# Arithmetic (for tmp = tmp + 1 patterns). Float operands trigger f32 interpretation.
|
||||
def __add__(s, o): return (_f32(s._val) + float(o)) if isinstance(o, float) else s._val + int(o)
|
||||
def __radd__(s, o): return (float(o) + _f32(s._val)) if isinstance(o, float) else int(o) + s._val
|
||||
def __sub__(s, o): return (_f32(s._val) - float(o)) if isinstance(o, float) else s._val - int(o)
|
||||
def __rsub__(s, o): return (float(o) - _f32(s._val)) if isinstance(o, float) else int(o) - s._val
|
||||
def __mul__(s, o): return (_f32(s._val) * float(o)) if isinstance(o, float) else s._val * int(o)
|
||||
def __rmul__(s, o): return (float(o) * _f32(s._val)) if isinstance(o, float) else int(o) * s._val
|
||||
def __and__(s, o): return s._val & int(o)
|
||||
def __rand__(s, o): return int(o) & s._val
|
||||
def __or__(s, o): return s._val | int(o)
|
||||
def __ror__(s, o): return int(o) | s._val
|
||||
def __xor__(s, o): return s._val ^ int(o)
|
||||
def __rxor__(s, o): return int(o) ^ s._val
|
||||
def __lshift__(s, o): n = int(o); return s._val << n if 0 <= n < 64 else 0
|
||||
def __rshift__(s, o): n = int(o); return s._val >> n if 0 <= n < 64 else 0
|
||||
def __invert__(s): return ~s._val
|
||||
|
||||
# Comparison (for tmp >= 0x100000000 patterns)
|
||||
def __lt__(s, o): return s._val < int(o)
|
||||
def __le__(s, o): return s._val <= int(o)
|
||||
def __gt__(s, o): return s._val > int(o)
|
||||
def __ge__(s, o): return s._val >= int(o)
|
||||
def __eq__(s, o): return s._val == int(o)
|
||||
def __ne__(s, o): return s._val != int(o)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PSEUDOCODE API - Functions and constants from AMD ISA pseudocode
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
# Rounding and float operations
|
||||
trunc, floor, ceil = _fpop(math.trunc), _fpop(math.floor), _fpop(math.ceil)
|
||||
def sqrt(x): return _SafeFloat(math.sqrt(x)) if x >= 0 else _SafeFloat(float("nan"))
|
||||
def log2(x): return math.log2(x) if x > 0 else (float("-inf") if x == 0 else float("nan"))
|
||||
def fract(x): return x - math.floor(x)
|
||||
def sin(x): return _trig(math.sin, x)
|
||||
def cos(x): return _trig(math.cos, x)
|
||||
def pow(a, b):
|
||||
try: return a ** b
|
||||
except OverflowError: return float("inf") if b > 0 else 0.0
|
||||
def isEven(x):
|
||||
x = float(x)
|
||||
if math.isinf(x) or math.isnan(x): return False
|
||||
return int(x) % 2 == 0
|
||||
def mantissa(f):
|
||||
if f == 0.0 or math.isinf(f) or math.isnan(f): return f
|
||||
m, _ = math.frexp(f)
|
||||
return m # AMD V_FREXP_MANT returns mantissa in [0.5, 1.0) range
|
||||
def signext_from_bit(val, bit):
|
||||
bit = int(bit)
|
||||
if bit == 0: return 0
|
||||
mask = (1 << bit) - 1
|
||||
val = int(val) & mask
|
||||
if val & (1 << (bit - 1)): return val - (1 << bit)
|
||||
return val
|
||||
|
||||
# Type conversions
|
||||
i32_to_f32 = u32_to_f32 = i32_to_f64 = u32_to_f64 = f32_to_f64 = f64_to_f32 = float
|
||||
def f32_to_i32(f): return _f_to_int(f, -2147483648, 2147483647)
|
||||
def f32_to_u32(f): return _f_to_int(f, 0, 4294967295)
|
||||
f64_to_i32, f64_to_u32 = f32_to_i32, f32_to_u32
|
||||
def f32_to_f16(f):
|
||||
f = float(f)
|
||||
if math.isnan(f): return 0x7e00 # f16 NaN
|
||||
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00 # f16 ±infinity
|
||||
try: return struct.unpack("<H", struct.pack("<e", f))[0]
|
||||
except OverflowError: return 0x7c00 if f > 0 else 0xfc00 # overflow -> ±infinity
|
||||
def f16_to_f32(v): return v if isinstance(v, float) else _f16_to_f32_bits(v)
|
||||
def i16_to_f16(v): return f32_to_f16(float(_sext(int(v) & 0xffff, 16)))
|
||||
def u16_to_f16(v): return f32_to_f16(float(int(v) & 0xffff))
|
||||
def f16_to_i16(bits): f = _f16_to_f32_bits(bits); return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
|
||||
def f16_to_u16(bits): f = _f16_to_f32_bits(bits); return max(0, min(65535, int(f))) if not math.isnan(f) else 0
|
||||
def bf16_to_f32(v): return _bf16(v) if isinstance(v, int) else float(v)
|
||||
def f32_to_bf16(f): return _ibf16(f)
|
||||
def u8_to_u32(v): return int(v) & 0xff
|
||||
def u4_to_u32(v): return int(v) & 0xf
|
||||
def u32_to_u16(u): return int(u) & 0xffff
|
||||
def i32_to_i16(i): return ((int(i) + 32768) & 0xffff) - 32768
|
||||
def f16_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
|
||||
def f16_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
|
||||
def f32_to_snorm(f): return max(-32768, min(32767, int(round(max(-1.0, min(1.0, f)) * 32767))))
|
||||
def f32_to_unorm(f): return max(0, min(65535, int(round(max(0.0, min(1.0, f)) * 65535))))
|
||||
def v_cvt_i16_f32(f): return max(-32768, min(32767, int(f))) if not math.isnan(f) else 0
|
||||
def v_cvt_u16_f32(f): return max(0, min(65535, int(f))) if not math.isnan(f) else 0
|
||||
def SAT8(v): return max(0, min(255, int(v)))
|
||||
def f32_to_u8(f): return max(0, min(255, int(f))) if not math.isnan(f) else 0
|
||||
|
||||
# Min/max operations
|
||||
def v_min_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _lt_neg_zero(a, b) else b)
|
||||
def v_max_f32(a, b): return a if math.isnan(b) else b if math.isnan(a) else (a if _gt_neg_zero(a, b) else b)
|
||||
v_min_f16, v_max_f16 = v_min_f32, v_max_f32
|
||||
v_min_i32, v_max_i32 = min, max
|
||||
v_min_i16, v_max_i16 = min, max
|
||||
def v_min_u32(a, b): return min(a & MASK32, b & MASK32)
|
||||
def v_max_u32(a, b): return max(a & MASK32, b & MASK32)
|
||||
def v_min_u16(a, b): return min(a & 0xffff, b & 0xffff)
|
||||
def v_max_u16(a, b): return max(a & 0xffff, b & 0xffff)
|
||||
def v_min3_f32(a, b, c): return v_min_f32(v_min_f32(a, b), c)
|
||||
def v_max3_f32(a, b, c): return v_max_f32(v_max_f32(a, b), c)
|
||||
v_min3_f16, v_max3_f16 = v_min3_f32, v_max3_f32
|
||||
v_min3_i32, v_max3_i32, v_min3_i16, v_max3_i16 = min, max, min, max
|
||||
def v_min3_u32(a, b, c): return min(a & MASK32, b & MASK32, c & MASK32)
|
||||
def v_max3_u32(a, b, c): return max(a & MASK32, b & MASK32, c & MASK32)
|
||||
def v_min3_u16(a, b, c): return min(a & 0xffff, b & 0xffff, c & 0xffff)
|
||||
def v_max3_u16(a, b, c): return max(a & 0xffff, b & 0xffff, c & 0xffff)
|
||||
|
||||
# SAD/MSAD operations
|
||||
def ABSDIFF(a, b): return abs(int(a) - int(b))
|
||||
def v_sad_u8(s0, s1, s2):
|
||||
"""V_SAD_U8: Sum of absolute differences of 4 byte pairs plus accumulator."""
|
||||
s0, s1, s2 = int(s0), int(s1), int(s2)
|
||||
result = s2
|
||||
for i in range(4):
|
||||
a = (s0 >> (i * 8)) & 0xff
|
||||
b = (s1 >> (i * 8)) & 0xff
|
||||
result += abs(a - b)
|
||||
return result & 0xffffffff
|
||||
def v_msad_u8(s0, s1, s2):
|
||||
"""V_MSAD_U8: Masked sum of absolute differences (skip if reference byte is 0)."""
|
||||
s0, s1, s2 = int(s0), int(s1), int(s2)
|
||||
result = s2
|
||||
for i in range(4):
|
||||
a = (s0 >> (i * 8)) & 0xff
|
||||
b = (s1 >> (i * 8)) & 0xff
|
||||
if b != 0: # Only add diff if reference (s1) byte is non-zero
|
||||
result += abs(a - b)
|
||||
return result & 0xffffffff
|
||||
|
||||
def BYTE_PERMUTE(data, sel):
|
||||
"""Select a byte from 64-bit data based on selector value."""
|
||||
sel = int(sel) & 0xff
|
||||
if sel <= 7: return (int(data) >> (sel * 8)) & 0xff
|
||||
if sel == 8: return 0xff if ((int(data) >> 15) & 1) else 0x00
|
||||
if sel == 9: return 0xff if ((int(data) >> 31) & 1) else 0x00
|
||||
if sel == 10: return 0xff if ((int(data) >> 47) & 1) else 0x00
|
||||
if sel == 11: return 0xff if ((int(data) >> 63) & 1) else 0x00
|
||||
if sel == 12: return 0x00
|
||||
return 0xff
|
||||
|
||||
# Pseudocode functions
|
||||
def s_ff1_i32_b32(v): return _ctz(v, 32)
|
||||
def s_ff1_i32_b64(v): return _ctz(v, 64)
|
||||
GT_NEG_ZERO, LT_NEG_ZERO = _gt_neg_zero, _lt_neg_zero
|
||||
def isNAN(x):
|
||||
try: return math.isnan(float(x))
|
||||
except (TypeError, ValueError): return False
|
||||
def isQuietNAN(x): return _check_nan_type(x, 1, True)
|
||||
def isSignalNAN(x): return _check_nan_type(x, 0, False)
|
||||
def fma(a, b, c):
|
||||
try: return math.fma(a, b, c)
|
||||
except ValueError: return float('nan')
|
||||
def ldexp(m, e): return math.ldexp(m, e)
|
||||
def sign(f): return 1 if math.copysign(1.0, f) < 0 else 0
|
||||
def exponent(f):
|
||||
if hasattr(f, '_bits') and hasattr(f, '_float') and f._float:
|
||||
raw = f._val
|
||||
if f._bits == 16: return (raw >> 10) & 0x1f
|
||||
if f._bits == 32: return (raw >> 23) & 0xff
|
||||
if f._bits == 64: return (raw >> 52) & 0x7ff
|
||||
f = float(f)
|
||||
if math.isinf(f) or math.isnan(f): return 255
|
||||
if f == 0.0: return 0
|
||||
try: bits = struct.unpack("<I", struct.pack("<f", f))[0]; return (bits >> 23) & 0xff
|
||||
except: return 0
|
||||
def signext(x): return int(x)
|
||||
def cvtToQuietNAN(x): return float('nan')
|
||||
|
||||
def F(x):
|
||||
"""32'F(x) or 64'F(x) - interpret x as float. If x is int, treat as bit pattern."""
|
||||
if isinstance(x, int): return _f32(x)
|
||||
if isinstance(x, TypedView): return x
|
||||
return float(x)
|
||||
|
||||
# Constants
|
||||
PI = math.pi
|
||||
WAVE32, WAVE64 = True, False
|
||||
OVERFLOW_F32, UNDERFLOW_F32 = float('inf'), 0.0
|
||||
OVERFLOW_F64, UNDERFLOW_F64 = float('inf'), 0.0
|
||||
MAX_FLOAT_F32 = 3.4028235e+38
|
||||
INF = _Inf()
|
||||
ROUND_MODE = _RoundMode()
|
||||
WAVE_MODE = _WaveMode()
|
||||
DENORM = _Denorm()
|
||||
|
||||
# 2/PI with 1201 bits of precision for V_TRIG_PREOP_F64
|
||||
TWO_OVER_PI_1201 = Reg(0x0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# COMPILER: pseudocode -> Python (minimal transforms)
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _filter_pseudocode(pseudocode: str) -> str:
|
||||
"""Filter raw PDF pseudocode to only include actual code lines."""
|
||||
pcode_lines, in_lambda, depth = [], 0, 0
|
||||
for line in pseudocode.split('\n'):
|
||||
s = line.strip()
|
||||
if not s: continue
|
||||
if '=>' in s or re.match(r'^[A-Z_]+\(', s): continue # Skip example lines
|
||||
if '= lambda(' in s: in_lambda += 1; continue # Skip lambda definitions
|
||||
if in_lambda > 0:
|
||||
if s.endswith(');'): in_lambda -= 1
|
||||
continue
|
||||
# Only include lines that look like pseudocode
|
||||
is_code = (any(p in s for p in ['D0.', 'D1.', 'S0.', 'S1.', 'S2.', 'SCC =', 'SCC ?', 'VCC', 'EXEC', 'tmp =', 'tmp[', 'lane =', 'PC =',
|
||||
'D0[', 'D1[', 'S0[', 'S1[', 'S2[', 'MEM[', 'RETURN_DATA', 'VADDR', 'VDATA', 'VDST', 'SADDR', 'OFFSET']) or
|
||||
s.startswith(('if ', 'else', 'elsif', 'endif', 'declare ', 'for ', 'endfor', '//')) or
|
||||
re.match(r'^[a-z_]+\s*=', s) or re.match(r'^[a-z_]+\[', s) or (depth > 0 and '=' in s))
|
||||
if s.startswith('if '): depth += 1
|
||||
elif s.startswith('endif'): depth = max(0, depth - 1)
|
||||
if is_code: pcode_lines.append(s)
|
||||
return '\n'.join(pcode_lines)
|
||||
|
||||
def _compile_pseudocode(pseudocode: str) -> str:
|
||||
"""Compile pseudocode to Python. Transforms are minimal - most syntax just works."""
|
||||
pseudocode = re.sub(r'\bpass\b', 'pass_', pseudocode) # 'pass' is Python keyword
|
||||
raw_lines = pseudocode.strip().split('\n')
|
||||
joined_lines: list[str] = []
|
||||
for line in raw_lines:
|
||||
line = line.strip()
|
||||
if joined_lines and (joined_lines[-1].rstrip().endswith(('||', '&&', '(', ',')) or
|
||||
(joined_lines[-1].count('(') > joined_lines[-1].count(')'))):
|
||||
joined_lines[-1] = joined_lines[-1].rstrip() + ' ' + line
|
||||
else:
|
||||
joined_lines.append(line)
|
||||
|
||||
lines = []
|
||||
indent, need_pass, in_first_match_loop = 0, False, False
|
||||
for line in joined_lines:
|
||||
line = line.split('//')[0].strip() # Strip C-style comments
|
||||
if not line: continue
|
||||
if line.startswith('if '):
|
||||
lines.append(' ' * indent + f"if {_expr(line[3:].rstrip(' then'))}:")
|
||||
indent += 1
|
||||
need_pass = True
|
||||
elif line.startswith('elsif '):
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
lines.append(' ' * indent + f"elif {_expr(line[6:].rstrip(' then'))}:")
|
||||
indent += 1
|
||||
need_pass = True
|
||||
elif line == 'else':
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
lines.append(' ' * indent + "else:")
|
||||
indent += 1
|
||||
need_pass = True
|
||||
elif line.startswith('endif'):
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
need_pass = False
|
||||
elif line.startswith('endfor'):
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
indent -= 1
|
||||
need_pass, in_first_match_loop = False, False
|
||||
elif line.startswith('declare '):
|
||||
pass
|
||||
elif m := re.match(r'for (\w+) in (.+?)\s*:\s*(.+?) do', line):
|
||||
start, end = _expr(m[2].strip()), _expr(m[3].strip())
|
||||
lines.append(' ' * indent + f"for {m[1]} in range({start}, int({end})+1):")
|
||||
indent += 1
|
||||
need_pass, in_first_match_loop = True, True
|
||||
elif '=' in line and not line.startswith('=='):
|
||||
need_pass = False
|
||||
line = line.rstrip(';')
|
||||
if m := re.match(r'\{\s*D1\.[ui]1\s*,\s*D0\.[ui]64\s*\}\s*=\s*(.+)', line):
|
||||
rhs = _expr(m[1])
|
||||
lines.append(' ' * indent + f"_full = {rhs}")
|
||||
lines.append(' ' * indent + f"D0.u64 = int(_full) & 0xffffffffffffffff")
|
||||
lines.append(' ' * indent + f"D1 = Reg((int(_full) >> 64) & 1)")
|
||||
elif any(op in line for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^=')):
|
||||
for op in ('+=', '-=', '*=', '/=', '|=', '&=', '^='):
|
||||
if op in line:
|
||||
lhs, rhs = line.split(op, 1)
|
||||
lines.append(' ' * indent + f"{lhs.strip()} {op} {_expr(rhs.strip())}")
|
||||
break
|
||||
else:
|
||||
lhs, rhs = line.split('=', 1)
|
||||
lhs_s, rhs_s = _expr(lhs.strip()), rhs.strip()
|
||||
stmt = _assign(lhs_s, _expr(rhs_s))
|
||||
if in_first_match_loop and rhs_s == 'i' and (lhs_s == 'tmp' or lhs_s == 'D0.i32'):
|
||||
stmt += "; break"
|
||||
lines.append(' ' * indent + stmt)
|
||||
if need_pass: lines.append(' ' * indent + "pass")
|
||||
return '\n'.join(lines)
|
||||
|
||||
def _assign(lhs: str, rhs: str) -> str:
|
||||
if lhs in ('tmp', 'SCC', 'VCC', 'EXEC', 'D0', 'D1', 'saveexec', 'PC'):
|
||||
return f"{lhs} = Reg({rhs})"
|
||||
return f"{lhs} = {rhs}"
|
||||
|
||||
def _expr(e: str) -> str:
|
||||
e = e.strip()
|
||||
e = e.replace('&&', ' and ').replace('||', ' or ').replace('<>', ' != ')
|
||||
e = re.sub(r'!([^=])', r' not \1', e)
|
||||
e = re.sub(r'\{\s*(\w+\.u32)\s*,\s*(\w+\.u32)\s*\}', r'_pack32(\1, \2)', e)
|
||||
def pack(m):
|
||||
hi, lo = _expr(m[1].strip()), _expr(m[2].strip())
|
||||
return f'_pack({hi}, {lo})'
|
||||
e = re.sub(r'\{\s*([^,{}]+)\s*,\s*([^,{}]+)\s*\}', pack, e)
|
||||
e = re.sub(r"1201'B\(2\.0\s*/\s*PI\)", "TWO_OVER_PI_1201", e)
|
||||
e = re.sub(r"\d+'([0-9a-fA-Fx]+)[UuFf]*", r'\1', e)
|
||||
e = re.sub(r"\d+'[FIBU]\(", "(", e)
|
||||
e = re.sub(r'\bB\(', '(', e)
|
||||
e = re.sub(r'([0-9a-fA-Fx])ULL\b', r'\1', e)
|
||||
e = re.sub(r'([0-9a-fA-Fx])LL\b', r'\1', e)
|
||||
e = re.sub(r'([0-9a-fA-Fx])U\b', r'\1', e)
|
||||
e = re.sub(r'(\d\.?\d*)F\b', r'\1', e)
|
||||
e = re.sub(r'(\[laneId\])\.[uib]\d+', r'\1', e)
|
||||
e = e.replace('+INF', 'INF').replace('-INF', '(-INF)')
|
||||
e = re.sub(r'NAN\.f\d+', 'float("nan")', e)
|
||||
def convert_verilog_slice(m):
|
||||
start, width = m.group(1).strip(), m.group(2).strip()
|
||||
return f'[({start}) + ({width}) - 1 : ({start})]'
|
||||
e = re.sub(r'\[([^:\[\]]+)\s*\+:\s*([^:\[\]]+)\]', convert_verilog_slice, e)
|
||||
def process_brackets(s):
|
||||
result, i = [], 0
|
||||
while i < len(s):
|
||||
if s[i] == '[':
|
||||
depth, start = 1, i + 1
|
||||
j = start
|
||||
while j < len(s) and depth > 0:
|
||||
if s[j] == '[': depth += 1
|
||||
elif s[j] == ']': depth -= 1
|
||||
j += 1
|
||||
inner = _expr(s[start:j-1])
|
||||
result.append('[' + inner + ']')
|
||||
i = j
|
||||
else:
|
||||
result.append(s[i])
|
||||
i += 1
|
||||
return ''.join(result)
|
||||
e = process_brackets(e)
|
||||
while '?' in e:
|
||||
depth, bracket, q = 0, 0, -1
|
||||
for i, c in enumerate(e):
|
||||
if c == '(': depth += 1
|
||||
elif c == ')': depth -= 1
|
||||
elif c == '[': bracket += 1
|
||||
elif c == ']': bracket -= 1
|
||||
elif c == '?' and depth == 0 and bracket == 0: q = i; break
|
||||
if q < 0: break
|
||||
depth, bracket, col = 0, 0, -1
|
||||
for i in range(q + 1, len(e)):
|
||||
if e[i] == '(': depth += 1
|
||||
elif e[i] == ')': depth -= 1
|
||||
elif e[i] == '[': bracket += 1
|
||||
elif e[i] == ']': bracket -= 1
|
||||
elif e[i] == ':' and depth == 0 and bracket == 0: col = i; break
|
||||
if col < 0: break
|
||||
cond, t, f = e[:q].strip(), e[q+1:col].strip(), e[col+1:].strip()
|
||||
e = f'(({t}) if ({cond}) else ({f}))'
|
||||
return e
|
||||
|
||||
def _apply_pseudocode_fixes(op_name: str, code: str) -> str:
|
||||
"""Apply known fixes for PDF pseudocode bugs."""
|
||||
if op_name == 'V_DIV_FMAS_F32':
|
||||
code = code.replace('D0.f32 = 2.0 ** 32 * fma(S0.f32, S1.f32, S2.f32)',
|
||||
'D0.f32 = (2.0 ** 64 if exponent(S2.f32) > 127 else 2.0 ** -64) * fma(S0.f32, S1.f32, S2.f32)')
|
||||
if op_name == 'V_DIV_FMAS_F64':
|
||||
code = code.replace('D0.f64 = 2.0 ** 64 * fma(S0.f64, S1.f64, S2.f64)',
|
||||
'D0.f64 = (2.0 ** 128 if exponent(S2.f64) > 1023 else 2.0 ** -128) * fma(S0.f64, S1.f64, S2.f64)')
|
||||
if op_name == 'V_DIV_SCALE_F32':
|
||||
code = code.replace('D0.f32 = float("nan")', 'VCC = Reg(1 << laneId); D0.f32 = float("nan")')
|
||||
code = code.replace('elif S1.f32 == DENORM.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif False:\n pass')
|
||||
code += '\nif S1.f32 == DENORM.f32:\n D0.f32 = float("nan")'
|
||||
code = code.replace('elif exponent(S2.f32) <= 23:\n D0.f32 = ldexp(S0.f32, 64)', 'elif exponent(S2.f32) <= 23:\n VCC = Reg(1 << laneId); D0.f32 = ldexp(S0.f32, 64)')
|
||||
code = code.replace('elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(0x1)\n if S0.f32 == S2.f32:\n D0.f32 = ldexp(S0.f32, 64)', 'elif S2.f32 / S1.f32 == DENORM.f32:\n VCC = Reg(1 << laneId)')
|
||||
if op_name == 'V_DIV_SCALE_F64':
|
||||
code = code.replace('D0.f64 = float("nan")', 'VCC = Reg(1 << laneId); D0.f64 = float("nan")')
|
||||
code = code.replace('elif S1.f64 == DENORM.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif False:\n pass')
|
||||
code += '\nif S1.f64 == DENORM.f64:\n D0.f64 = float("nan")'
|
||||
code = code.replace('elif exponent(S2.f64) <= 52:\n D0.f64 = ldexp(S0.f64, 128)', 'elif exponent(S2.f64) <= 52:\n VCC = Reg(1 << laneId); D0.f64 = ldexp(S0.f64, 128)')
|
||||
code = code.replace('elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(0x1)\n if S0.f64 == S2.f64:\n D0.f64 = ldexp(S0.f64, 128)', 'elif S2.f64 / S1.f64 == DENORM.f64:\n VCC = Reg(1 << laneId)')
|
||||
if op_name == 'V_DIV_FIXUP_F32':
|
||||
code = code.replace('D0.f32 = ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))',
|
||||
'D0.f32 = ((-OVERFLOW_F32) if (sign_out) else (OVERFLOW_F32)) if isNAN(S0.f32) else ((-abs(S0.f32)) if (sign_out) else (abs(S0.f32)))')
|
||||
if op_name == 'V_DIV_FIXUP_F64':
|
||||
code = code.replace('D0.f64 = ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))',
|
||||
'D0.f64 = ((-OVERFLOW_F64) if (sign_out) else (OVERFLOW_F64)) if isNAN(S0.f64) else ((-abs(S0.f64)) if (sign_out) else (abs(S0.f64)))')
|
||||
if op_name == 'V_TRIG_PREOP_F64':
|
||||
code = code.replace('result = F((TWO_OVER_PI_1201[1200 : 0] << shift.u32) & 0x1fffffffffffff)',
|
||||
'result = float(((TWO_OVER_PI_1201[1200 : 0] << int(shift)) >> (1201 - 53)) & 0x1fffffffffffff)')
|
||||
return code
|
||||
|
||||
def _generate_function(cls_name: str, op_name: str, pc: str, code: str) -> str:
|
||||
"""Generate a single compiled pseudocode function.
|
||||
Functions take int parameters and return dict of int values.
|
||||
Reg wrapping happens inside the function, only for registers actually used."""
|
||||
has_d1 = '{ D1' in pc
|
||||
is_cmpx = (cls_name in ('VOPCOp', 'VOP3Op')) and 'EXEC.u64[laneId]' in pc
|
||||
is_div_scale = 'DIV_SCALE' in op_name
|
||||
has_sdst = cls_name == 'VOP3SDOp' and ('VCC.u64[laneId]' in pc or is_div_scale)
|
||||
is_ds = cls_name == 'DSOp'
|
||||
is_flat = cls_name in ('FLATOp', 'GLOBALOp', 'SCRATCHOp')
|
||||
is_smem = cls_name == 'SMEMOp'
|
||||
has_s_array = 'S[i]' in pc # FMA_MIX style: S[0], S[1], S[2] array access
|
||||
combined = code + pc
|
||||
|
||||
fn_name = f"_{cls_name}_{op_name}"
|
||||
|
||||
# Detect which registers are used/modified
|
||||
def needs_init(name): return name in combined and not re.search(rf'^\s*{name}\s*=\s*Reg\(', code, re.MULTILINE)
|
||||
modifies_d0 = is_div_scale or bool(re.search(r'\bD0\b[.\[]', combined))
|
||||
modifies_exec = is_cmpx or bool(re.search(r'EXEC\.(u32|u64|b32|b64)\s*=', combined))
|
||||
modifies_vcc = has_sdst or bool(re.search(r'VCC\.(u32|u64|b32|b64)\s*=|VCC\.u64\[laneId\]\s*=', combined))
|
||||
modifies_scc = bool(re.search(r'\bSCC\s*=', combined))
|
||||
modifies_pc = bool(re.search(r'\bPC\s*=', combined))
|
||||
|
||||
# Build function signature and Reg init lines
|
||||
if is_smem:
|
||||
lines = [f"def {fn_name}(MEM, addr):"]
|
||||
reg_inits = ["ADDR=Reg(addr)", "SDATA=Reg(0)"]
|
||||
special_regs = []
|
||||
elif is_ds:
|
||||
lines = [f"def {fn_name}(MEM, addr, data0, data1, offset0, offset1):"]
|
||||
reg_inits = ["ADDR=Reg(addr)", "DATA0=Reg(data0)", "DATA1=Reg(data1)", "OFFSET0=Reg(offset0)", "OFFSET1=Reg(offset1)", "RETURN_DATA=Reg(0)"]
|
||||
special_regs = [('DATA', 'DATA0'), ('DATA2', 'DATA1'), ('OFFSET', 'OFFSET0'), ('ADDR_BASE', 'ADDR')]
|
||||
elif is_flat:
|
||||
lines = [f"def {fn_name}(MEM, addr, vdata, vdst):"]
|
||||
reg_inits = ["ADDR=addr", "VDATA=Reg(vdata)", "VDST=Reg(vdst)", "RETURN_DATA=Reg(0)"]
|
||||
special_regs = [('DATA', 'VDATA')]
|
||||
elif has_s_array:
|
||||
# FMA_MIX style: needs S[i] array, opsel, opsel_hi for source selection (neg/neg_hi applied in emu.py before call)
|
||||
lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None, opsel=0, opsel_hi=0):"]
|
||||
reg_inits = ["S0=Reg(s0)", "S1=Reg(s1)", "S2=Reg(s2)", "S=[S0,S1,S2]", "D0=Reg(d0)", "OPSEL=Reg(opsel)", "OPSEL_HI=Reg(opsel_hi)"]
|
||||
special_regs = []
|
||||
# Detect array declarations like "declare in : 32'F[3]" and create them (rename 'in' to 'ins' since 'in' is a keyword)
|
||||
if "in[" in combined:
|
||||
reg_inits.append("ins=[Reg(0),Reg(0),Reg(0)]")
|
||||
code = code.replace("in[", "ins[")
|
||||
else:
|
||||
lines = [f"def {fn_name}(s0, s1, s2, d0, scc, vcc, laneId, exec_mask, literal, VGPR, src0_idx=0, vdst_idx=0, pc=None):"]
|
||||
# Only create Regs for registers actually used in the pseudocode
|
||||
reg_inits = []
|
||||
if 'S0' in combined: reg_inits.append("S0=Reg(s0)")
|
||||
if 'S1' in combined: reg_inits.append("S1=Reg(s1)")
|
||||
if 'S2' in combined: reg_inits.append("S2=Reg(s2)")
|
||||
if modifies_d0 or 'D0' in combined: reg_inits.append("D0=Reg(s0)" if is_div_scale else "D0=Reg(d0)")
|
||||
if modifies_scc or 'SCC' in combined: reg_inits.append("SCC=Reg(scc)")
|
||||
if modifies_vcc or 'VCC' in combined: reg_inits.append("VCC=Reg(vcc)")
|
||||
if modifies_exec or 'EXEC' in combined: reg_inits.append("EXEC=Reg(exec_mask)")
|
||||
if modifies_pc or 'PC' in combined: reg_inits.append("PC=Reg(pc) if pc is not None else None")
|
||||
special_regs = [('D1', 'Reg(0)'), ('SIMM16', 'Reg(literal)'), ('SIMM32', 'Reg(literal)'),
|
||||
('SRC0', 'Reg(src0_idx)'), ('VDST', 'Reg(vdst_idx)')]
|
||||
if needs_init('tmp'): special_regs.insert(0, ('tmp', 'Reg(0)'))
|
||||
if needs_init('saveexec'): special_regs.insert(0, ('saveexec', 'Reg(EXEC._val)'))
|
||||
|
||||
# Build init code
|
||||
init_parts = reg_inits.copy()
|
||||
for name, init in special_regs:
|
||||
if name in combined: init_parts.append(f"{name}={init}")
|
||||
if 'EXEC_LO' in code: init_parts.append("EXEC_LO=TypedView(EXEC, 31, 0)")
|
||||
if 'EXEC_HI' in code: init_parts.append("EXEC_HI=TypedView(EXEC, 63, 32)")
|
||||
if 'VCCZ' in code and not re.search(r'^\s*VCCZ\s*=', code, re.MULTILINE): init_parts.append("VCCZ=Reg(1 if VCC._val == 0 else 0)")
|
||||
if 'EXECZ' in code and not re.search(r'^\s*EXECZ\s*=', code, re.MULTILINE): init_parts.append("EXECZ=Reg(1 if EXEC._val == 0 else 0)")
|
||||
|
||||
# Add init line and separator
|
||||
if init_parts: lines.append(f" {'; '.join(init_parts)}")
|
||||
|
||||
# Add compiled pseudocode
|
||||
for line in code.split('\n'):
|
||||
if line.strip(): lines.append(f" {line}")
|
||||
|
||||
# Build result dict
|
||||
result_items = []
|
||||
if modifies_d0: result_items.append("'D0': D0._val")
|
||||
if modifies_scc: result_items.append("'SCC': SCC._val")
|
||||
if modifies_vcc: result_items.append("'VCC': VCC._val")
|
||||
if modifies_exec: result_items.append("'EXEC': EXEC._val")
|
||||
if has_d1: result_items.append("'D1': D1._val")
|
||||
if modifies_pc: result_items.append("'PC': PC._val")
|
||||
if is_smem and 'SDATA' in combined and re.search(r'^\s*SDATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'SDATA': SDATA._val")
|
||||
if is_ds and 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'RETURN_DATA': RETURN_DATA._val")
|
||||
if is_flat:
|
||||
if 'RETURN_DATA' in combined and re.search(r'^\s*RETURN_DATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'RETURN_DATA': RETURN_DATA._val")
|
||||
if re.search(r'^\s*VDATA[\.\[].*=', code, re.MULTILINE):
|
||||
result_items.append("'VDATA': VDATA._val")
|
||||
lines.append(f" return {{{', '.join(result_items)}}}")
|
||||
return '\n'.join(lines)
|
||||
|
||||
# Build the globals dict for exec() - includes all pcode symbols
|
||||
_PCODE_GLOBALS = {
|
||||
'Reg': Reg, 'TypedView': TypedView, '_pack': _pack, '_pack32': _pack32,
|
||||
'ABSDIFF': ABSDIFF, 'BYTE_PERMUTE': BYTE_PERMUTE, 'DENORM': DENORM, 'F': F,
|
||||
'GT_NEG_ZERO': GT_NEG_ZERO, 'LT_NEG_ZERO': LT_NEG_ZERO, 'INF': INF,
|
||||
'MAX_FLOAT_F32': MAX_FLOAT_F32, 'OVERFLOW_F32': OVERFLOW_F32, 'OVERFLOW_F64': OVERFLOW_F64,
|
||||
'UNDERFLOW_F32': UNDERFLOW_F32, 'UNDERFLOW_F64': UNDERFLOW_F64,
|
||||
'PI': PI, 'ROUND_MODE': ROUND_MODE, 'WAVE_MODE': WAVE_MODE,
|
||||
'WAVE32': WAVE32, 'WAVE64': WAVE64, 'TWO_OVER_PI_1201': TWO_OVER_PI_1201,
|
||||
'SAT8': SAT8, 'trunc': trunc, 'floor': floor, 'ceil': ceil, 'sqrt': sqrt,
|
||||
'log2': log2, 'fract': fract, 'sin': sin, 'cos': cos, 'pow': pow,
|
||||
'isEven': isEven, 'mantissa': mantissa, 'signext_from_bit': signext_from_bit,
|
||||
'i32_to_f32': i32_to_f32, 'u32_to_f32': u32_to_f32, 'i32_to_f64': i32_to_f64,
|
||||
'u32_to_f64': u32_to_f64, 'f32_to_f64': f32_to_f64, 'f64_to_f32': f64_to_f32,
|
||||
'f32_to_i32': f32_to_i32, 'f32_to_u32': f32_to_u32, 'f64_to_i32': f64_to_i32,
|
||||
'f64_to_u32': f64_to_u32, 'f32_to_f16': f32_to_f16, 'f16_to_f32': f16_to_f32,
|
||||
'i16_to_f16': i16_to_f16, 'u16_to_f16': u16_to_f16, 'f16_to_i16': f16_to_i16,
|
||||
'f16_to_u16': f16_to_u16, 'bf16_to_f32': bf16_to_f32, 'f32_to_bf16': f32_to_bf16,
|
||||
'u8_to_u32': u8_to_u32, 'u4_to_u32': u4_to_u32, 'u32_to_u16': u32_to_u16,
|
||||
'i32_to_i16': i32_to_i16, 'f16_to_snorm': f16_to_snorm, 'f16_to_unorm': f16_to_unorm,
|
||||
'f32_to_snorm': f32_to_snorm, 'f32_to_unorm': f32_to_unorm,
|
||||
'v_cvt_i16_f32': v_cvt_i16_f32, 'v_cvt_u16_f32': v_cvt_u16_f32, 'f32_to_u8': f32_to_u8,
|
||||
'v_min_f32': v_min_f32, 'v_max_f32': v_max_f32, 'v_min_f16': v_min_f16, 'v_max_f16': v_max_f16,
|
||||
'v_min_i32': v_min_i32, 'v_max_i32': v_max_i32, 'v_min_i16': v_min_i16, 'v_max_i16': v_max_i16,
|
||||
'v_min_u32': v_min_u32, 'v_max_u32': v_max_u32, 'v_min_u16': v_min_u16, 'v_max_u16': v_max_u16,
|
||||
'v_min3_f32': v_min3_f32, 'v_max3_f32': v_max3_f32, 'v_min3_f16': v_min3_f16, 'v_max3_f16': v_max3_f16,
|
||||
'v_min3_i32': v_min3_i32, 'v_max3_i32': v_max3_i32, 'v_min3_i16': v_min3_i16, 'v_max3_i16': v_max3_i16,
|
||||
'v_min3_u32': v_min3_u32, 'v_max3_u32': v_max3_u32, 'v_min3_u16': v_min3_u16, 'v_max3_u16': v_max3_u16,
|
||||
'v_sad_u8': v_sad_u8, 'v_msad_u8': v_msad_u8,
|
||||
's_ff1_i32_b32': s_ff1_i32_b32, 's_ff1_i32_b64': s_ff1_i32_b64,
|
||||
'isNAN': isNAN, 'isQuietNAN': isQuietNAN, 'isSignalNAN': isSignalNAN,
|
||||
'fma': fma, 'ldexp': ldexp, 'sign': sign, 'exponent': exponent,
|
||||
'signext': signext, 'cvtToQuietNAN': cvtToQuietNAN,
|
||||
}
|
||||
|
||||
@functools.cache
|
||||
def compile_pseudocode(cls_name: str, op_name: str, pseudocode: str):
|
||||
"""Compile pseudocode string to executable function. Cached for performance."""
|
||||
filtered = _filter_pseudocode(pseudocode)
|
||||
code = _compile_pseudocode(filtered)
|
||||
code = _apply_pseudocode_fixes(op_name, code)
|
||||
fn_code = _generate_function(cls_name, op_name, filtered, code)
|
||||
fn_name = f"_{cls_name}_{op_name}"
|
||||
local_ns = {}
|
||||
exec(fn_code, _PCODE_GLOBALS, local_ns)
|
||||
return local_ns[fn_name]
|
||||
@@ -6,7 +6,9 @@ from pathlib import Path
|
||||
# Set AMD=1 before importing tinygrad
|
||||
os.environ["AMD"] = "1"
|
||||
|
||||
from extra.assembly.amd.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program
|
||||
from extra.assembly.amd.emu2 import run_asm as python_run_asm, decode_program, _get_inst_sink, _get_inst_prg
|
||||
from extra.assembly.amd.decode import decode_inst
|
||||
from extra.assembly.amd.autogen.rdna3.ins import SOPP, SOPPOp
|
||||
|
||||
REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
|
||||
if not REMU_PATH.exists():
|
||||
@@ -64,6 +66,131 @@ 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 instructions and return sorted by render time."""
|
||||
from extra.assembly.amd.emu2 import _get_inst_prg, _get_inst_sink, _canonical_prg_cache
|
||||
from tinygrad.codegen import get_program
|
||||
from extra.assembly.amd.emu2 import _emu_renderer
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
# Clear caches to measure fresh
|
||||
_get_inst_sink.cache_clear()
|
||||
_get_inst_prg.cache_clear()
|
||||
_canonical_prg_cache.clear()
|
||||
decode_program.cache_clear()
|
||||
|
||||
# Collect instruction bytes and names
|
||||
inst_data = []
|
||||
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__}>"
|
||||
inst_data.append((inst_bytes, inst_str, type(inst).__name__))
|
||||
i += inst.size()
|
||||
|
||||
# Profile each instruction
|
||||
from extra.assembly.amd.emu2 import _match_canonical
|
||||
results = []
|
||||
for inst_bytes, inst_str, inst_type in inst_data:
|
||||
# Check canonical cache BEFORE building sink (matches real behavior)
|
||||
inst_size = decode_inst(inst_bytes).size()
|
||||
inst_int = int.from_bytes(inst_bytes[:inst_size], 'little')
|
||||
is_cache_hit = _match_canonical(inst_int, inst_size) is not None
|
||||
|
||||
if is_cache_hit:
|
||||
# Skip build and render entirely for cache hits
|
||||
build_time, render_time, uop_count = 0, 0, 0
|
||||
else:
|
||||
# Build sink
|
||||
build_start = time.perf_counter()
|
||||
sink, ctx = _get_inst_sink(inst_bytes)
|
||||
build_time = time.perf_counter() - build_start
|
||||
|
||||
# Count UOps in sink
|
||||
uop_count = len(sink.toposort())
|
||||
|
||||
# Render
|
||||
render_start = time.perf_counter()
|
||||
with Context(NOOPT=1, IGNORE_OOB=1, TUPLE_ORDER=0):
|
||||
prg = get_program(sink, _emu_renderer)
|
||||
render_time = time.perf_counter() - render_start
|
||||
|
||||
# Update canonical cache
|
||||
base, mask, size = ctx.canonical_mask(inst_bytes)
|
||||
_canonical_prg_cache.append((base, mask, size, prg))
|
||||
|
||||
results.append({
|
||||
'inst_str': inst_str + (' [HIT]' if is_cache_hit else ''),
|
||||
'inst_type': inst_type,
|
||||
'uop_count': uop_count,
|
||||
'build_ms': build_time * 1000,
|
||||
'render_ms': render_time * 1000,
|
||||
})
|
||||
|
||||
# Sort by render time descending
|
||||
return sorted(results, key=lambda x: x['render_ms'], reverse=True)
|
||||
|
||||
def benchmark_python_split(kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
|
||||
"""Benchmark Python emulator with build/render/compile/execution times separated."""
|
||||
from extra.assembly.amd.emu2 import _emu_renderer, _emu_compiler, _elf_symbol_offsets
|
||||
from extra.assembly.amd.emu2 import _get_inst_prg, _get_inst_sink, _canonical_prg_cache
|
||||
from tinygrad.codegen import get_program
|
||||
from tinygrad.helpers import Context
|
||||
from tinygrad.runtime.support.elf import jit_loader
|
||||
|
||||
# Clear caches to measure fresh
|
||||
_get_inst_sink.cache_clear()
|
||||
_get_inst_prg.cache_clear()
|
||||
_canonical_prg_cache.clear()
|
||||
decode_program.cache_clear()
|
||||
|
||||
# Collect instruction bytes
|
||||
inst_bytes_list = []
|
||||
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_list.append(bytes(kernel[i:i + inst.size() + 4]))
|
||||
i += inst.size()
|
||||
|
||||
# Measure build time (UOp sink generation, cached)
|
||||
build_start = time.perf_counter()
|
||||
for inst_bytes in inst_bytes_list:
|
||||
_get_inst_sink(inst_bytes)
|
||||
build_time = time.perf_counter() - build_start
|
||||
|
||||
# Measure render time (uses cached sinks, handles canonical dedup)
|
||||
render_start = time.perf_counter()
|
||||
cache_before = len(_canonical_prg_cache)
|
||||
prgs = [_get_inst_prg(inst_bytes) for inst_bytes in inst_bytes_list]
|
||||
render_count = len(_canonical_prg_cache) - cache_before # number of unique renders
|
||||
render_time = time.perf_counter() - render_start
|
||||
|
||||
# Measure compile time (clang/llvm compile C to native)
|
||||
compile_start = time.perf_counter()
|
||||
# Deduplicate by function name (same as decode_program does)
|
||||
seen = set()
|
||||
unique_srcs = []
|
||||
for prg in prgs:
|
||||
if prg.function_name not in seen:
|
||||
seen.add(prg.function_name)
|
||||
unique_srcs.append(prg.src)
|
||||
combined_src = "\n".join(unique_srcs)
|
||||
obj = _emu_compiler.compile_to_obj(combined_src)
|
||||
_elf_symbol_offsets(obj)
|
||||
jit_loader(obj)
|
||||
compile_time = time.perf_counter() - compile_start
|
||||
|
||||
# Execution time (need to populate cache first)
|
||||
decode_program(kernel)
|
||||
exec_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, iterations)
|
||||
return build_time, render_time, render_count, compile_time, exec_time
|
||||
|
||||
def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
|
||||
"""Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
|
||||
try:
|
||||
@@ -119,14 +246,40 @@ def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], d
|
||||
print(f" Error getting kernel: {e}")
|
||||
return None
|
||||
|
||||
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
|
||||
TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "sin", "gelu", "matmul_small"]
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
|
||||
parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
|
||||
parser.add_argument("--profile", type=str, default=None, help="Profile instructions for a specific kernel (e.g. 'sin')")
|
||||
parser.add_argument("--top", type=int, default=20, help="Number of top instructions to show in profile")
|
||||
parser.add_argument("--sort-build", action="store_true", help="Sort profile by build time instead of render time")
|
||||
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("=" * 140)
|
||||
results = profile_instructions(kernel)
|
||||
if args.sort_build:
|
||||
results = sorted(results, key=lambda x: x['build_ms'], reverse=True)
|
||||
print(f"{'Instruction':<90} {'UOps':>6} {'Build(ms)':>10} {'Render(ms)':>10}")
|
||||
print("-" * 140)
|
||||
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['uop_count']:>6} {r['build_ms']:>10.3f} {r['render_ms']:>10.3f}")
|
||||
print("-" * 140)
|
||||
total_build = sum(r['build_ms'] for r in results)
|
||||
total_render = sum(r['render_ms'] for r in results)
|
||||
print(f"{'TOTAL':<90} {'':>6} {total_build:>10.3f} {total_render:>10.3f}")
|
||||
return
|
||||
|
||||
rust_remu = get_rust_remu()
|
||||
if rust_remu is None:
|
||||
print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
|
||||
@@ -149,44 +302,50 @@ def main():
|
||||
continue
|
||||
|
||||
kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
|
||||
n_insts = count_instructions(kernel)
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
|
||||
|
||||
# Benchmark Python emulator (must be first to measure compile time before cache is populated)
|
||||
py_build, py_render, render_count, py_compile, py_exec = benchmark_python_split(kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
|
||||
|
||||
n_insts = count_instructions(kernel) # uses cached decode_program
|
||||
n_workgroups = global_size[0] * global_size[1] * global_size[2]
|
||||
n_threads = local_size[0] * local_size[1] * local_size[2]
|
||||
total_work = n_insts * n_workgroups * n_threads
|
||||
|
||||
print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
|
||||
|
||||
buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
|
||||
set_valid_mem_ranges(ranges)
|
||||
|
||||
py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
|
||||
rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, rsrc2, args.iterations) if rust_remu else None
|
||||
|
||||
if py_time:
|
||||
py_rate = total_work / py_time / 1e6
|
||||
print(f" Python: {py_time*1000:8.3f} ms ({py_rate:7.2f} M ops/s)")
|
||||
if py_build is not None:
|
||||
py_exec_rate = total_work / py_exec / 1e6
|
||||
print(f" Build: {py_build*1000:8.3f} ms")
|
||||
print(f" Render: {py_render*1000:8.3f} ms ({render_count} unique)")
|
||||
print(f" Compile: {py_compile*1000:8.3f} ms")
|
||||
print(f" Exec: {py_exec*1000:8.3f} ms ({py_exec_rate:7.2f} M ops/s)")
|
||||
if rust_time:
|
||||
rust_rate = total_work / rust_time / 1e6
|
||||
speedup = py_time / rust_time if py_time else 0
|
||||
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
|
||||
speedup = py_exec / rust_time if py_exec else 0
|
||||
print(f" Rust: {rust_time*1000:8.3f} ms ({rust_rate:7.2f} M ops/s) [{speedup:.1f}x faster]")
|
||||
|
||||
results.append((op_name, n_insts, n_workgroups, py_time, rust_time))
|
||||
results.append((op_name, n_insts, n_workgroups, py_build, py_render, render_count, py_compile, py_exec, rust_time))
|
||||
|
||||
# Summary table
|
||||
print("\n" + "=" * 90)
|
||||
print("\n" + "=" * 140)
|
||||
print("SUMMARY")
|
||||
print("=" * 90)
|
||||
print(f"{'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
|
||||
print("-" * 90)
|
||||
print("=" * 140)
|
||||
print(f"{'Name':<16} {'Insts':<6} {'WGs':<5} {'Build (ms)':<12} {'Render (ms)':<16} {'Compile (ms)':<14} {'Exec (ms)':<12} {'Rust (ms)':<12} {'Speedup':<10}")
|
||||
print("-" * 140)
|
||||
|
||||
for name, n_insts, n_wgs, py_time, rust_time in results:
|
||||
py_ms = f"{py_time*1000:.3f}" if py_time else "error"
|
||||
for name, n_insts, n_wgs, py_build, py_render, render_count, py_compile, py_exec, rust_time in results:
|
||||
build_ms = f"{py_build*1000:.3f}" if py_build else "error"
|
||||
render_ms = f"{py_render*1000:.3f} ({render_count})" if py_render else "error"
|
||||
compile_ms = f"{py_compile*1000:.3f}" if py_compile else "error"
|
||||
exec_ms = f"{py_exec*1000:.3f}" if py_exec else "error"
|
||||
if rust_time:
|
||||
rust_ms = f"{rust_time*1000:.3f}"
|
||||
speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
|
||||
speedup = f"{py_exec/rust_time:.1f}x" if py_exec else "N/A"
|
||||
else:
|
||||
rust_ms, speedup = "N/A", "N/A"
|
||||
print(f"{name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
|
||||
print(f"{name:<16} {n_insts:<6} {n_wgs:<5} {build_ms:<12} {render_ms:<16} {compile_ms:<14} {exec_ms:<12} {rust_ms:<12} {speedup:<10}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -35,46 +35,4 @@ TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items
|
||||
def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
|
||||
|
||||
def get_mattr(arch:str) -> str:
|
||||
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# EXECUTION CONTEXT (for testing compiled pseudocode)
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class ExecContext:
|
||||
"""Context for running compiled pseudocode in tests."""
|
||||
def __init__(self, s0=0, s1=0, s2=0, d0=0, scc=0, vcc=0, lane=0, exec_mask=0xffffffff, literal=0, vgprs=None, src0_idx=0, vdst_idx=0):
|
||||
from extra.assembly.amd.pcode import Reg, MASK32, MASK64, TypedView
|
||||
self._Reg, self._MASK64, self._TypedView = Reg, MASK64, TypedView
|
||||
self.S0, self.S1, self.S2 = Reg(s0), Reg(s1), Reg(s2)
|
||||
self.D0, self.D1 = Reg(d0), Reg(0)
|
||||
self.SCC, self.VCC, self.EXEC = Reg(scc), Reg(vcc), Reg(exec_mask)
|
||||
self.tmp, self.saveexec = Reg(0), Reg(exec_mask)
|
||||
self.lane, self.laneId, self.literal = lane, lane, literal
|
||||
self.SIMM16, self.SIMM32 = Reg(literal), Reg(literal)
|
||||
self.VGPR = vgprs if vgprs is not None else {}
|
||||
self.SRC0, self.VDST = Reg(src0_idx), Reg(vdst_idx)
|
||||
|
||||
def run(self, code: str):
|
||||
"""Execute compiled code."""
|
||||
import extra.assembly.amd.pcode as pcode
|
||||
ns = {k: getattr(pcode, k) for k in dir(pcode) if not k.startswith('_')}
|
||||
# Also include underscore-prefixed helpers that compiled pseudocode uses
|
||||
for k in ['_pack', '_pack32']:
|
||||
if hasattr(pcode, k): ns[k] = getattr(pcode, k)
|
||||
ns.update({
|
||||
'S0': self.S0, 'S1': self.S1, 'S2': self.S2, 'D0': self.D0, 'D1': self.D1,
|
||||
'SCC': self.SCC, 'VCC': self.VCC, 'EXEC': self.EXEC,
|
||||
'EXEC_LO': self._TypedView(self.EXEC, 31, 0), 'EXEC_HI': self._TypedView(self.EXEC, 63, 32),
|
||||
'tmp': self.tmp, 'saveexec': self.saveexec,
|
||||
'lane': self.lane, 'laneId': self.laneId, 'literal': self.literal,
|
||||
'SIMM16': self.SIMM16, 'SIMM32': self.SIMM32, 'VGPR': self.VGPR, 'SRC0': self.SRC0, 'VDST': self.VDST,
|
||||
})
|
||||
exec(code, ns)
|
||||
def _sync(ctx_reg, ns_val):
|
||||
if isinstance(ns_val, self._Reg): ctx_reg._val = ns_val._val
|
||||
else: ctx_reg._val = int(ns_val) & self._MASK64
|
||||
for name in ('SCC', 'VCC', 'EXEC', 'D0', 'D1', 'tmp', 'saveexec'):
|
||||
if ns.get(name) is not getattr(self, name): _sync(getattr(self, name), ns[name])
|
||||
|
||||
def result(self) -> dict: return {"d0": self.D0._val, "scc": self.SCC._val & 1}
|
||||
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
|
||||
@@ -6,7 +6,7 @@ Set USE_HW=1 to run on both emulator and hardware, comparing results.
|
||||
import ctypes, math, os, struct
|
||||
from extra.assembly.amd.autogen.rdna3.ins import *
|
||||
|
||||
from extra.assembly.amd.emu import run_asm
|
||||
from extra.assembly.amd.emu2 import run_asm
|
||||
from extra.assembly.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
|
||||
|
||||
def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
|
||||
@@ -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)
|
||||
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2, scratch_size)
|
||||
assert result == 0, f"run_asm failed with {result}"
|
||||
|
||||
return parse_output(bytes(out_buf), n_lanes)
|
||||
@@ -204,6 +204,7 @@ amdhsa.kernels:
|
||||
prg = AMDProgram(dev, "test", lib)
|
||||
|
||||
out_gpu = dev.allocator.alloc(OUT_BYTES)
|
||||
assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
|
||||
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
|
||||
|
||||
out_buf = bytearray(OUT_BYTES)
|
||||
|
||||
@@ -523,5 +523,157 @@ class TestD16HiLoads(unittest.TestCase):
|
||||
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
|
||||
|
||||
|
||||
class TestGlobalOffset(unittest.TestCase):
|
||||
"""Tests for GLOBAL instructions with different offsets.
|
||||
|
||||
These tests verify that instruction deduplication correctly handles different offset values.
|
||||
If offset is made dynamic incorrectly, instructions with different offsets may load/store wrong data.
|
||||
"""
|
||||
|
||||
def test_global_load_different_offsets(self):
|
||||
"""Load from two different offsets and verify correct values."""
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
# Store 0xAAAAAAAA at offset 100
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
|
||||
# Store 0xBBBBBBBB at offset 200
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
|
||||
s_waitcnt(vmcnt=0),
|
||||
# Load from offset 100 -> should get 0xAAAAAAAA
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[3], saddr=SrcEnum.NULL, offset=100),
|
||||
# Load from offset 200 -> should get 0xBBBBBBBB
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=200),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[3]),
|
||||
v_mov_b32_e32(v[1], v[4]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
|
||||
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
|
||||
|
||||
def test_global_store_different_offsets(self):
|
||||
"""Store to two different offsets and verify correct values."""
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
# Store 0x11111111 at offset 300
|
||||
s_mov_b32(s[0], 0x11111111),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=300),
|
||||
# Store 0x22222222 at offset 400
|
||||
s_mov_b32(s[0], 0x22222222),
|
||||
v_mov_b32_e32(v[3], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[3], saddr=SrcEnum.NULL, offset=400),
|
||||
s_waitcnt(vmcnt=0),
|
||||
# Load back to verify
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=300),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[5], saddr=SrcEnum.NULL, offset=400),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[4]),
|
||||
v_mov_b32_e32(v[1], v[5]),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0x11111111, f"offset 300: expected 0x11111111, got 0x{st.vgpr[0][0]:08x}")
|
||||
self.assertEqual(st.vgpr[0][1], 0x22222222, f"offset 400: expected 0x22222222, got 0x{st.vgpr[0][1]:08x}")
|
||||
|
||||
def test_global_negative_offset_no_saddr(self):
|
||||
"""Test negative offset without saddr (VGPR pair for address).
|
||||
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
|
||||
Load with offset -100 from vaddr pointing to base+200 -> should get 0xAAAA (at 100).
|
||||
Load with offset -100 from vaddr pointing to base+300 -> should get 0xBBBB (at 200)."""
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], s[2]),
|
||||
v_mov_b32_e32(v[1], s[3]),
|
||||
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
|
||||
s_waitcnt(vmcnt=0),
|
||||
# vaddr = base+200, load with offset -100 -> should get value at 100
|
||||
s_add_u32(s[4], s[2], 200),
|
||||
s_addc_u32(s[5], s[3], 0),
|
||||
v_mov_b32_e32(v[4], s[4]),
|
||||
v_mov_b32_e32(v[5], s[5]),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[6], saddr=SrcEnum.NULL, offset=-100),
|
||||
# vaddr = base+300, load with offset -100 -> should get value at 200
|
||||
s_add_u32(s[4], s[2], 300),
|
||||
s_addc_u32(s[5], s[3], 0),
|
||||
v_mov_b32_e32(v[4], s[4]),
|
||||
v_mov_b32_e32(v[5], s[5]),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[7], saddr=SrcEnum.NULL, offset=-100),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[6]),
|
||||
v_mov_b32_e32(v[1], v[7]),
|
||||
v_mov_b32_e32(v[4], 0),
|
||||
v_mov_b32_e32(v[5], 0),
|
||||
v_mov_b32_e32(v[6], 0),
|
||||
v_mov_b32_e32(v[7], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
s_mov_b32(s[4], 0),
|
||||
s_mov_b32(s[5], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
|
||||
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
|
||||
|
||||
def test_global_negative_offset_with_saddr(self):
|
||||
"""Test negative offset with saddr (SGPR pair for base address).
|
||||
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
|
||||
Load with offset -100 from saddr pointing to base+200 -> should get 0xAAAA (at 100).
|
||||
Load with offset -100 from saddr pointing to base+300 -> should get 0xBBBB (at 200)."""
|
||||
instructions = [
|
||||
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
|
||||
s_waitcnt(lgkmcnt=0),
|
||||
v_mov_b32_e32(v[0], 0),
|
||||
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
|
||||
s_mov_b32(s[0], 0xAAAAAAAA),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
|
||||
s_mov_b32(s[0], 0xBBBBBBBB),
|
||||
v_mov_b32_e32(v[2], s[0]),
|
||||
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
|
||||
s_waitcnt(vmcnt=0),
|
||||
# saddr = base+200, load with offset -100 -> should get value at 100
|
||||
s_add_u32(s[4], s[2], 200),
|
||||
s_addc_u32(s[5], s[3], 0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[6], saddr=s[4:5], offset=-100),
|
||||
# saddr = base+300, load with offset -100 -> should get value at 200
|
||||
s_add_u32(s[4], s[2], 300),
|
||||
s_addc_u32(s[5], s[3], 0),
|
||||
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[7], saddr=s[4:5], offset=-100),
|
||||
s_waitcnt(vmcnt=0),
|
||||
v_mov_b32_e32(v[0], v[6]),
|
||||
v_mov_b32_e32(v[1], v[7]),
|
||||
v_mov_b32_e32(v[6], 0),
|
||||
v_mov_b32_e32(v[7], 0),
|
||||
s_mov_b32(s[2], 0),
|
||||
s_mov_b32(s[3], 0),
|
||||
s_mov_b32(s[4], 0),
|
||||
s_mov_b32(s[5], 0),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
|
||||
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,282 @@
|
||||
"""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 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()
|
||||
@@ -0,0 +1,341 @@
|
||||
"""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 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,5 +619,105 @@ 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)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
@@ -237,6 +237,32 @@ class TestF16Ops(unittest.TestCase):
|
||||
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
|
||||
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_max_f16_basic(self):
|
||||
"""V_MAX_F16 returns the maximum of two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_max_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x4000, f"Expected 0x4000 (f16 2.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_min_f16_basic(self):
|
||||
"""V_MIN_F16 returns the minimum of two f16 values."""
|
||||
instructions = [
|
||||
s_mov_b32(s[0], 0x3c00), # f16 1.0
|
||||
s_mov_b32(s[1], 0x4000), # f16 2.0
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[1]),
|
||||
v_min_f16_e32(v[2], v[0], v[1]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
result = st.vgpr[0][2] & 0xffff
|
||||
self.assertEqual(result, 0x3c00, f"Expected 0x3c00 (f16 1.0), got 0x{result:04x}")
|
||||
|
||||
def test_v_fmaak_f16_basic(self):
|
||||
"""V_FMAAK_F16: d = a * b + K."""
|
||||
instructions = [
|
||||
@@ -810,6 +836,81 @@ class TestCarryOps(unittest.TestCase):
|
||||
self.assertEqual(st.vgpr[0][2], 0) # Overflowed to 0
|
||||
self.assertEqual(st.vcc, 1) # Carry out
|
||||
|
||||
def test_v_add_co_ci_u32_clears_carry(self):
|
||||
"""V_ADD_CO_CI_U32: VCC must be updated even when no carry is generated.
|
||||
|
||||
This tests the case where VCC=1 going in (carry-in consumed) but the addition
|
||||
does not overflow, so VCC must be cleared to 0.
|
||||
|
||||
Regression test for: VCC not being written by v_add_co_ci_u32_e32.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
|
||||
v_mov_b32_e32(v[0], 1), # S0 = 1
|
||||
v_mov_b32_e32(v[1], 1), # S1 = 1
|
||||
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][2], 3) # 1 + 1 + 1 = 3
|
||||
self.assertEqual(st.vcc, 0) # No carry out - VCC must be cleared
|
||||
|
||||
def test_v_add_co_ci_u32_multilane_clears_vcc(self):
|
||||
"""V_ADD_CO_CI_U32 with multiple lanes: VCC bits must be updated per-lane.
|
||||
|
||||
When VCC has multiple bits set (one per active lane), and the addition doesn't
|
||||
overflow for any lane, all VCC bits must be cleared.
|
||||
|
||||
Regression test for: VCC not being written by v_add_co_ci_u32_e32 in multi-lane case.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 0b11), # VCC = 0b11 (lanes 0,1 have carry-in)
|
||||
v_mov_b32_e32(v[0], 1), # S0 = 1 for all lanes
|
||||
v_mov_b32_e32(v[1], 1), # S1 = 1 for all lanes
|
||||
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
|
||||
]
|
||||
st = run_program(instructions, n_lanes=2)
|
||||
self.assertEqual(st.vgpr[0][2], 3) # lane 0: 1 + 1 + 1 = 3
|
||||
self.assertEqual(st.vgpr[1][2], 3) # lane 1: 1 + 1 + 1 = 3
|
||||
self.assertEqual(st.vcc, 0) # No carry out for any lane - all VCC bits must be cleared
|
||||
|
||||
def test_v_add_co_ci_u32_preserves_inactive_vcc_bits(self):
|
||||
"""V_ADD_CO_CI_U32: VCC carry-out overwrites entire VCC register.
|
||||
|
||||
VOP2 carry instructions write ALL VCC bits based on carry-out, clearing
|
||||
bits for lanes that don't overflow regardless of EXEC mask.
|
||||
|
||||
Note: This differs from VOPC which only writes active lane bits.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 0x00010000), # VCC bit 16 set
|
||||
v_mov_b32_e32(v[0], 1), # S0 = 1
|
||||
v_mov_b32_e32(v[1], 1), # S1 = 1
|
||||
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 0 = 2 (no carry)
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
self.assertEqual(st.vgpr[0][2], 2) # lane 0: 1 + 1 + 0 = 2
|
||||
# VCC should be completely cleared (all lanes have no carry-out)
|
||||
self.assertEqual(st.vcc, 0)
|
||||
|
||||
def test_v_add_co_ci_u32_all_lanes_same_result(self):
|
||||
"""V_ADD_CO_CI_U32: all active lanes should produce the same result.
|
||||
|
||||
When the same constant inputs are used across all lanes, each lane should
|
||||
compute the same result and write to its own VGPR slot.
|
||||
|
||||
Regression test for: VGPR writes not happening for all lanes.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 0), # No carry-in
|
||||
v_mov_b32_e32(v[0], 3), # inline constant 3
|
||||
v_mov_b32_e32(v[1], 5), # value 5
|
||||
v_add_co_ci_u32_e32(v[1], 3, v[1]), # v[1] = 3 + v[1] + 0 = 3 + 5 = 8
|
||||
]
|
||||
st = run_program(instructions, n_lanes=4)
|
||||
# All 4 lanes should have v[1] = 8
|
||||
for lane in range(4):
|
||||
self.assertEqual(st.vgpr[lane][1], 8, f"lane {lane} should have v[1]=8")
|
||||
|
||||
def test_v_sub_co_ci_u32_no_borrow(self):
|
||||
"""V_SUB_CO_CI_U32: D0 = S0 - S1 - VCC_IN, when VCC_IN=0."""
|
||||
instructions = [
|
||||
@@ -860,6 +961,23 @@ class TestCarryOps(unittest.TestCase):
|
||||
self.assertEqual(st.vgpr[0][0], 16)
|
||||
self.assertEqual(st.sgpr[10], 0) # No carry out
|
||||
|
||||
def test_v_add_co_ci_u32_vop3sd_null_sdst(self):
|
||||
"""VOP3SD V_ADD_CO_CI_U32 with sdst=NULL: carry output is discarded.
|
||||
|
||||
When sdst=NULL (register 124), the carry-out should NOT be written anywhere.
|
||||
We verify this by checking that VCC (which we set to a sentinel value) is unchanged.
|
||||
"""
|
||||
instructions = [
|
||||
s_mov_b32(VCC_LO, 0xDEADBEEF), # Sentinel value in VCC
|
||||
s_mov_b32(s[6], 0), # carry-in = 0
|
||||
# VOP3SD with NULL sdst: carry-out should be discarded
|
||||
# Uses 0xFFFFFFFF + 1 + 0 = 0 with carry-out=1, but carry should not be written
|
||||
v_add_co_ci_u32(v[0], NULL, 0xFFFFFFFF, 1, s[6]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vgpr[0][0], 0) # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
|
||||
self.assertEqual(st.vcc, 0xDEADBEEF) # VCC unchanged - carry was discarded
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
@@ -404,8 +404,99 @@ 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."""
|
||||
"""Tests for WMMA (Wave Matrix Multiply-Accumulate) instructions with F32 output."""
|
||||
|
||||
def test_v_wmma_f32_16x16x16_f16_all_ones(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with all ones produces 16.0."""
|
||||
@@ -440,6 +531,75 @@ class TestWMMA(unittest.TestCase):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
|
||||
|
||||
def test_v_wmma_f32_16x16x16_f16_high_registers(self):
|
||||
"""V_WMMA_F32_16X16X16_F16 with high register indices.
|
||||
|
||||
Regression test: WMMA was using static register indices instead of dynamic,
|
||||
causing incorrect results when registers weren't at the default positions.
|
||||
This test uses v[64:71] for A, v[80:87] for B, v[96:103] for C/D.
|
||||
"""
|
||||
instructions = []
|
||||
instructions.append(s_mov_b32(s[0], 0x3c003c00)) # packed f16 1.0
|
||||
# Initialize A matrix in v[64:71]
|
||||
for i in range(64, 72):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
# Initialize B matrix in v[80:87]
|
||||
for i in range(80, 88):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
# Initialize C (accumulator) in v[96:103] to zero
|
||||
for i in range(96, 104):
|
||||
instructions.append(v_mov_b32_e32(v[i], 0))
|
||||
# WMMA: D = A @ B + C, result in v[96:103]
|
||||
instructions.append(v_wmma_f32_16x16x16_f16(v[96:103], v[64:71], v[80:87], v[96:103]))
|
||||
# Copy results to v[0:7] for checking
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], v[96+i]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(16.0)
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
|
||||
|
||||
|
||||
class TestWMMABF16(unittest.TestCase):
|
||||
"""Tests for WMMA BF16 instructions."""
|
||||
|
||||
def test_v_wmma_f32_16x16x16_bf16_all_ones(self):
|
||||
"""V_WMMA_F32_16X16X16_BF16 with all ones produces 16.0."""
|
||||
instructions = []
|
||||
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
|
||||
instructions.append(s_mov_b32(s[0], 0x3f803f80))
|
||||
for i in range(16, 32):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], 0))
|
||||
instructions.append(v_wmma_f32_16x16x16_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(16.0)
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 16.0, got {i2f(result)}")
|
||||
|
||||
def test_v_wmma_f32_16x16x16_bf16_with_accumulator(self):
|
||||
"""V_WMMA_F32_16X16X16_BF16 with non-zero accumulator."""
|
||||
instructions = []
|
||||
# BF16 1.0 = 0x3f80, packed = 0x3f803f80
|
||||
instructions.append(s_mov_b32(s[0], 0x3f803f80))
|
||||
instructions.append(s_mov_b32(s[1], f2i(5.0)))
|
||||
for i in range(16, 32):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[0]))
|
||||
for i in range(8):
|
||||
instructions.append(v_mov_b32_e32(v[i], s[1]))
|
||||
instructions.append(v_wmma_f32_16x16x16_bf16(v[0:7], v[16:23], v[24:31], v[0:7]))
|
||||
st = run_program(instructions, n_lanes=32)
|
||||
expected = f2i(21.0) # 16 + 5
|
||||
for lane in range(32):
|
||||
for reg in range(8):
|
||||
result = st.vgpr[lane][reg]
|
||||
self.assertEqual(result, expected, f"v[{reg}] lane {lane}: expected 21.0, got {i2f(result)}")
|
||||
|
||||
|
||||
class TestSpecialOps(unittest.TestCase):
|
||||
"""Tests for special operations (SAD, PERM, DOT2)."""
|
||||
|
||||
@@ -731,6 +731,111 @@ class TestVCCBehavior(unittest.TestCase):
|
||||
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
|
||||
|
||||
|
||||
class TestCmpNge(unittest.TestCase):
|
||||
"""Tests for V_CMP_NGE (not-greater-or-equal) with NaN semantics.
|
||||
|
||||
NGE = !(a >= b). With NaN inputs:
|
||||
- If either input is NaN, a >= b is false, so !(false) = true
|
||||
- This differs from a < b which returns false for NaN inputs
|
||||
"""
|
||||
|
||||
def test_v_cmp_nge_f32_normal_values(self):
|
||||
"""v_cmp_nge_f32: basic comparison with normal floats."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[1], f2i(2.0)),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 2.0) = !(false) = true
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= 2.0) should be true")
|
||||
|
||||
def test_v_cmp_nge_f32_equal_values(self):
|
||||
"""v_cmp_nge_f32: equal values should return false."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[1], f2i(1.0)),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 1.0) = !(true) = false
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "!(1.0 >= 1.0) should be false")
|
||||
|
||||
def test_v_cmp_nge_f32_greater_value(self):
|
||||
"""v_cmp_nge_f32: greater value should return false."""
|
||||
instructions = [
|
||||
v_mov_b32_e32(v[0], f2i(2.0)),
|
||||
v_mov_b32_e32(v[1], f2i(1.0)),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(2.0 >= 1.0) = !(true) = false
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 0, "!(2.0 >= 1.0) should be false")
|
||||
|
||||
def test_v_cmp_nge_f32_neg_inf(self):
|
||||
"""v_cmp_nge_f32: -inf compared to normal value."""
|
||||
neg_inf = 0xff800000 # -inf
|
||||
instructions = [
|
||||
s_mov_b32(s[0], neg_inf),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], f2i(1.0)),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(-inf >= 1.0) = !(false) = true
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "!(-inf >= 1.0) should be true")
|
||||
|
||||
def test_v_cmp_nge_f32_clears_inactive_vcc_bits(self):
|
||||
"""v_cmp_nge_f32 with partial EXEC clears inactive VCC bits (hardware behavior)."""
|
||||
neg_inf = 0xff800000 # -inf
|
||||
instructions = [
|
||||
# Set VCC to all 1s first
|
||||
s_mov_b32(VCC_LO, 0xFFFFFFFF),
|
||||
# Set EXEC to only lane 0
|
||||
s_mov_b32(EXEC_LO, 0x00000001),
|
||||
# v0 = 1.0 for lane 0
|
||||
v_mov_b32_e32(v[0], f2i(1.0)),
|
||||
# Compare: !(-inf >= 1.0) = true for lane 0
|
||||
v_cmp_nge_f32_e32(neg_inf, v[0]),
|
||||
]
|
||||
st = run_program(instructions, n_lanes=16)
|
||||
# Hardware clears inactive lane bits, only active lane results remain
|
||||
# Lane 0 result = 1 (true), lanes 1-15 = 0 (cleared)
|
||||
self.assertEqual(st.vcc, 0x00000001, "VCC should only have active lane results")
|
||||
|
||||
def test_v_cmp_nge_f32_nan_src0(self):
|
||||
"""v_cmp_nge_f32: NaN in src0 should return true (NaN >= x is false)."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], f2i(1.0)),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= 1.0) = !(false) = true
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "!(NaN >= 1.0) should be true")
|
||||
|
||||
def test_v_cmp_nge_f32_nan_src1(self):
|
||||
"""v_cmp_nge_f32: NaN in src1 should return true (x >= NaN is false)."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], f2i(1.0)),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= NaN) = !(false) = true
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= NaN) should be true")
|
||||
|
||||
def test_v_cmp_nge_f32_both_nan(self):
|
||||
"""v_cmp_nge_f32: both NaN should return true."""
|
||||
quiet_nan = 0x7fc00000
|
||||
instructions = [
|
||||
s_mov_b32(s[0], quiet_nan),
|
||||
v_mov_b32_e32(v[0], s[0]),
|
||||
v_mov_b32_e32(v[1], s[0]),
|
||||
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= NaN) = !(false) = true
|
||||
]
|
||||
st = run_program(instructions, n_lanes=1)
|
||||
self.assertEqual(st.vcc & 1, 1, "!(NaN >= NaN) should be true")
|
||||
|
||||
|
||||
class TestCmpxPartialWavefront(unittest.TestCase):
|
||||
"""Tests for V_CMPX with partial wavefronts (fewer than 32 active lanes).
|
||||
|
||||
|
||||
@@ -9,10 +9,13 @@ 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.emu2 import WaveState, decode_program, WAVE_SIZE, MASK32, PC_LO_IDX, PC_HI_IDX, SCC_IDX, VCC_LO, EXEC_LO
|
||||
from extra.assembly.amd.decode import decode_inst
|
||||
from extra.assembly.amd.test.helpers import KernelInfo
|
||||
from extra.assembly.amd.test.bench_emu import REMU_PATH
|
||||
|
||||
def set_valid_mem_ranges(ranges): pass # emu2 doesn't need this
|
||||
|
||||
def _is_f32_nan(bits: int) -> bool:
|
||||
"""Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
|
||||
return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
|
||||
@@ -91,33 +94,61 @@ class PythonEmulator:
|
||||
def __init__(self):
|
||||
self.state: WaveState | None = None
|
||||
self.program: dict | None = None
|
||||
self.vmem_buf = None
|
||||
self.lds_buf = None
|
||||
self.kernel_buf = None # Keep kernel bytes alive
|
||||
self.lib_addr = 0 # Base address of kernel code
|
||||
|
||||
def create(self, kernel: bytes, n_lanes: int):
|
||||
self.program = decode_program(kernel)
|
||||
self.state = WaveState(LDSMem(bytearray(65536)), n_lanes)
|
||||
self.state.exec_mask = (1 << n_lanes) - 1
|
||||
import ctypes
|
||||
from tinygrad.device import Buffer, BufferSpec
|
||||
from tinygrad.dtype import dtypes
|
||||
# Store kernel in a ctypes buffer so generic instructions can read from vmem at actual PC address
|
||||
self.kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
|
||||
self.lib_addr = ctypes.addressof(self.kernel_buf)
|
||||
# Remap program dict to use actual addresses (like run_asm does)
|
||||
program_raw = decode_program(kernel)
|
||||
self.program = {self.lib_addr + offset: val for offset, val in program_raw.items()}
|
||||
self.state = WaveState(n_lanes)
|
||||
self.state.pc = self.lib_addr # Set PC to code base address
|
||||
self.vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
|
||||
self.lds_buf = Buffer('CPU', 65536 // 4, dtypes.uint32).ensure_allocated()
|
||||
|
||||
def step(self) -> int:
|
||||
import ctypes
|
||||
assert self.program is not None and self.state is not None
|
||||
return self.program[self.state.pc]._dispatch(self.state, self.program[self.state.pc])
|
||||
pc = self.state.pc
|
||||
if pc == 0xFFFFFFFFFFFFFFFF or pc not in self.program: return -1
|
||||
name, fxn, globals_list, _runner = self.program[pc]
|
||||
if fxn is None: return 1 # unsupported instruction
|
||||
buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,
|
||||
2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}
|
||||
# Direct ctypes call - bypasses HCQ overhead
|
||||
fxn(*[ctypes.c_uint64(buf_addrs[g]) for g in globals_list], ctypes.c_int32(0))
|
||||
return -1 if self.state.pc == 0xFFFFFFFFFFFFFFFF else 0
|
||||
|
||||
def set_sgpr(self, idx: int, val: int):
|
||||
assert self.state is not None
|
||||
self.state.sgpr[idx] = val & 0xffffffff
|
||||
self.state._write_sgpr(idx, val)
|
||||
def set_vgpr(self, lane: int, idx: int, val: int):
|
||||
assert self.state is not None
|
||||
self.state.vgpr[lane][idx] = val & 0xffffffff
|
||||
self.state._write_vgpr(idx, lane, val)
|
||||
|
||||
def get_snapshot(self) -> StateSnapshot:
|
||||
assert self.state is not None
|
||||
return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
|
||||
exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
|
||||
vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
|
||||
sgpr = [self.state._read_sgpr(i) for i in range(128)]
|
||||
vgpr = [[self.state._read_vgpr(reg, lane) for reg in range(256)] for lane in range(WAVE_SIZE)]
|
||||
# Convert actual PC address to word offset for comparison with Rust emulator
|
||||
pc_offset = (self.state.pc - self.lib_addr) // 4 if self.state.pc != 0xFFFFFFFFFFFFFFFF else 0xFFFFFFFFFFFFFFFF
|
||||
return StateSnapshot(pc=pc_offset, scc=self.state._read_sgpr(SCC_IDX), vcc=sgpr[VCC_LO.offset],
|
||||
exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
|
||||
|
||||
def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
|
||||
program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
|
||||
max_workgroups: int = 8) -> tuple[bool, str, int]:
|
||||
local_size: tuple[int, int, int], program, max_steps: int, debug: bool, trace_len: int,
|
||||
kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, int]:
|
||||
"""Run a single kernel through both emulators. Returns (success, message, total_steps)."""
|
||||
gx, gy, gz = global_size
|
||||
lx, ly, lz = local_size
|
||||
total_steps = 0
|
||||
wg_count = 0
|
||||
|
||||
@@ -140,28 +171,52 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
|
||||
emu.set_sgpr(13, gidx)
|
||||
emu.set_sgpr(14, gidy)
|
||||
emu.set_sgpr(15, gidz)
|
||||
# Initialize v[0] with packed workitem IDs for each lane
|
||||
for lane in range(n_lanes):
|
||||
tid = lane
|
||||
z, y, x = tid // (lx * ly), (tid // lx) % ly, tid % lx
|
||||
emu.set_vgpr(lane, 0, (z << 20) | (y << 10) | x)
|
||||
|
||||
step = 0
|
||||
trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
|
||||
prev_sync_after = False # Track if previous instruction had known Rust bugs
|
||||
try:
|
||||
while step < max_steps:
|
||||
rust_before = rust.get_snapshot()
|
||||
python_before = python.get_snapshot()
|
||||
|
||||
inst = program.get(python_before.pc)
|
||||
inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
|
||||
inst_info = python.program.get(python.lib_addr + python_before.pc * 4) # Convert word offset to actual address
|
||||
inst_hex_name = inst_info[0] if inst_info else f"unknown at PC={python_before.pc}"
|
||||
# Decode the instruction to get mnemonic for sync_after checks
|
||||
try:
|
||||
# Format is mnemonic_hexbytes, e.g. v_exp_f32_e32_014b027e -> hex is 014b027e
|
||||
parts = inst_hex_name.rsplit('_', 1)
|
||||
inst_bytes_hex = parts[1] if len(parts) == 2 else ""
|
||||
inst_bytes = bytes.fromhex(inst_bytes_hex) if inst_bytes_hex else b''
|
||||
decoded = decode_inst(inst_bytes) if inst_bytes else None
|
||||
inst_mnemonic = repr(decoded).split('(')[0] if decoded else ""
|
||||
except:
|
||||
inst_mnemonic = ""
|
||||
# For generic instructions, use function name for sync_after check
|
||||
if not inst_mnemonic: inst_mnemonic = inst_hex_name
|
||||
inst_str = inst_hex_name
|
||||
trace.append((step, python_before.pc, inst_str, rust_before, python_before))
|
||||
if len(trace) > trace_len: trace.pop(0)
|
||||
|
||||
if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
|
||||
|
||||
# Instructions with known Rust emulator bugs - sync Python to Rust after execution
|
||||
# Instructions with known Rust emulator bugs or precision differences - sync Python to Rust after execution
|
||||
# v_div_scale/v_div_fixup: Rust has different VCC handling
|
||||
# v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
|
||||
# s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
|
||||
sync_after = any(x in inst_str for x in ('v_div_scale_f32', 'v_div_scale_f64', 'v_div_fixup_f32', 'v_div_fixup_f64',
|
||||
'v_cvt_f16_f32', 's_add_i32', 's_sub_i32'))
|
||||
diffs = rust_before.diff(python_before, n_lanes)
|
||||
# v_exp_f32/v_log_f32/v_ldexp_f32: precision differences in transcendental functions
|
||||
# s_delay_alu: Rust handles differently
|
||||
# v_add_co_ci_u32/v_sub_co_ci_u32/v_subrev_co_ci_u32: Rust preserves inactive VCC bits, but hardware clears all bits
|
||||
sync_after = any(x in inst_mnemonic.lower() for x in ('v_div_scale', 'v_div_fixup', 'v_cvt_f16_f32', 's_add_i32', 's_sub_i32',
|
||||
'v_exp_f32', 'v_log_f32', 'v_ldexp_f32', 's_delay_alu',
|
||||
'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'))
|
||||
# Skip comparison if previous instruction had known Rust bugs (states were synced but may still differ slightly)
|
||||
diffs = rust_before.diff(python_before, n_lanes) if not prev_sync_after else []
|
||||
if diffs:
|
||||
trace_lines = []
|
||||
for idx, (s, pc, d, rb, pb) in enumerate(trace):
|
||||
@@ -200,7 +255,12 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
|
||||
for lane in range(n_lanes):
|
||||
for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
|
||||
assert python.state is not None
|
||||
python.state.pc, python.state.scc, python.state.vcc, python.state.exec_mask = rust_after.pc, rust_after.scc, rust_after.vcc, rust_after.exec_mask
|
||||
# Convert Rust's word-based PC to Python's actual address
|
||||
python.state.pc = python.lib_addr + rust_after.pc * 4
|
||||
python.state._write_sgpr(SCC_IDX, rust_after.scc)
|
||||
python.state._write_sgpr(VCC_LO.offset, rust_after.vcc)
|
||||
python.state._write_sgpr(EXEC_LO.offset, rust_after.exec_mask)
|
||||
prev_sync_after = sync_after
|
||||
|
||||
if rust_result == -1:
|
||||
total_steps += step + 1
|
||||
@@ -254,7 +314,7 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
|
||||
|
||||
ok, msg, steps = run_single_kernel(
|
||||
kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
|
||||
program, max_steps, debug, trace_len, ki
|
||||
kernel.local_size, program, max_steps, debug, trace_len, ki
|
||||
)
|
||||
total_steps += steps
|
||||
if not ok:
|
||||
@@ -281,7 +341,8 @@ def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list,
|
||||
set_valid_mem_ranges(ranges)
|
||||
|
||||
program = decode_program(kernel)
|
||||
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
|
||||
# Legacy wrapper assumes local_size = (n_lanes, 1, 1)
|
||||
ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), program, max_steps, debug, trace_len)
|
||||
return ok, msg
|
||||
|
||||
def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
|
||||
@@ -387,6 +448,7 @@ class TestTinygradKernels(unittest.TestCase):
|
||||
from tinygrad import dtypes
|
||||
self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
|
||||
def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
|
||||
def test_exp(self): self._test_kernel(lambda T: T.empty(1024).exp())
|
||||
def test_cross_entropy(self):
|
||||
import numpy as np
|
||||
np.random.seed(0)
|
||||
@@ -398,5 +460,48 @@ 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_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()
|
||||
|
||||
@@ -0,0 +1,273 @@
|
||||
"""Tests for the pcode parser."""
|
||||
import unittest
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from extra.assembly.amd.emu2 import parse_pcode
|
||||
from extra.assembly.amd.expr_parser 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
|
||||
|
||||
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])
|
||||
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])
|
||||
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])
|
||||
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])
|
||||
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),
|
||||
}
|
||||
_, assigns = parse_pcode(pcode, srcs, lane=UOp.const(dtypes.uint32, 0))
|
||||
# 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,
|
||||
}
|
||||
_, assigns = parse_pcode(pcode, srcs, lane=UOp.const(dtypes.uint32, 0))
|
||||
# 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),
|
||||
}
|
||||
_, assigns = parse_pcode(pcode, srcs, lane=UOp.const(dtypes.uint32, 0))
|
||||
# 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),
|
||||
}
|
||||
_, assigns = parse_pcode(pcode, srcs, lane=UOp.const(dtypes.uint32, 0))
|
||||
_, 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)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,403 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for the RDNA3 pseudocode DSL."""
|
||||
import unittest
|
||||
from extra.assembly.amd.pcode import (Reg, TypedView, TypedView, MASK32, MASK64,
|
||||
_f32, _i32, _f16, _i16, f32_to_f16, isNAN, _bf16, _ibf16, bf16_to_f32, f32_to_bf16,
|
||||
BYTE_PERMUTE, v_sad_u8, v_msad_u8, _compile_pseudocode, _expr, compile_pseudocode)
|
||||
from extra.assembly.amd.test.helpers import ExecContext
|
||||
from extra.assembly.amd.autogen.rdna3.str_pcode import PCODE
|
||||
from extra.assembly.amd.autogen.rdna3.enum import VOP3SDOp, VOPCOp
|
||||
|
||||
# Compile pseudocode functions on demand for regression tests
|
||||
_VOP3SDOp_V_DIV_SCALE_F32 = compile_pseudocode('VOP3SDOp', 'V_DIV_SCALE_F32', PCODE[VOP3SDOp.V_DIV_SCALE_F32])
|
||||
_VOPCOp_V_CMP_CLASS_F32 = compile_pseudocode('VOPCOp', 'V_CMP_CLASS_F32', PCODE[VOPCOp.V_CMP_CLASS_F32_E32])
|
||||
|
||||
class TestReg(unittest.TestCase):
|
||||
def test_u32_read(self):
|
||||
r = Reg(0xDEADBEEF)
|
||||
self.assertEqual(int(r.u32), 0xDEADBEEF)
|
||||
|
||||
def test_u32_write(self):
|
||||
r = Reg(0)
|
||||
r.u32 = 0x12345678
|
||||
self.assertEqual(r._val, 0x12345678)
|
||||
|
||||
def test_f32_read(self):
|
||||
r = Reg(0x40400000) # 3.0f
|
||||
self.assertAlmostEqual(float(r.f32), 3.0)
|
||||
|
||||
def test_f32_write(self):
|
||||
r = Reg(0)
|
||||
r.f32 = 3.0
|
||||
self.assertEqual(r._val, 0x40400000)
|
||||
|
||||
def test_i32_signed(self):
|
||||
r = Reg(0xFFFFFFFF) # -1 as signed
|
||||
self.assertEqual(int(r.i32), -1)
|
||||
|
||||
def test_u64(self):
|
||||
r = Reg(0xDEADBEEFCAFEBABE)
|
||||
self.assertEqual(int(r.u64), 0xDEADBEEFCAFEBABE)
|
||||
|
||||
def test_f64(self):
|
||||
r = Reg(0x4008000000000000) # 3.0 as f64
|
||||
self.assertAlmostEqual(float(r.f64), 3.0)
|
||||
|
||||
class TestTypedView(unittest.TestCase):
|
||||
def test_bit_slice(self):
|
||||
r = Reg(0xDEADBEEF)
|
||||
# Slices return TypedView which supports .u32, .u16 etc (matching pseudocode like S1.u32[1:0].u32)
|
||||
self.assertEqual(r.u32[7:0].u32, 0xEF)
|
||||
self.assertEqual(r.u32[15:8].u32, 0xBE)
|
||||
self.assertEqual(r.u32[23:16].u32, 0xAD)
|
||||
self.assertEqual(r.u32[31:24].u32, 0xDE)
|
||||
# Also works with int() for arithmetic
|
||||
self.assertEqual(int(r.u32[7:0]), 0xEF)
|
||||
|
||||
def test_single_bit_read(self):
|
||||
r = Reg(0b11010101)
|
||||
self.assertEqual(r.u32[0], 1)
|
||||
self.assertEqual(r.u32[1], 0)
|
||||
self.assertEqual(r.u32[2], 1)
|
||||
self.assertEqual(r.u32[3], 0)
|
||||
|
||||
def test_single_bit_write(self):
|
||||
r = Reg(0)
|
||||
r.u32[5] = 1
|
||||
r.u32[3] = 1
|
||||
self.assertEqual(r._val, 0b00101000)
|
||||
|
||||
def test_nested_bit_access(self):
|
||||
# S0.u32[S1.u32[4:0]] - access bit at position from another register
|
||||
s0 = Reg(0b11010101)
|
||||
s1 = Reg(3)
|
||||
bit_pos = s1.u32[4:0] # TypedView, int value = 3
|
||||
bit_val = s0.u32[int(bit_pos)] # bit 3 of s0 = 0
|
||||
self.assertEqual(int(bit_pos), 3)
|
||||
self.assertEqual(bit_val, 0)
|
||||
|
||||
def test_arithmetic(self):
|
||||
r1 = Reg(0x40400000) # 3.0f
|
||||
r2 = Reg(0x40800000) # 4.0f
|
||||
result = r1.f32 + r2.f32
|
||||
self.assertAlmostEqual(result, 7.0)
|
||||
|
||||
def test_comparison(self):
|
||||
r1 = Reg(5)
|
||||
r2 = Reg(3)
|
||||
self.assertTrue(r1.u32 > r2.u32)
|
||||
self.assertFalse(r1.u32 < r2.u32)
|
||||
self.assertTrue(r1.u32 != r2.u32)
|
||||
|
||||
class TestTypedView(unittest.TestCase):
|
||||
def test_slice_read(self):
|
||||
r = Reg(0x56781234)
|
||||
self.assertEqual(r[15:0].u16, 0x1234)
|
||||
self.assertEqual(r[31:16].u16, 0x5678)
|
||||
|
||||
def test_slice_write(self):
|
||||
r = Reg(0)
|
||||
r[15:0].u16 = 0x1234
|
||||
r[31:16].u16 = 0x5678
|
||||
self.assertEqual(r._val, 0x56781234)
|
||||
|
||||
def test_slice_f16(self):
|
||||
r = Reg(0)
|
||||
r[15:0].f16 = 3.0
|
||||
self.assertAlmostEqual(_f16(r._val & 0xffff), 3.0, places=2)
|
||||
|
||||
class TestCompiler(unittest.TestCase):
|
||||
def test_ternary(self):
|
||||
result = _expr("a > b ? 1 : 0")
|
||||
self.assertIn("if", result)
|
||||
self.assertIn("else", result)
|
||||
|
||||
def test_type_prefix_strip(self):
|
||||
self.assertEqual(_expr("1'0U"), "0")
|
||||
self.assertEqual(_expr("32'1"), "1")
|
||||
self.assertEqual(_expr("16'0xFFFF"), "0xFFFF")
|
||||
|
||||
def test_suffix_strip(self):
|
||||
self.assertEqual(_expr("0ULL"), "0")
|
||||
self.assertEqual(_expr("1LL"), "1")
|
||||
self.assertEqual(_expr("5U"), "5")
|
||||
self.assertEqual(_expr("3.14F"), "3.14")
|
||||
|
||||
def test_boolean_ops(self):
|
||||
self.assertIn("and", _expr("a && b"))
|
||||
self.assertIn("or", _expr("a || b"))
|
||||
self.assertIn("!=", _expr("a <> b"))
|
||||
|
||||
def test_pack16(self):
|
||||
result = _expr("{ a, b }")
|
||||
self.assertIn("_pack", result)
|
||||
|
||||
def test_type_cast_strip(self):
|
||||
self.assertEqual(_expr("64'U(x)"), "(x)")
|
||||
self.assertEqual(_expr("32'I(y)"), "(y)")
|
||||
|
||||
class TestExecContext(unittest.TestCase):
|
||||
def test_float_add(self):
|
||||
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
|
||||
ctx.D0.f32 = ctx.S0.f32 + ctx.S1.f32
|
||||
self.assertAlmostEqual(_f32(ctx.D0._val), 7.0)
|
||||
|
||||
def test_float_mul(self):
|
||||
ctx = ExecContext(s0=0x40400000, s1=0x40800000) # 3.0f, 4.0f
|
||||
ctx.run("D0.f32 = S0.f32 * S1.f32")
|
||||
self.assertAlmostEqual(_f32(ctx.D0._val), 12.0)
|
||||
|
||||
def test_scc_comparison(self):
|
||||
ctx = ExecContext(s0=42, s1=42)
|
||||
ctx.run("SCC = S0.u32 == S1.u32")
|
||||
self.assertEqual(ctx.SCC._val, 1)
|
||||
|
||||
def test_scc_comparison_false(self):
|
||||
ctx = ExecContext(s0=42, s1=43)
|
||||
ctx.run("SCC = S0.u32 == S1.u32")
|
||||
self.assertEqual(ctx.SCC._val, 0)
|
||||
|
||||
def test_ternary(self):
|
||||
code = _compile_pseudocode("D0.u32 = S0.u32 > S1.u32 ? 1'1U : 1'0U")
|
||||
ctx = ExecContext(s0=5, s1=3)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 1)
|
||||
|
||||
def test_pack(self):
|
||||
code = _compile_pseudocode("D0 = { S1[15:0].u16, S0[15:0].u16 }")
|
||||
ctx = ExecContext(s0=0x1234, s1=0x5678)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 0x56781234)
|
||||
|
||||
def test_tmp_with_typed_access(self):
|
||||
code = _compile_pseudocode("""tmp = S0.u32 + S1.u32
|
||||
D0.u32 = tmp.u32""")
|
||||
ctx = ExecContext(s0=100, s1=200)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 300)
|
||||
|
||||
def test_s_add_u32_pattern(self):
|
||||
# Real pseudocode pattern from S_ADD_U32
|
||||
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
|
||||
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
|
||||
D0.u32 = tmp.u32""")
|
||||
# Test overflow case
|
||||
ctx = ExecContext(s0=0xFFFFFFFF, s1=0x00000001)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 0) # Wraps to 0
|
||||
self.assertEqual(ctx.SCC._val, 1) # Carry set
|
||||
|
||||
def test_s_add_u32_no_overflow(self):
|
||||
code = _compile_pseudocode("""tmp = 64'U(S0.u32) + 64'U(S1.u32)
|
||||
SCC = tmp >= 0x100000000ULL ? 1'1U : 1'0U
|
||||
D0.u32 = tmp.u32""")
|
||||
ctx = ExecContext(s0=100, s1=200)
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val, 300)
|
||||
self.assertEqual(ctx.SCC._val, 0) # No carry
|
||||
|
||||
def test_vcc_lane_read(self):
|
||||
ctx = ExecContext(vcc=0b1010, lane=1)
|
||||
# Lane 1 is set
|
||||
self.assertEqual(ctx.VCC.u64[1], 1)
|
||||
self.assertEqual(ctx.VCC.u64[2], 0)
|
||||
|
||||
def test_vcc_lane_write(self):
|
||||
ctx = ExecContext(vcc=0, lane=0)
|
||||
ctx.VCC.u64[3] = 1
|
||||
ctx.VCC.u64[1] = 1
|
||||
self.assertEqual(ctx.VCC._val, 0b1010)
|
||||
|
||||
def test_for_loop(self):
|
||||
# CTZ pattern - find first set bit
|
||||
code = _compile_pseudocode("""tmp = -1
|
||||
for i in 0 : 31 do
|
||||
if S0.u32[i] == 1 then
|
||||
tmp = i
|
||||
endif
|
||||
endfor
|
||||
D0.i32 = tmp""")
|
||||
ctx = ExecContext(s0=0b1000) # Bit 3 is set
|
||||
ctx.run(code)
|
||||
self.assertEqual(ctx.D0._val & MASK32, 3)
|
||||
|
||||
def test_result_dict(self):
|
||||
ctx = ExecContext(s0=5, s1=3)
|
||||
ctx.D0.u32 = 42
|
||||
ctx.SCC._val = 1
|
||||
result = ctx.result()
|
||||
self.assertEqual(result['d0'], 42)
|
||||
self.assertEqual(result['scc'], 1)
|
||||
|
||||
class TestPseudocodeRegressions(unittest.TestCase):
|
||||
"""Regression tests for pseudocode instruction emulation bugs."""
|
||||
|
||||
def test_v_div_scale_f32_vcc_always_returned(self):
|
||||
"""V_DIV_SCALE_F32 must always return VCC, even when VCC=0 (no scaling needed).
|
||||
Bug: when VCC._val == vcc (both 0), VCC wasn't returned, so VCC bits weren't written.
|
||||
This caused division to produce wrong results for multiple lanes."""
|
||||
# Normal case: 1.0 / 3.0, no scaling needed, VCC should be 0
|
||||
s0 = 0x3f800000 # 1.0
|
||||
s1 = 0x40400000 # 3.0
|
||||
s2 = 0x3f800000 # 1.0 (numerator)
|
||||
result = _VOP3SDOp_V_DIV_SCALE_F32(s0, s1, s2, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
# Must always have VCC in result
|
||||
self.assertIn('VCC', result, "V_DIV_SCALE_F32 must always return VCC")
|
||||
self.assertEqual(result['VCC'] & 1, 0, "VCC lane 0 should be 0 when no scaling needed")
|
||||
|
||||
def test_v_cmp_class_f32_detects_quiet_nan(self):
|
||||
"""V_CMP_CLASS_F32 must correctly identify quiet NaN vs signaling NaN.
|
||||
Bug: isQuietNAN and isSignalNAN both used math.isnan which can't distinguish them."""
|
||||
quiet_nan = 0x7fc00000 # quiet NaN: exponent=255, bit22=1
|
||||
signal_nan = 0x7f800001 # signaling NaN: exponent=255, bit22=0
|
||||
# Test quiet NaN detection (bit 1 in mask)
|
||||
s1_quiet = 0b0000000010 # bit 1 = quiet NaN
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 1, "Should detect quiet NaN with quiet NaN mask")
|
||||
# Test signaling NaN detection (bit 0 in mask)
|
||||
s1_signal = 0b0000000001 # bit 0 = signaling NaN
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 1, "Should detect signaling NaN with signaling NaN mask")
|
||||
# Test that quiet NaN doesn't match signaling NaN mask
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(quiet_nan, s1_signal, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 0, "Quiet NaN should not match signaling NaN mask")
|
||||
# Test that signaling NaN doesn't match quiet NaN mask
|
||||
result = _VOPCOp_V_CMP_CLASS_F32(signal_nan, s1_quiet, 0, 0, 0, 0, 0, 0xffffffff, 0, None)
|
||||
self.assertEqual(result['D0'] & 1, 0, "Signaling NaN should not match quiet NaN mask")
|
||||
|
||||
def testisNAN_with_typed_view(self):
|
||||
"""isNAN must work with TypedView objects, not just Python floats.
|
||||
Bug: isNAN checked isinstance(x, float) which returned False for TypedView."""
|
||||
nan_reg = Reg(0x7fc00000) # quiet NaN
|
||||
normal_reg = Reg(0x3f800000) # 1.0
|
||||
inf_reg = Reg(0x7f800000) # +inf
|
||||
self.assertTrue(isNAN(nan_reg.f32), "isNAN should return True for NaN TypedView")
|
||||
self.assertFalse(isNAN(normal_reg.f32), "isNAN should return False for normal TypedView")
|
||||
self.assertFalse(isNAN(inf_reg.f32), "isNAN should return False for inf TypedView")
|
||||
|
||||
class TestBF16(unittest.TestCase):
|
||||
"""Tests for BF16 (bfloat16) support."""
|
||||
|
||||
def test_bf16_conversion(self):
|
||||
"""Test bf16 <-> f32 conversion."""
|
||||
# bf16 is just the top 16 bits of f32
|
||||
# 1.0f = 0x3f800000, bf16 = 0x3f80
|
||||
self.assertAlmostEqual(_bf16(0x3f80), 1.0, places=2)
|
||||
self.assertEqual(_ibf16(1.0), 0x3f80)
|
||||
# 2.0f = 0x40000000, bf16 = 0x4000
|
||||
self.assertAlmostEqual(_bf16(0x4000), 2.0, places=2)
|
||||
self.assertEqual(_ibf16(2.0), 0x4000)
|
||||
# -1.0f = 0xbf800000, bf16 = 0xbf80
|
||||
self.assertAlmostEqual(_bf16(0xbf80), -1.0, places=2)
|
||||
self.assertEqual(_ibf16(-1.0), 0xbf80)
|
||||
|
||||
def test_bf16_special_values(self):
|
||||
"""Test bf16 special values (inf, nan)."""
|
||||
import math
|
||||
# +inf: f32 = 0x7f800000, bf16 = 0x7f80
|
||||
self.assertTrue(math.isinf(_bf16(0x7f80)))
|
||||
self.assertEqual(_ibf16(float('inf')), 0x7f80)
|
||||
# -inf: f32 = 0xff800000, bf16 = 0xff80
|
||||
self.assertTrue(math.isinf(_bf16(0xff80)))
|
||||
self.assertEqual(_ibf16(float('-inf')), 0xff80)
|
||||
# NaN: quiet NaN bf16 = 0x7fc0
|
||||
self.assertTrue(math.isnan(_bf16(0x7fc0)))
|
||||
self.assertEqual(_ibf16(float('nan')), 0x7fc0)
|
||||
|
||||
def test_bf16_register_property(self):
|
||||
"""Test Reg.bf16 property."""
|
||||
r = Reg(0)
|
||||
r.bf16 = 3.0 # 3.0f = 0x40400000, bf16 = 0x4040
|
||||
self.assertEqual(r._val & 0xffff, 0x4040)
|
||||
self.assertAlmostEqual(float(r.bf16), 3.0, places=1)
|
||||
|
||||
def test_bf16_slice_property(self):
|
||||
"""Test TypedView.bf16 property."""
|
||||
r = Reg(0x40404040) # Two bf16 3.0 values
|
||||
self.assertAlmostEqual(r[15:0].bf16, 3.0, places=1)
|
||||
self.assertAlmostEqual(r[31:16].bf16, 3.0, places=1)
|
||||
|
||||
class TestBytePermute(unittest.TestCase):
|
||||
"""Tests for BYTE_PERMUTE helper function (V_PERM_B32)."""
|
||||
|
||||
def test_byte_select_0_to_7(self):
|
||||
"""Test selecting bytes 0-7 from 64-bit data."""
|
||||
# data = {s0, s1} where s0 is bytes 0-3, s1 is bytes 4-7
|
||||
# Combined: 0x0706050403020100 (byte 0 = 0x00, byte 7 = 0x07)
|
||||
data = 0x0706050403020100
|
||||
for i in range(8):
|
||||
self.assertEqual(BYTE_PERMUTE(data, i), i, f"byte {i} should be {i}")
|
||||
|
||||
def test_sign_extend_bytes(self):
|
||||
"""Test sign extension selectors 8-11."""
|
||||
# sel 8: sign of byte 1 (bits 15:8)
|
||||
# sel 9: sign of byte 3 (bits 31:24)
|
||||
# sel 10: sign of byte 5 (bits 47:40)
|
||||
# sel 11: sign of byte 7 (bits 63:56)
|
||||
data = 0x8000800080008000 # All relevant bytes have sign bit set
|
||||
self.assertEqual(BYTE_PERMUTE(data, 8), 0xff)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 9), 0xff)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 10), 0xff)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 11), 0xff)
|
||||
data = 0x7f007f007f007f00 # No sign bits set
|
||||
self.assertEqual(BYTE_PERMUTE(data, 8), 0x00)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 9), 0x00)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 10), 0x00)
|
||||
self.assertEqual(BYTE_PERMUTE(data, 11), 0x00)
|
||||
|
||||
def test_constant_zero(self):
|
||||
"""Test selector 12 returns 0x00."""
|
||||
self.assertEqual(BYTE_PERMUTE(0xffffffffffffffff, 12), 0x00)
|
||||
|
||||
def test_constant_ff(self):
|
||||
"""Test selectors >= 13 return 0xFF."""
|
||||
for sel in [13, 14, 15, 255]:
|
||||
self.assertEqual(BYTE_PERMUTE(0, sel), 0xff, f"sel {sel} should be 0xff")
|
||||
|
||||
class TestSADHelpers(unittest.TestCase):
|
||||
"""Tests for V_SAD_U8 and V_MSAD_U8 helper functions."""
|
||||
|
||||
def test_v_sad_u8_basic(self):
|
||||
"""Test v_sad_u8 with simple values."""
|
||||
# s0 = 0x04030201, s1 = 0x04030201 -> diff = 0 for all bytes
|
||||
result = v_sad_u8(0x04030201, 0x04030201, 0)
|
||||
self.assertEqual(result, 0)
|
||||
# s0 = 0x05040302, s1 = 0x04030201 -> diff = 1+1+1+1 = 4
|
||||
result = v_sad_u8(0x05040302, 0x04030201, 0)
|
||||
self.assertEqual(result, 4)
|
||||
|
||||
def test_v_sad_u8_with_accumulator(self):
|
||||
"""Test v_sad_u8 with non-zero accumulator."""
|
||||
# s0 = 0x05040302, s1 = 0x04030201, s2 = 100 -> 4 + 100 = 104
|
||||
result = v_sad_u8(0x05040302, 0x04030201, 100)
|
||||
self.assertEqual(result, 104)
|
||||
|
||||
def test_v_sad_u8_large_diff(self):
|
||||
"""Test v_sad_u8 with maximum byte differences."""
|
||||
# s0 = 0xffffffff, s1 = 0x00000000 -> diff = 255*4 = 1020
|
||||
result = v_sad_u8(0xffffffff, 0x00000000, 0)
|
||||
self.assertEqual(result, 1020)
|
||||
|
||||
def test_v_msad_u8_basic(self):
|
||||
"""Test v_msad_u8 masks when reference byte is 0."""
|
||||
# s0 = 0x10101010, s1 = 0x00000000 -> all masked, result = 0
|
||||
result = v_msad_u8(0x10101010, 0x00000000, 0)
|
||||
self.assertEqual(result, 0)
|
||||
# s0 = 0x10101010, s1 = 0x01010101 -> diff = |0x10-0x01|*4 = 15*4 = 60
|
||||
result = v_msad_u8(0x10101010, 0x01010101, 0)
|
||||
self.assertEqual(result, 60)
|
||||
|
||||
def test_v_msad_u8_partial_mask(self):
|
||||
"""Test v_msad_u8 with partial masking."""
|
||||
# s0 = 0x10101010, s1 = 0x00010001 -> bytes 1 and 3 masked
|
||||
# diff = |0x10-0x01| + |0x10-0x01| = 15 + 15 = 30
|
||||
result = v_msad_u8(0x10101010, 0x00010001, 0)
|
||||
self.assertEqual(result, 30)
|
||||
|
||||
def test_v_msad_u8_with_accumulator(self):
|
||||
"""Test v_msad_u8 with non-zero accumulator."""
|
||||
result = v_msad_u8(0x10101010, 0x01010101, 50)
|
||||
self.assertEqual(result, 110) # 60 + 50
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -19,12 +19,11 @@ 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, 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)
|
||||
from extra.assembly.amd.emu2 import run_asm
|
||||
return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size)
|
||||
|
||||
def _try_dlopen_remu():
|
||||
# Use Python emulator only if PYTHON_REMU=1
|
||||
|
||||
@@ -285,10 +285,13 @@ pm_render = PatternMatcher([
|
||||
lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:])
|
||||
if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
|
||||
# Where after gated load becomes alt value
|
||||
# NOTE: if a is CAST and a.src[0].dtype == l.dtype, use a.src[0] to avoid roundtrip cast (e.g. uint->float->uint)
|
||||
(UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
|
||||
UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
|
||||
UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype))+
|
||||
l.src[2:]).cast(a.dtype)),
|
||||
(UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),),
|
||||
allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
|
||||
allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype
|
||||
else a.cast(l.dtype))+l.src[2:]).cast(a.dtype)),
|
||||
])
|
||||
|
||||
# *** Ops.REDUCE -> Ops.DEFINE_ACC ***
|
||||
|
||||
Reference in New Issue
Block a user