forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Prebuilt Release @ ab07000
This commit is contained in:
75
tinygrad_repo/tinygrad/renderer/amd/__init__.py
Normal file
75
tinygrad_repo/tinygrad/renderer/amd/__init__.py
Normal file
@@ -0,0 +1,75 @@
|
||||
# Instruction format detection and decoding
|
||||
from __future__ import annotations
|
||||
from tinygrad.renderer.amd.dsl import Inst, FixedBitField, EnumBitField
|
||||
|
||||
# SDWA/DPP variant detection: src0 field (bits 0-8) encodes the variant
|
||||
# 0xf9 (249) = SDWA, 0xfa (250) = DPP16 for CDNA (GFX9)
|
||||
_VARIANT_SRC0 = {"_SDWA_SDST": 0xf9, "_SDWA": 0xf9, "_DPP16": 0xfa}
|
||||
|
||||
def _matches(data: bytes, cls: type[Inst]) -> bool:
|
||||
"""Check if data matches all FixedBitFields and op is in allowed."""
|
||||
for _, field in cls._fields:
|
||||
dword_idx = field.lo // 32
|
||||
if len(data) < (dword_idx + 1) * 4: return False
|
||||
word = int.from_bytes(data[dword_idx*4:(dword_idx+1)*4], 'little')
|
||||
field_lo = field.lo % 32
|
||||
if isinstance(field, FixedBitField):
|
||||
if ((word >> field_lo) & field.mask) != field.default: return False
|
||||
if isinstance(field, EnumBitField) and field.allowed is not None:
|
||||
try: opcode = field.decode((word >> field_lo) & field.mask)
|
||||
except ValueError: return False # opcode not in enum
|
||||
if opcode not in field.allowed: return False
|
||||
# Check SDWA/DPP variant based on src0 field (bits 0-8) - only for variant classes
|
||||
name = cls.__name__
|
||||
word = int.from_bytes(data[:4], 'little')
|
||||
for suffix, expected_src0 in _VARIANT_SRC0.items():
|
||||
if name.endswith(suffix): return (word & 0x1ff) == expected_src0
|
||||
return True
|
||||
|
||||
# Lazy-load instruction format tables to avoid circular imports (ins.py imports dsl.py which is in this package)
|
||||
_FORMATS: dict[str, list[type[Inst]]] | None = None
|
||||
|
||||
def _load_formats() -> dict[str, list[type[Inst]]]:
|
||||
global _FORMATS
|
||||
if _FORMATS is not None: return _FORMATS
|
||||
from tinygrad.runtime.autogen.amd.rdna3.ins import (VOP1, VOP1_SDST, VOP1_DPP16, VOP1_LIT, VOP2, VOP2_DPP16, VOP2_LIT, VOP3, VOP3_SDST,
|
||||
VOP3SD, VOP3P, VOPC, VOPC_DPP16, VOPD, VINTERP, SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL,
|
||||
SCRATCH)
|
||||
from tinygrad.runtime.autogen.amd.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST, VOP1_DPP16 as R4_VOP1_DPP16,
|
||||
VOP1_LIT as R4_VOP1_LIT, VOP2 as R4_VOP2, VOP2_DPP16 as R4_VOP2_DPP16, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3,
|
||||
VOP3_SDST as R4_VOP3_SDST, VOP3SD as R4_VOP3SD, VOP3P as R4_VOP3P, VOPC as R4_VOPC, VOPC_DPP16 as R4_VOPC_DPP16,
|
||||
VOPD as R4_VOPD, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT,
|
||||
SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT, SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT,
|
||||
SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP,
|
||||
SMEM as R4_SMEM, DS as R4_DS, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
|
||||
from tinygrad.runtime.autogen.amd.cdna.ins import (VOP1 as C_VOP1, VOP1_SDWA as C_VOP1_SDWA, VOP1_DPP16 as C_VOP1_DPP16,
|
||||
VOP2 as C_VOP2, VOP2_LIT as C_VOP2_LIT, VOP2_SDWA as C_VOP2_SDWA, VOP2_DPP16 as C_VOP2_DPP16,
|
||||
VOPC as C_VOPC, VOPC_SDWA_SDST as C_VOPC_SDWA_SDST,
|
||||
VOP3 as C_VOP3, VOP3_SDST as C_VOP3_SDST, VOP3SD as C_VOP3SD, VOP3P as C_VOP3P, VOP3P_MFMA as C_VOP3P_MFMA, VOP3PX2 as C_VOP3PX2,
|
||||
SOP1 as C_SOP1, SOP2 as C_SOP2, SOPC as C_SOPC, SOPK as C_SOPK, SOPK_LIT as C_SOPK_LIT, SOPP as C_SOPP, SMEM as C_SMEM, DS as C_DS,
|
||||
FLAT as C_FLAT, GLOBAL as C_GLOBAL, SCRATCH as C_SCRATCH, MUBUF as C_MUBUF)
|
||||
# Order matters: more specific encodings first, catch-alls (SOP2, VOP2) last
|
||||
# Order: base before _LIT (base matches regular ops, _LIT catches lit-only ops excluded from base)
|
||||
_FORMATS = {
|
||||
"rdna3": [VOPD, VOP3P, VINTERP, VOP3SD, VOP3_SDST, VOP3, DS, GLOBAL, SCRATCH, FLAT, SMEM,
|
||||
SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPK, SOPK_LIT, SOPP, VOPC_DPP16, VOPC, VOP1_SDST, VOP1_DPP16, VOP1, VOP1_LIT,
|
||||
VOP2_DPP16, VOP2, VOP2_LIT],
|
||||
"rdna4": [R4_VOPD, R4_VOP3P, R4_VINTERP, R4_VOP3SD, R4_VOP3_SDST, R4_VOP3, R4_DS, R4_GLOBAL, R4_SCRATCH, R4_FLAT, R4_SMEM,
|
||||
R4_SOP1, R4_SOP1_LIT, R4_SOPC, R4_SOPC_LIT, R4_SOPP, R4_SOPK, R4_SOPK_LIT, R4_VOPC_DPP16, R4_VOPC, R4_VOP1_SDST,
|
||||
R4_VOP1_DPP16, R4_VOP1, R4_VOP1_LIT, R4_SOP2, R4_SOP2_LIT, R4_VOP2_DPP16, R4_VOP2, R4_VOP2_LIT],
|
||||
"cdna": [C_VOP3PX2, C_VOP3P_MFMA, C_VOP3P, C_VOP3SD, C_VOP3_SDST, C_VOP3, C_DS, C_GLOBAL, C_SCRATCH, C_FLAT, C_MUBUF, C_SMEM,
|
||||
C_SOP1, C_SOPC, C_SOPP, C_SOPK, C_SOPK_LIT, C_VOPC_SDWA_SDST, C_VOPC,
|
||||
C_VOP1_DPP16, C_VOP1_SDWA, C_VOP1, C_VOP2_DPP16, C_VOP2_SDWA, C_SOP2, C_VOP2, C_VOP2_LIT],
|
||||
}
|
||||
return _FORMATS
|
||||
|
||||
def detect_format(data: bytes, arch: str = "rdna3") -> type[Inst]:
|
||||
"""Detect instruction format from machine code bytes."""
|
||||
assert len(data) >= 4, f"need at least 4 bytes, got {len(data)}"
|
||||
for cls in _load_formats()[arch]:
|
||||
if _matches(data, cls): return cls
|
||||
raise ValueError(f"unknown {arch} format word={int.from_bytes(data[:4], 'little'):#010x}")
|
||||
|
||||
def decode_inst(data: bytes, arch: str = "rdna3") -> Inst:
|
||||
"""Decode machine code bytes into an instruction."""
|
||||
return detect_format(data, arch).from_bytes(data)
|
||||
456
tinygrad_repo/tinygrad/renderer/amd/dsl.py
Normal file
456
tinygrad_repo/tinygrad/renderer/amd/dsl.py
Normal file
@@ -0,0 +1,456 @@
|
||||
# dsl.py - clean DSL for AMD assembly
|
||||
from typing import Any
|
||||
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
# Registers - unified src encoding space (0-511)
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
|
||||
class Reg:
|
||||
# Register names vary by arch: RDNA has NULL@124/M0@125, CDNA has M0@124/reserved@125
|
||||
# RDNA4 has DPP8@233, CDNA has SDWA@249/DPP@250/VCCZ@251/EXECZ@252
|
||||
_NAMES = {102: "FLAT_SCRATCH_LO", 103: "FLAT_SCRATCH_HI", 104: "XNACK_MASK_LO", 105: "XNACK_MASK_HI",
|
||||
106: "VCC_LO", 107: "VCC_HI", 124: "NULL", 125: "M0", 126: "EXEC_LO", 127: "EXEC_HI",
|
||||
233: "DPP8", 234: "DPP8FI", 235: "SHARED_BASE", 236: "SHARED_LIMIT", 237: "PRIVATE_BASE", 238: "PRIVATE_LIMIT",
|
||||
240: "0.5", 241: "-0.5", 242: "1.0", 243: "-1.0", 244: "2.0", 245: "-2.0", 246: "4.0", 247: "-4.0",
|
||||
248: "INV_2PI", 249: "SDWA", 250: "DPP", 251: "VCCZ", 252: "EXECZ", 253: "SCC", 254: "SRC_LDS_DIRECT", 255: "LIT"}
|
||||
_PAIRS = {106: "VCC", 126: "EXEC"}
|
||||
|
||||
def __init__(self, offset: int = 0, sz: int = 512, *, neg: bool = False, abs_: bool = False, hi: bool = False):
|
||||
self.offset, self.sz = offset, sz
|
||||
self.neg, self.abs_, self.hi = neg, abs_, hi
|
||||
|
||||
def __hash__(self): return hash((self.offset, self.sz, self.neg, self.abs_, self.hi))
|
||||
def __getitem__(self, key):
|
||||
if isinstance(key, slice):
|
||||
start, stop = key.start or 0, key.stop or (self.sz - 1)
|
||||
if start < 0 or stop >= self.sz: raise RuntimeError(f"slice [{start}:{stop}] out of bounds for size {self.sz}")
|
||||
return Reg(self.offset + start, stop - start + 1)
|
||||
if key < 0 or key >= self.sz: raise RuntimeError(f"index {key} out of bounds for size {self.sz}")
|
||||
return Reg(self.offset + key, 1)
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, Reg):
|
||||
return (self.offset == other.offset and self.sz == other.sz and
|
||||
self.neg == other.neg and self.abs_ == other.abs_ and self.hi == other.hi)
|
||||
return NotImplemented
|
||||
def __add__(self, other):
|
||||
if isinstance(other, int): return Reg(self.offset + other, self.sz)
|
||||
return NotImplemented
|
||||
def __neg__(self) -> 'Reg': return Reg(self.offset, self.sz, neg=not self.neg, abs_=self.abs_, hi=self.hi)
|
||||
def __abs__(self) -> 'Reg': return Reg(self.offset, self.sz, neg=self.neg, abs_=True, hi=self.hi)
|
||||
@property
|
||||
def h(self) -> 'Reg': return Reg(self.offset, self.sz, neg=self.neg, abs_=self.abs_, hi=True)
|
||||
@property
|
||||
def l(self) -> 'Reg': return Reg(self.offset, self.sz, neg=self.neg, abs_=self.abs_, hi=False)
|
||||
def fmt(self, sz=None, parens=False, upper=False) -> str:
|
||||
o, sz = self.offset, sz or self.sz
|
||||
l, r = ("[", "]") if parens or sz > 1 else ("", "") # brackets for multi-reg or when parens=True
|
||||
if 256 <= o < 512:
|
||||
idx = o - 256
|
||||
base = f"v{l}{idx}{r}" if sz == 1 else f"v[{idx}:{idx + sz - 1}]"
|
||||
elif o < 106: base = f"s{l}{o}{r}" if sz == 1 else f"s[{o}:{o + sz - 1}]"
|
||||
elif sz == 2 and o in self._PAIRS: base = self._PAIRS[o] if upper else self._PAIRS[o].lower()
|
||||
elif o in self._NAMES: base = self._NAMES[o] if upper else self._NAMES[o].lower() # special regs (any sz)
|
||||
elif 108 <= o < 124:
|
||||
idx = o - 108
|
||||
base = f"ttmp{l}{idx}{r}" if sz == 1 else f"ttmp[{idx}:{idx + sz - 1}]"
|
||||
elif 128 <= o <= 192: base = str(o - 128) # inline int constants (0-64)
|
||||
elif 193 <= o <= 208: base = str(-(o - 192)) # inline negative int constants (-1 to -16)
|
||||
else: raise RuntimeError(f"unknown register: offset={o}, sz={sz}")
|
||||
if self.hi: base += ".h"
|
||||
if self.abs_: base = f"abs({base})" if upper else f"|{base}|"
|
||||
if self.neg: base = f"-{base}"
|
||||
return base
|
||||
def __repr__(self): return self.fmt(parens=True, upper=True)
|
||||
|
||||
# Full src encoding space
|
||||
src = Reg(0, 512)
|
||||
|
||||
# Slices for each region (inclusive end)
|
||||
s = src[0:105] # SGPR0-105
|
||||
VCC_LO = src[106]
|
||||
VCC_HI = src[107]
|
||||
VCC = src[106:107]
|
||||
ttmp = src[108:123] # TTMP0-15
|
||||
NULL = OFF = src[124]
|
||||
M0 = src[125]
|
||||
EXEC_LO = src[126]
|
||||
EXEC_HI = src[127]
|
||||
EXEC = src[126:127]
|
||||
# 128: 0, 129-192: integers 1-64, 193-208: integers -1 to -16
|
||||
# 240-248: float constants (0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0, -4.0, 1/(2*PI))
|
||||
INV_2PI = src[248]
|
||||
SDWA = src[249]
|
||||
DPP = DPP16 = src[250]
|
||||
VCCZ = src[251]
|
||||
EXECZ = src[252]
|
||||
SCC = src[253]
|
||||
SRC_LDS_DIRECT = src[254]
|
||||
LIT = src[255] # literal constant marker
|
||||
v = src[256:511] # VGPR0-255
|
||||
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
# BitField
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
|
||||
class _Bits:
|
||||
"""Helper for defining bit fields with slice syntax: bits[hi:lo] or bits[n]."""
|
||||
def __getitem__(self, key) -> 'BitField': return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
|
||||
bits = _Bits()
|
||||
|
||||
class BitField:
|
||||
name: str | None
|
||||
def __init__(self, hi: int, lo: int, default = 0):
|
||||
self.hi, self.lo, self.default, self.name, self.mask = hi, lo, default, None, (1 << (hi - lo + 1)) - 1
|
||||
def __set_name__(self, owner, name: str): self.name = name
|
||||
def __eq__(self, other) -> 'FixedBitField': # type: ignore[override]
|
||||
if isinstance(other, int): return FixedBitField(self.hi, self.lo, other)
|
||||
raise TypeError(f"BitField.__eq__ expects int, got {type(other).__name__}")
|
||||
def enum(self, enum_cls) -> 'EnumBitField': return EnumBitField(self.hi, self.lo, enum_cls)
|
||||
def encode(self, val) -> int:
|
||||
assert isinstance(val, int), f"BitField.encode expects int, got {type(val).__name__}"
|
||||
return val
|
||||
def decode(self, val): return val
|
||||
def set(self, raw: int, val) -> int:
|
||||
if val is None: val = self.default
|
||||
encoded = self.encode(val)
|
||||
# Handle signed values: convert negative to 2's complement
|
||||
if encoded < 0: encoded = encoded & self.mask
|
||||
if encoded < 0 or encoded > self.mask: raise RuntimeError(f"field '{self.name}': value {encoded} doesn't fit in {self.hi - self.lo + 1} bits")
|
||||
return (raw & ~(self.mask << self.lo)) | (encoded << self.lo)
|
||||
def __get__(self, obj, objtype=None):
|
||||
if obj is None: return self
|
||||
return self.decode((obj._raw >> self.lo) & self.mask)
|
||||
def __set__(self, obj, val): obj._raw = self.set(obj._raw, val)
|
||||
|
||||
class FixedBitField(BitField):
|
||||
def set(self, raw: int, val=None) -> int:
|
||||
assert val is None, f"FixedBitField does not accept values, got {val}"
|
||||
return super().set(raw, self.default)
|
||||
|
||||
class EnumBitField(BitField):
|
||||
def __init__(self, hi: int, lo: int, enum_cls, allowed: set | None = None):
|
||||
super().__init__(hi, lo)
|
||||
self._enum = enum_cls
|
||||
self.allowed = allowed # if set, only these enum values are valid for this encoding
|
||||
def encode(self, val) -> int:
|
||||
if not isinstance(val, self._enum): raise RuntimeError(f"expected {self._enum.__name__}, got {type(val).__name__}")
|
||||
if self.allowed is not None and val not in self.allowed:
|
||||
raise RuntimeError(f"opcode {val.name} not allowed in this encoding")
|
||||
return val.value
|
||||
def decode(self, raw): return self._enum(raw)
|
||||
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
# Typed fields
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
|
||||
import struct
|
||||
def _f32(f: float) -> int: return struct.unpack('I', struct.pack('f', f))[0]
|
||||
|
||||
class SrcField(BitField):
|
||||
_valid_range = (0, 511) # inclusive
|
||||
_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}
|
||||
|
||||
def __init__(self, hi: int, lo: int, default=s[0]):
|
||||
super().__init__(hi, lo, default)
|
||||
expected_size = self._valid_range[1] - self._valid_range[0] + 1
|
||||
actual_size = 1 << (hi - lo + 1)
|
||||
if actual_size != expected_size:
|
||||
raise RuntimeError(f"{self.__class__.__name__}: field size {hi - lo + 1} bits ({actual_size}) "
|
||||
f"doesn't match range {self._valid_range} ({expected_size})")
|
||||
|
||||
def encode(self, val) -> int:
|
||||
"""Encode value. Returns 255 (literal marker) for out-of-range values."""
|
||||
if isinstance(val, Reg): offset = val.offset
|
||||
elif isinstance(val, float): offset = self._FLOAT_ENC.get(val, 255)
|
||||
elif isinstance(val, int) and 0 <= val <= 64: offset = 128 + val
|
||||
elif isinstance(val, int) and -16 <= val < 0: offset = 192 - val
|
||||
elif isinstance(val, int): offset = 255 # literal
|
||||
else: raise TypeError(f"invalid src value {val}")
|
||||
if not (self._valid_range[0] <= offset <= self._valid_range[1]):
|
||||
raise TypeError(f"{self.__class__.__name__}: {val} (offset {offset}) out of range {self._valid_range}")
|
||||
return offset - self._valid_range[0]
|
||||
|
||||
def decode(self, raw): return src[raw + self._valid_range[0]]
|
||||
|
||||
def __get__(self, obj, objtype=None):
|
||||
if obj is None: return self
|
||||
reg = self.decode((obj._raw >> self.lo) & self.mask)
|
||||
# Resize register based on operand info (skip non-resizable special registers)
|
||||
# VCC/EXEC pairs (106, 126), NULL (124), M0 (125), float constants (240-255)
|
||||
if reg.offset not in (124, 125) and not 240 <= reg.offset <= 255:
|
||||
# Map variant field names (vsrc0->src0, vsrc1->src1, etc.) for DPP/SDWA classes
|
||||
assert self.name is not None
|
||||
name = self.name[1:] if self.name.startswith('v') and self.name[1:] in obj.op_regs else self.name
|
||||
if sz := obj.op_regs.get(name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
|
||||
return reg
|
||||
|
||||
class VGPRField(SrcField):
|
||||
_valid_range = (256, 511)
|
||||
def __init__(self, hi: int, lo: int, default=v[0]): super().__init__(hi, lo, default)
|
||||
def encode(self, val) -> int:
|
||||
if not isinstance(val, Reg): raise TypeError(f"VGPRField requires Reg, got {type(val).__name__}")
|
||||
# For 8-bit vdst fields in VOP1/VOP2 16-bit ops, bit 7 is opsel for dest half
|
||||
encoded = super().encode(val)
|
||||
if val.hi and (self.hi - self.lo + 1) == 8:
|
||||
if encoded >= 128:
|
||||
raise ValueError(f"VGPRField: v[{encoded}].h not encodable in 8-bit field (v[0:127] only for .h)")
|
||||
encoded |= 0x80
|
||||
return encoded
|
||||
class SGPRField(SrcField): _valid_range = (0, 127)
|
||||
class SSrcField(SrcField): _valid_range = (0, 255)
|
||||
|
||||
class AlignedSGPRField(BitField):
|
||||
"""SGPR field with alignment requirement. Encoded as sgpr_index // alignment."""
|
||||
_align: int = 2
|
||||
def encode(self, val):
|
||||
if isinstance(val, int) and val == 0: return 0 # default: encode as s[0]
|
||||
if not isinstance(val, Reg): raise TypeError(f"{self.__class__.__name__} requires Reg, got {type(val).__name__}")
|
||||
if not (0 <= val.offset < 128): raise ValueError(f"{self.__class__.__name__} requires SGPR, got offset {val.offset}")
|
||||
if val.offset & (self._align - 1): raise ValueError(f"{self.__class__.__name__} requires {self._align}-aligned SGPR, got s[{val.offset}]")
|
||||
return val.offset >> (self._align.bit_length() - 1)
|
||||
def decode(self, raw): return src[raw << (self._align.bit_length() - 1)]
|
||||
def __get__(self, obj, objtype=None):
|
||||
if obj is None: return self
|
||||
reg = self.decode((obj._raw >> self.lo) & self.mask)
|
||||
if sz := obj.op_regs.get(self.name, 1): reg = Reg(reg.offset, sz, neg=reg.neg, abs_=reg.abs_, hi=reg.hi)
|
||||
return reg
|
||||
|
||||
class SBaseField(AlignedSGPRField): _align = 2
|
||||
class SRsrcField(AlignedSGPRField): _align = 4
|
||||
|
||||
class VDSTYField(BitField):
|
||||
"""VOPD vdsty: encoded = vgpr_idx >> 1. Actual vgpr = (encoded << 1) | ((vdstx & 1) ^ 1)."""
|
||||
def encode(self, val):
|
||||
if not isinstance(val, Reg): raise TypeError(f"VDSTYField requires Reg, got {type(val).__name__}")
|
||||
if not (256 <= val.offset < 512): raise ValueError(f"VDSTYField requires VGPR, got offset {val.offset}")
|
||||
return (val.offset - 256) >> 1
|
||||
def __get__(self, obj, objtype=None):
|
||||
if obj is None: return self
|
||||
raw = (obj._raw >> self.lo) & self.mask
|
||||
vdstx_bit0 = (obj.vdstx.offset - 256) & 1
|
||||
vgpr_idx = (raw << 1) | (vdstx_bit0 ^ 1)
|
||||
return Reg(256 + vgpr_idx, 1)
|
||||
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
# Operand info from XML
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
|
||||
import functools
|
||||
from tinygrad.runtime.autogen.amd.rdna3.operands import OPERANDS as OPERANDS_RDNA3
|
||||
from tinygrad.runtime.autogen.amd.rdna4.operands import OPERANDS as OPERANDS_RDNA4
|
||||
from tinygrad.runtime.autogen.amd.cdna.operands import OPERANDS as OPERANDS_CDNA
|
||||
OPERANDS = {**OPERANDS_CDNA, **OPERANDS_RDNA3, **OPERANDS_RDNA4}
|
||||
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
# Inst base class
|
||||
# ══════════════════════════════════════════════════════════════
|
||||
|
||||
def _needs_literal(val) -> bool:
|
||||
"""Check if a value needs a literal constant (can't be encoded inline)."""
|
||||
if val is None or isinstance(val, Reg): return False
|
||||
if isinstance(val, float): return val not in SrcField._FLOAT_ENC
|
||||
if isinstance(val, int): return not (0 <= val <= 64 or -16 <= val < 0)
|
||||
return False
|
||||
|
||||
def _get_variant(cls, suffix: str):
|
||||
"""Get a variant class by suffix (e.g., '_LIT') via module lookup."""
|
||||
import sys
|
||||
module = sys.modules.get(cls.__module__)
|
||||
return getattr(module, f"{cls.__name__}{suffix}", None) if module else None
|
||||
|
||||
def _canonical_name(name: str) -> str | None:
|
||||
"""Map operand name to canonical name."""
|
||||
if name in ('src0', 'vsrc0', 'ssrc0'): return 's0'
|
||||
if name in ('src1', 'vsrc1', 'ssrc1'): return 's1'
|
||||
if name == 'src2': return 's2'
|
||||
if name in ('vdst', 'sdst', 'sdata'): return 'd'
|
||||
if name in ('data', 'vdata', 'data0', 'vsrc'): return 'data'
|
||||
return None
|
||||
|
||||
class Inst:
|
||||
_fields: list[tuple[str, BitField]]
|
||||
_base_size: int
|
||||
|
||||
def __init_subclass__(cls):
|
||||
# Collect fields from all parent classes, then override with this class's fields
|
||||
inherited = {}
|
||||
for base in reversed(cls.__mro__[1:]):
|
||||
if hasattr(base, '_fields'):
|
||||
inherited.update(dict(base._fields))
|
||||
inherited.update({name: val for name, val in cls.__dict__.items() if isinstance(val, BitField)})
|
||||
cls._fields = list(inherited.items())
|
||||
cls._base_size = (max(f.hi for _, f in cls._fields) + 8) // 8
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
# Auto-upgrade to variant if needed (only for base classes, not variants)
|
||||
if not any(cls.__name__.endswith(sfx) for sfx in ('_LIT', '_DPP16', '_DPP8', '_SDWA', '_SDWA_SDST', '_MFMA')):
|
||||
args_iter = iter(args)
|
||||
for name, field in cls._fields:
|
||||
if isinstance(field, FixedBitField): continue
|
||||
val = kwargs.get(name) if name in kwargs else next(args_iter, None)
|
||||
if not isinstance(field, SrcField): continue
|
||||
if isinstance(val, Reg) and val.offset == 255 and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
|
||||
if isinstance(val, Reg) and val.offset == 249:
|
||||
if (sdwa_cls := _get_variant(cls, '_SDWA') or _get_variant(cls, '_SDWA_SDST')): return sdwa_cls(*args, **kwargs)
|
||||
if isinstance(val, Reg) and val.offset == 250 and (dpp_cls := _get_variant(cls, '_DPP16')): return dpp_cls(*args, **kwargs)
|
||||
if _needs_literal(val) and (lit_cls := _get_variant(cls, '_LIT')): return lit_cls(*args, **kwargs)
|
||||
return object.__new__(cls)
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self._raw = 0
|
||||
# Map positional args to field names (skip FixedBitFields)
|
||||
args_iter = iter(args)
|
||||
vals: dict[str, Any] = {}
|
||||
for name, field in self._fields:
|
||||
if isinstance(field, FixedBitField): vals[name] = None
|
||||
elif name in kwargs: vals[name] = kwargs[name]
|
||||
else: vals[name] = next(args_iter, None)
|
||||
assert not (remaining := list(args_iter)), f"too many positional args: {remaining}"
|
||||
known_field_names = [name for name,field in self._fields if not isinstance(field, FixedBitField)]
|
||||
for name in kwargs:
|
||||
if name not in known_field_names: raise TypeError(f"{self.__class__.__name__}() got an unexpected keyword argument {name!r}")
|
||||
# Extract modifiers from Reg objects and merge into neg/abs/opsel
|
||||
neg_bits, abs_bits, opsel_bits = 0, 0, 0
|
||||
for name, bit in [('src0', 0), ('src1', 1), ('src2', 2)]:
|
||||
if name in vals and isinstance(vals[name], Reg):
|
||||
reg = vals[name]
|
||||
if reg.neg: neg_bits |= (1 << bit)
|
||||
if reg.abs_: abs_bits |= (1 << bit)
|
||||
if reg.hi: opsel_bits |= (1 << bit)
|
||||
if 'vdst' in vals and isinstance(vals['vdst'], Reg) and vals['vdst'].hi:
|
||||
opsel_bits |= (1 << 3)
|
||||
if neg_bits: vals['neg'] = (vals.get('neg') or 0) | neg_bits
|
||||
if abs_bits: vals['abs'] = (vals.get('abs') or 0) | abs_bits
|
||||
if opsel_bits: vals['opsel'] = (vals.get('opsel') or 0) | opsel_bits
|
||||
# For _LIT classes, capture literal value from SrcFields that encode to 255
|
||||
literal_val = None
|
||||
for name, field in self._fields:
|
||||
val = vals[name]
|
||||
if isinstance(field, SrcField) and val is not None and _needs_literal(val):
|
||||
literal_val = _f32(val) if isinstance(val, float) else val & 0xFFFFFFFF
|
||||
if literal_val is not None and 'literal' in vals:
|
||||
vals['literal'] = literal_val
|
||||
# Set all field values
|
||||
for name, field in self._fields:
|
||||
self._raw = field.set(self._raw, vals[name])
|
||||
# Validate register sizes against operand info (skip special registers like NULL, VCC, EXEC, SDWA/DPP markers)
|
||||
for name, expected in self.op_regs.items():
|
||||
if (val := vals.get(name)) is None: continue
|
||||
if isinstance(val, Reg) and val.sz != expected and not (106 <= val.offset <= 127 or 249 <= val.offset <= 255):
|
||||
raise TypeError(f"{name} expects {expected} register(s), got {val.sz}")
|
||||
|
||||
@property
|
||||
def op_name(self) -> str: return getattr(self, 'op').name
|
||||
@property
|
||||
def operands(self) -> dict: return OPERANDS.get(getattr(self, 'op'), {}) if hasattr(self, 'op') else {}
|
||||
def _is_cdna(self) -> bool: return 'cdna' in type(self).__module__
|
||||
|
||||
@functools.cached_property
|
||||
def op_bits(self) -> dict[str, int]:
|
||||
"""Get bit widths for each operand field, with WAVE32 and addr/saddr adjustments."""
|
||||
if not hasattr(self, 'op'): return {k: v[1] for k, v in self.operands.items()}
|
||||
bits = {k: v[1] for k, v in self.operands.items()}
|
||||
# RDNA (WAVE32): condition masks, carry flags, and compare results are 32-bit
|
||||
if not self._is_cdna():
|
||||
name = self.op_name.lower()
|
||||
if 'cndmask' in name and 'src2' in bits: bits['src2'] = 32
|
||||
if '_co_ci_' in name and 'src2' in bits: bits['src2'] = 32 # carry-in source
|
||||
# VOP3SD: sdst is always wavefront-size dependent (carry-out or condition mask)
|
||||
if 'VOP3SD' in type(self).__name__ and 'sdst' in bits: bits['sdst'] = 32
|
||||
if 'cmp' in name and 'vdst' in bits: bits['vdst'] = 32
|
||||
# GLOBAL/FLAT: addr is 32-bit if saddr is valid SGPR, 64-bit if saddr is NULL
|
||||
# SCRATCH: addr is always 32-bit (offset from scratch base, not absolute address)
|
||||
if 'addr' in bits and (saddr_field := getattr(type(self), 'saddr', None)) and type(self).__name__ not in ('SCRATCH', 'VSCRATCH'):
|
||||
saddr_val = (self._raw >> saddr_field.lo) & saddr_field.mask # access _raw directly to avoid recursion
|
||||
bits['addr'] = 64 if saddr_val in (124, 125) else 32 # 124=NULL, 125=M0
|
||||
# MUBUF/MTBUF: vaddr size depends on offen/idxen (1 or 2 regs)
|
||||
if 'vaddr' in bits and hasattr(self, 'offen') and hasattr(self, 'idxen'):
|
||||
bits['vaddr'] = max(1, self.offen + self.idxen) * 32
|
||||
# F8F6F4 MFMA: CBSZ selects matrix A format, BLGP selects matrix B format
|
||||
# VGPRs: FP8/BF8(0,1)=8, FP6/BF6(2,3)=6, FP4(4)=4
|
||||
if 'f8f6f4' in getattr(self, 'op_name', '').lower():
|
||||
# Use explicit fields if available (VOP3PX2), else extract from VOP3P-MAI bit positions
|
||||
cbsz = getattr(self, 'cbsz') if hasattr(type(self), 'cbsz') else (self._raw >> 8) & 0x7
|
||||
blgp = getattr(self, 'blgp') if hasattr(type(self), 'blgp') else (self._raw >> 61) & 0x7
|
||||
vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
|
||||
bits['src0'], bits['src1'] = vgprs.get(cbsz, 8) * 32, vgprs.get(blgp, 8) * 32
|
||||
return bits
|
||||
@property
|
||||
def op_regs(self) -> dict[str, int]:
|
||||
"""Get register counts for each operand field."""
|
||||
return {k: max(1, v // 32) for k, v in self.op_bits.items()}
|
||||
|
||||
@functools.cached_property
|
||||
def canonical_op_bits(self) -> dict[str, int]:
|
||||
"""Get bit widths with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
|
||||
bits = {'d': 32, 's0': 32, 's1': 32, 's2': 32, 'data': 32}
|
||||
for name, val in self.op_bits.items():
|
||||
if (cn := _canonical_name(name)): bits[cn] = val
|
||||
return bits
|
||||
|
||||
@functools.cached_property
|
||||
def canonical_operands(self) -> dict:
|
||||
"""Get operands with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
|
||||
result = {}
|
||||
for name, val in self.operands.items():
|
||||
if (cn := _canonical_name(name)): result[cn] = val
|
||||
return result
|
||||
|
||||
@property
|
||||
def canonical_op_regs(self) -> dict[str, int]:
|
||||
"""Get register counts with canonical names: {'s0', 's1', 's2', 'd', 'data'}."""
|
||||
return {k: max(1, v // 32) for k, v in self.canonical_op_bits.items()}
|
||||
|
||||
def num_srcs(self) -> int:
|
||||
"""Get number of source operands from operand info."""
|
||||
ops = self.operands
|
||||
if 'src2' in ops: return 3
|
||||
if 'src1' in ops or 'vsrc1' in ops or 'ssrc1' in ops: return 2
|
||||
if 'src0' in ops or 'vsrc0' in ops or 'ssrc0' in ops: return 1
|
||||
return 0
|
||||
@classmethod
|
||||
def _size(cls) -> int: return cls._base_size
|
||||
def size(self) -> int: return self._base_size
|
||||
def disasm(self) -> str: raise NotImplementedError("disasm is no longer supported")
|
||||
|
||||
def to_bytes(self) -> bytes: return self._raw.to_bytes(self._base_size, 'little')
|
||||
|
||||
@property
|
||||
def _literal(self) -> int | None:
|
||||
"""Get the literal value if this instruction has one."""
|
||||
return getattr(self, 'literal', None)
|
||||
|
||||
def _variant_suffix(self) -> str | None:
|
||||
"""Check if instruction needs a variant class (_LIT, _DPP8, _DPP16, _SDWA). Returns suffix or None."""
|
||||
cls_name = type(self).__name__
|
||||
# Don't check for variants if we're already a variant class
|
||||
if any(s in cls_name for s in ('_LIT', '_DPP8', '_DPP16', '_SDWA')): return None
|
||||
# VOPD: FMAMK/FMAAK opcodes always require literal (check by name since enum may differ across archs)
|
||||
for name in ('opx', 'opy'):
|
||||
if hasattr(self, name) and any(x in getattr(self, name).name for x in ('FMAMK', 'FMAAK')): return '_LIT'
|
||||
for name, field in self._fields:
|
||||
if isinstance(field, SrcField):
|
||||
off = getattr(self, name).offset
|
||||
if off == 255: return '_LIT'
|
||||
if off == 249: return '_SDWA' if self._is_cdna() else '_DPP8'
|
||||
if off == 250: return '_DPP16'
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, data: bytes):
|
||||
inst = object.__new__(cls)
|
||||
inst._raw = int.from_bytes(data[:cls._base_size], 'little')
|
||||
# Upgrade to variant class if needed (_LIT, _DPP8, _DPP16, _SDWA)
|
||||
if (suffix := inst._variant_suffix()) and (var_cls := _get_variant(cls, suffix)) is not None:
|
||||
return var_cls.from_bytes(data)
|
||||
return inst
|
||||
|
||||
def __eq__(self, other): return type(self) is type(other) and self._raw == other._raw
|
||||
def __hash__(self): return hash((type(self), self._raw))
|
||||
|
||||
def __repr__(self):
|
||||
# collect (repr, is_default) pairs, strip trailing defaults so repr roundtrips with eval
|
||||
name = self.op.name.lower() if hasattr(self, 'op') else type(self).__name__
|
||||
parts = [(repr(v := getattr(self, n)), v == f.default) for n, f in self._fields if n != 'op' and not isinstance(f, FixedBitField)]
|
||||
while parts and parts[-1][1]: parts.pop()
|
||||
return f"{name}({', '.join(p[0] for p in parts)})"
|
||||
111
tinygrad_repo/tinygrad/renderer/amd/elf.py
Normal file
111
tinygrad_repo/tinygrad/renderer/amd/elf.py
Normal file
@@ -0,0 +1,111 @@
|
||||
# minimal amdgpu elf packer
|
||||
import ctypes
|
||||
from tinygrad.helpers import ceildiv, round_up
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.runtime.autogen import amdgpu_kd, hsa, libc
|
||||
from tinygrad.renderer.amd.dsl import Reg, FixedBitField
|
||||
from tinygrad.runtime.autogen.amd.common import OpType
|
||||
|
||||
# instructions used for padding
|
||||
from tinygrad.runtime.autogen.amd.rdna3.ins import s_code_end # same encoding as RDNA4
|
||||
from tinygrad.runtime.autogen.amd.cdna.ins import s_nop as s_nop_cdna
|
||||
|
||||
_arch_map = {"gfx9": "cdna", "gfx10": "rdna3", "gfx11": "rdna3", "gfx12": "rdna4"}
|
||||
def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes:
|
||||
insts = [u.arg for u in lin.src]
|
||||
|
||||
# ** scan for max vgpr/sgpr/accvgpr
|
||||
max_vgpr, max_sgpr, max_accvgpr = 0, 0, 0
|
||||
_ACCVGPR_TYPES = {OpType.OPR_ACCVGPR, OpType.OPR_SRC_ACCVGPR}
|
||||
for inst in insts:
|
||||
# build set of field names that are AccVGPR for this instruction
|
||||
accvgpr_fields: set[str] = set()
|
||||
for opr_name, (_, _, opr_type) in inst.operands.items():
|
||||
if opr_type in _ACCVGPR_TYPES: accvgpr_fields.add(opr_name)
|
||||
elif opr_type in {OpType.OPR_VGPR_OR_ACCVGPR, OpType.OPR_SRC_VGPR_OR_ACCVGPR, OpType.OPR_SRC_VGPR_OR_ACCVGPR_OR_CONST}:
|
||||
if getattr(inst, 'acc_cd', 0) == 1: accvgpr_fields.add(opr_name)
|
||||
for name, field in inst._fields:
|
||||
if isinstance(field, FixedBitField): continue
|
||||
val = getattr(inst, name)
|
||||
if not isinstance(val, Reg): continue
|
||||
if 256 <= val.offset < 512:
|
||||
if name in accvgpr_fields: max_accvgpr = max(max_accvgpr, (val.offset - 256) + val.sz)
|
||||
else: max_vgpr = max(max_vgpr, (val.offset - 256) + val.sz)
|
||||
elif val.offset < 106: max_sgpr = max(max_sgpr, val.offset + val.sz)
|
||||
|
||||
# ** scan sink for metadata
|
||||
sink, n_bufs, n_vars, lds_size, gids = prg.src[0], 0, 0, 0, set()
|
||||
for u in sink.toposort():
|
||||
if u.op is Ops.PARAM: n_bufs += 1
|
||||
elif u.op is Ops.DEFINE_VAR: n_vars += 1
|
||||
elif u.op is Ops.DEFINE_LOCAL: lds_size += u.ptrdtype.size * u.ptrdtype.base.itemsize
|
||||
elif u.op is Ops.SPECIAL and u.arg.startswith("gidx"): gids.add(int(u.arg[-1]))
|
||||
code_bytes = b"".join(inst.to_bytes() for inst in insts)
|
||||
arch = next(v for k, v in _arch_map.items() if arch.startswith(k))
|
||||
is_cdna, is_rdna4 = arch == "cdna", arch == "rdna4"
|
||||
|
||||
# ** pad text to ISA alignment
|
||||
padding_inst = (s_nop_cdna(0) if is_cdna else s_code_end()).to_bytes()
|
||||
text = code_bytes + padding_inst * ((hsa.AMD_ISA_ALIGN_BYTES - len(code_bytes) % hsa.AMD_ISA_ALIGN_BYTES) % hsa.AMD_ISA_ALIGN_BYTES)
|
||||
text_offset = round_up(ctypes.sizeof(libc.Elf64_Ehdr), hsa.AMD_ISA_ALIGN_BYTES)
|
||||
|
||||
# ** pack kernel descriptor (rodata)
|
||||
# CDNA: total VGPRs = regular VGPRs + AccVGPRs, each rounded to granularity of 4
|
||||
accum_offset = round_up(max_vgpr, 4) if max_accvgpr > 0 else 0
|
||||
next_free_vgpr = round_up(accum_offset + max_accvgpr, 8) if max_accvgpr > 0 else round_up(max_vgpr, 8)
|
||||
next_free_sgpr = round_up(max_sgpr, 8)
|
||||
vgpr_granule = max(0, (next_free_vgpr + 7) // 8 - 1)
|
||||
# CDNA: add 6 for VCC(2) + FLAT_SCRATCH(2) + XNACK_MASK(2), next_free_sgpr is unused in RDNA.
|
||||
sgpr_granule = max(0, ceildiv(next_free_sgpr + 6, 8) - 1) if is_cdna else 0
|
||||
desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t()
|
||||
desc.group_segment_fixed_size = lds_size
|
||||
desc.kernarg_size = n_bufs * 8 + n_vars * 4
|
||||
desc.kernel_code_entry_byte_offset = -len(text)
|
||||
|
||||
# https://llvm.org/docs/AMDGPUUsage.html#amdgpu-amdhsa-compute-pgm-rsrc1-gfx6-gfx12-table
|
||||
# NOTE: CU mode is the default
|
||||
desc.compute_pgm_rsrc1 = (vgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT_SHIFT |
|
||||
sgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT_SHIFT |
|
||||
3 << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_16_64_SHIFT |
|
||||
(0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_DX10_CLAMP_SHIFT |
|
||||
(0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_IEEE_MODE_SHIFT |
|
||||
(0 if is_cdna else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX10_PLUS_MEM_ORDERED_SHIFT)
|
||||
desc.compute_pgm_rsrc2 = (2 << amdgpu_kd.COMPUTE_PGM_RSRC2_USER_SGPR_COUNT_SHIFT |
|
||||
int(0 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X_SHIFT |
|
||||
int(1 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y_SHIFT |
|
||||
int(2 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Z_SHIFT)
|
||||
desc.kernel_code_properties = (1 << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_KERNARG_SEGMENT_PTR_SHIFT |
|
||||
(0 if is_cdna else 1) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_WAVEFRONT_SIZE32_SHIFT)
|
||||
if is_cdna and max_accvgpr > 0:
|
||||
desc.compute_pgm_rsrc3 = max(0, accum_offset // 4 - 1) << amdgpu_kd.COMPUTE_PGM_RSRC3_GFX90A_ACCUM_OFFSET_SHIFT
|
||||
rodata = bytes(desc)
|
||||
|
||||
# ** pack ELF
|
||||
sh_names:list[int] = []
|
||||
strtab = bytearray(b"\x00")
|
||||
for name in [".text", ".rodata", ".strtab"]:
|
||||
sh_names.append(len(strtab))
|
||||
strtab += name.encode("ascii") + b"\x00"
|
||||
|
||||
rodata_offset = round_up(text_offset + (text_size := len(text)), hsa.AMD_KERNEL_CODE_ALIGN_BYTES)
|
||||
strtab_offset = rodata_offset + (rodata_size := len(rodata))
|
||||
shdr_offset = strtab_offset + (strtab_size := len(strtab))
|
||||
|
||||
sections = [(libc.SHT_PROGBITS, libc.SHF_ALLOC | libc.SHF_EXECINSTR, text_offset, text_offset, text_size),
|
||||
(libc.SHT_PROGBITS, libc.SHF_ALLOC, rodata_offset, rodata_offset, rodata_size),
|
||||
(libc.SHT_STRTAB, 0, 0, strtab_offset, strtab_size)]
|
||||
shdrs = (libc.Elf64_Shdr * len(sections))()
|
||||
for i, s in enumerate(sections): shdrs[i] = libc.Elf64_Shdr(sh_names[i], *s)
|
||||
|
||||
ehdr = libc.Elf64_Ehdr()
|
||||
ehdr.e_ident[:5], ehdr.e_shoff, ehdr.e_shnum, ehdr.e_shstrndx = b"\x7FELF\x02", shdr_offset, len(sections), 2
|
||||
|
||||
elf = bytearray(shdr_offset + ctypes.sizeof(shdrs))
|
||||
elf[0:ctypes.sizeof(ehdr)] = bytes(ehdr)
|
||||
elf[text_offset:text_offset+text_size] = text
|
||||
elf[rodata_offset:rodata_offset+rodata_size] = rodata
|
||||
elf[strtab_offset:strtab_offset+strtab_size] = strtab
|
||||
elf[shdr_offset:shdr_offset+ctypes.sizeof(shdrs)] = bytes(shdrs)
|
||||
binary = bytes(elf)
|
||||
|
||||
return binary
|
||||
531
tinygrad_repo/tinygrad/renderer/amd/generate.py
Normal file
531
tinygrad_repo/tinygrad/renderer/amd/generate.py
Normal file
@@ -0,0 +1,531 @@
|
||||
# AMD ISA code generator - generates enum.py, ins.py, operands.py, str_pcode.py
|
||||
# Sources: XML from https://gpuopen.com/download/machine-readable-isa/latest/
|
||||
# PDF manuals from AMD documentation
|
||||
import re, zlib, xml.etree.ElementTree as ET, zipfile, pathlib
|
||||
from tinygrad.helpers import fetch
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Configuration
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
ARCHS = {
|
||||
"rdna3": {"xml": "amdgpu_isa_rdna3_5.xml", "pdf": "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"},
|
||||
"rdna4": {"xml": "amdgpu_isa_rdna4.xml", "pdf": "https://docs.amd.com/api/khub/documents/uQpkEvk3pv~kfAb2x~j4uw/content"},
|
||||
"cdna": {"xml": "amdgpu_isa_cdna4.xml", "pdf": "https://www.amd.com/content/dam/amd/en/documents/instinct-tech-docs/instruction-set-architectures/amd-instinct-cdna4-instruction-set-architecture.pdf"},
|
||||
}
|
||||
# Pin the September 2025 XML bundle because newer `latest` changed WMMA format bit sizes across archs and breaks generation.
|
||||
XML_URL = "https://gpuopen.com/download/AMD_GPU_MR_ISA_XML_2025_09_05.zip"
|
||||
# Map XML encoding names to codebase names
|
||||
NAME_MAP = {"VOP3_SDST_ENC": "VOP3SD", "VOP3_SDST_ENC_LIT": "VOP3SD_LIT", "VOP3_SDST_ENC_DPP16": "VOP3SD_DPP16",
|
||||
"VOP3_SDST_ENC_DPP8": "VOP3SD_DPP8", "VOPDXY": "VOPD", "VOPDXY_LIT": "VOPD_LIT", "VDS": "DS"}
|
||||
# Instructions missing from XML but present in PDF
|
||||
FIXES = {"rdna3": {"SOPK": {22: "S_SUBVECTOR_LOOP_BEGIN", 23: "S_SUBVECTOR_LOOP_END"}, "FLAT": {55: "FLAT_ATOMIC_CSUB_U32"}},
|
||||
"rdna4": {"SOP1": {80: "S_GET_BARRIER_STATE", 81: "S_BARRIER_INIT", 82: "S_BARRIER_JOIN"}, "SOPP": {9: "S_WAITCNT", 21: "S_BARRIER_LEAVE"}},
|
||||
"cdna": {"DS": {152: "DS_GWS_SEMA_RELEASE_ALL", 154: "DS_GWS_SEMA_V", 156: "DS_GWS_SEMA_P"},
|
||||
"VOP3P": {44: "V_MFMA_LD_SCALE_B32", 62: "V_MFMA_F32_16X16X8_XF32", 63: "V_MFMA_F32_32X32X4_XF32"}}}
|
||||
# Fields missing from XML but present in hardware (format: {arch: {encoding: [(name, hi, lo), ...]}})
|
||||
FIELD_FIXES = {"cdna": {"VOP3P": [("opsel_hi2", 14, 14)]}}
|
||||
# Encoding suffixes to strip (variants we don't generate separate classes for)
|
||||
_ENC_SUFFIXES = ("_NSA1",)
|
||||
# Encoding suffix to class suffix mapping (for variants we DO generate)
|
||||
_ENC_SUFFIX_MAP = {"_INST_LITERAL": "_LIT", "_VOP_DPP16": "_DPP16", "_VOP_DPP": "_DPP16", "_VOP_DPP8": "_DPP8",
|
||||
"_VOP_SDWA": "_SDWA", "_VOP_SDWA_SDST_ENC": "_SDWA_SDST", "_MFMA": "_MFMA"}
|
||||
# Field name normalization
|
||||
_FIELD_RENAMES = {"opsel_hi_2": "opsel_hi2", "op_sel_hi_2": "opsel_hi2", "op_sel": "opsel", "bound_ctrl": "bc",
|
||||
"tgt": "target", "row_en": "row", "unorm": "unrm", "clamp": "clmp", "wait_exp": "waitexp",
|
||||
"simm32": "literal", "dpp_ctrl": "dpp", "acc_cd": "acc_cd", "acc": "acc",
|
||||
"dst_sel": "dst_sel", "dst_unused": "dst_unused", "src0_sel": "src0_sel", "src1_sel": "src1_sel"}
|
||||
# Encoding variants to skip entirely (NSA is for MIMG graphics instructions)
|
||||
_SKIP_ENCODINGS = ("NSA",)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# XML parsing helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _strip_enc(name: str) -> str:
|
||||
"""Strip ENC_ prefix and normalize encoding suffixes."""
|
||||
name = name.removeprefix("ENC_")
|
||||
for sfx in _ENC_SUFFIXES: name = name.replace(sfx, "")
|
||||
# Process longer suffixes first to avoid partial matches (e.g., _VOP_DPP8 before _VOP_DPP)
|
||||
for old, new in sorted(_ENC_SUFFIX_MAP.items(), key=lambda x: -len(x[0])): name = name.replace(old, new)
|
||||
return name
|
||||
|
||||
def _norm_field(name: str) -> str:
|
||||
"""Normalize field name to match expected names."""
|
||||
for old, new in _FIELD_RENAMES.items(): name = name.replace(old, new)
|
||||
return name
|
||||
|
||||
def _map_flat(enc_name: str, instr_name: str) -> str:
|
||||
"""Map FLAT/GLOBAL/SCRATCH encoding to correct enum based on instruction prefix."""
|
||||
if enc_name in ("FLAT_GLBL", "FLAT_GLOBAL"): return "GLOBAL"
|
||||
if enc_name == "FLAT_SCRATCH": return "SCRATCH"
|
||||
if enc_name in ("FLAT", "VFLAT", "VGLOBAL", "VSCRATCH"):
|
||||
v = "V" if enc_name.startswith("V") else ""
|
||||
if instr_name.startswith("GLOBAL_"): return f"{v}GLOBAL"
|
||||
if instr_name.startswith("SCRATCH_"): return f"{v}SCRATCH"
|
||||
return f"{v}FLAT"
|
||||
return enc_name
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# XML parsing
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def parse_xml(filename: str):
|
||||
root = ET.fromstring(zipfile.ZipFile(fetch(XML_URL)).read(filename))
|
||||
encodings, enums, types, fmts, op_types_set = {}, {}, {}, {}, set()
|
||||
# Extract HWREG and MSG enums from OperandTypes
|
||||
op_enum_map = {("OPR_HWREG", "ID"): "HWREG", ("OPR_SENDMSG_RTN", "MSG"): "MSG"}
|
||||
for ot in root.findall(".//OperandTypes/OperandType"):
|
||||
ot_name = ot.findtext("OperandTypeName")
|
||||
for field in ot.findall(".//Field"):
|
||||
key = (ot_name, field.findtext("FieldName"))
|
||||
if (enum_name := op_enum_map.get(key)): # type: ignore[arg-type]
|
||||
def _pv_val(pv: ET.Element) -> tuple[int, str]:
|
||||
v, n = pv.findtext("Value"), pv.findtext("Name")
|
||||
assert v is not None and n is not None
|
||||
return int(v), n.upper()
|
||||
enums[enum_name] = dict(_pv_val(pv) for pv in field.findall(".//PredefinedValue"))
|
||||
# Extract DataFormats with BitCount
|
||||
for df in root.findall("ISA/DataFormats/DataFormat"):
|
||||
name, bits = df.findtext("DataFormatName"), df.findtext("BitCount")
|
||||
if name and bits: fmts[name] = int(bits)
|
||||
# Extract encoding definitions
|
||||
for enc in root.findall("ISA/Encodings/Encoding"):
|
||||
name = enc.findtext("EncodingName")
|
||||
assert name is not None
|
||||
is_base = name.startswith("ENC_") or name in ("VOP3_SDST_ENC", "VOPDXY")
|
||||
is_variant = any(sfx in name for sfx in _ENC_SUFFIX_MAP)
|
||||
if not is_base and not is_variant: continue
|
||||
if any(s in name for s in _SKIP_ENCODINGS): continue
|
||||
fields: list[tuple[str, int, int]] = []
|
||||
for f in enc.findall(".//MicrocodeFormat/BitMap/Field"):
|
||||
br = f.find("BitLayout/Range")
|
||||
if br is None: continue
|
||||
fn = f.findtext("FieldName")
|
||||
assert fn is not None
|
||||
fields.append((_norm_field(fn.lower()),
|
||||
int(br.findtext("BitOffset") or 0) + int(br.findtext("BitCount") or 0) - 1, int(br.findtext("BitOffset") or 0)))
|
||||
ident_list = enc.findall("EncodingIdentifiers/EncodingIdentifier")
|
||||
ident = ident_list[0] if ident_list else None
|
||||
enc_field = next((f for f in fields if f[0] == "encoding"), None)
|
||||
# For multi-dword formats, encoding field may be in higher dword but identifier is always in dword0; use % 32
|
||||
enc_bits: str | None = None
|
||||
if ident is not None and ident.text is not None and enc_field:
|
||||
enc_bits = "".join(ident.text[len(ident.text)-1-b] for b in range(enc_field[1] % 32, (enc_field[2] % 32)-1, -1))
|
||||
base_name = _strip_enc(name)
|
||||
encodings[NAME_MAP.get(base_name, base_name)] = (fields, enc_bits)
|
||||
# Extract instruction opcodes and operand info
|
||||
# Track which encodings each opcode appears in (for detecting LIT-only ops)
|
||||
opcode_encs: dict[str, dict[int, set[str]]] = {} # {base_fmt: {opcode: {enc_names}}}
|
||||
for instr in root.findall("ISA/Instructions/Instruction"):
|
||||
name = instr.findtext("InstructionName")
|
||||
assert name is not None
|
||||
for enc in instr.findall("InstructionEncodings/InstructionEncoding"):
|
||||
if enc.findtext("EncodingCondition") != "default": continue
|
||||
enc_enc_name = enc.findtext("EncodingName")
|
||||
assert enc_enc_name is not None
|
||||
base, opcode = _map_flat(_strip_enc(enc_enc_name), name), int(enc.findtext("Opcode") or 0)
|
||||
enc_name = NAME_MAP.get(base, base)
|
||||
# Encoding variants use the same Op enum as the base format
|
||||
base_enum = enc_name
|
||||
for sfx in ("_SDWA_SDST", "_DPP16", "_DPP8", "_SDWA", "_LIT", "_MFMA"):
|
||||
base_enum = base_enum.replace(sfx, "")
|
||||
# Track which encodings this opcode appears in
|
||||
opcode_encs.setdefault(base_enum, {}).setdefault(opcode, set()).add(enc_name)
|
||||
# ADDTID instructions go in both FLAT and GLOBAL enums (pcode uses FLATOp for these)
|
||||
if "ADDTID" in name:
|
||||
if base == "GLOBAL": enums.setdefault("FLAT", {})[opcode] = name
|
||||
elif base == "VGLOBAL": enums.setdefault("VFLAT", {})[opcode] = name
|
||||
enums.setdefault(base_enum, {})[opcode] = name
|
||||
# Extract operand info
|
||||
op_info: dict[str, tuple[str | None, int, str | None]] = {}
|
||||
for op in enc.findall("Operands/Operand"):
|
||||
fn = op.findtext("FieldName")
|
||||
if fn: op_info[fn.lower()] = (op.findtext("DataFormatName"), int(op.findtext("OperandSize") or 0), op.findtext("OperandType"))
|
||||
for fmt, _, otype in op_info.values():
|
||||
if fmt and fmt not in fmts: fmts[fmt] = 0
|
||||
if otype: op_types_set.add(otype)
|
||||
if op_info: types[(name, base_enum)] = op_info
|
||||
# Find opcodes that only exist in a specific variant encoding (no base format version)
|
||||
suffix_only_ops: dict[str, dict[str, set[int]]] = {} # {suffix: {base_fmt: {opcodes}}}
|
||||
for base_fmt, opcodes in opcode_encs.items():
|
||||
for opcode, encs in opcodes.items():
|
||||
suffix = next((s for s in _ENC_SUFFIX_MAP.values() if all(s in e for e in encs)), None)
|
||||
if suffix is not None: suffix_only_ops.setdefault(suffix, {}).setdefault(base_fmt, set()).add(opcode)
|
||||
return encodings, enums, types, fmts, op_types_set, suffix_only_ops
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PDF parsing
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def extract_pdf_text(url: str) -> list[list[tuple[float, float, str, str]]]:
|
||||
"""Extract positioned text from PDF. Returns list of text elements (x, y, text, font) per page."""
|
||||
data = fetch(url).read_bytes()
|
||||
# Parse xref table to locate objects
|
||||
xref: dict[int, int] = {}
|
||||
xref_match = re.search(rb'startxref\s+(\d+)', data)
|
||||
assert xref_match is not None
|
||||
pos = int(xref_match.group(1)) + 4
|
||||
while data[pos:pos+7] != b'trailer':
|
||||
while data[pos:pos+1] in b' \r\n': pos += 1
|
||||
line_end = data.find(b'\n', pos)
|
||||
start_obj, count = map(int, data[pos:line_end].split()[:2])
|
||||
pos = line_end + 1
|
||||
for i in range(count):
|
||||
if data[pos+17:pos+18] == b'n' and (off := int(data[pos:pos+10])) > 0: xref[start_obj + i] = off
|
||||
pos += 20
|
||||
|
||||
def get_stream(n: int) -> bytes:
|
||||
obj = data[xref[n]:data.find(b'endobj', xref[n])]
|
||||
raw = obj[obj.find(b'stream\n') + 7:obj.find(b'\nendstream')]
|
||||
return zlib.decompress(raw) if b'/FlateDecode' in obj else raw
|
||||
|
||||
pages = []
|
||||
for n in sorted(xref):
|
||||
if b'/Type /Page' not in data[xref[n]:xref[n]+500]: continue
|
||||
if not (m := re.search(rb'/Contents (\d+) 0 R', data[xref[n]:xref[n]+500])): continue
|
||||
stream = get_stream(int(m.group(1))).decode('latin-1')
|
||||
elements, font = [], ''
|
||||
_RE_BT = (r'(/F[\d.]+) [\d.]+ Tf|([\d.+-]+) ([\d.+-]+) Td|[\d.+-]+ [\d.+-]+ [\d.+-]+ [\d.+-]+ ([\d.+-]+) ([\d.+-]+) Tm'
|
||||
r'|<([0-9A-Fa-f]+)>.*?Tj|\[([^\]]+)\] TJ')
|
||||
for bt in re.finditer(r'BT(.*?)ET', stream, re.S):
|
||||
x, y = 0.0, 0.0
|
||||
for sm in re.finditer(_RE_BT, bt.group(1)):
|
||||
if sm.group(1): font = sm.group(1)
|
||||
elif sm.group(2): x, y = x + float(sm.group(2)), y + float(sm.group(3))
|
||||
elif sm.group(4): x, y = float(sm.group(4)), float(sm.group(5))
|
||||
elif sm.group(6) and (t := bytes.fromhex(sm.group(6)).decode('latin-1')).strip():
|
||||
elements.append((x, y, t, font))
|
||||
elif sm.group(7):
|
||||
t = ''.join(bytes.fromhex(h).decode('latin-1') for h in re.findall(r'<([0-9A-Fa-f]+)>', sm.group(7)))
|
||||
if t.strip(): elements.append((x, y, t, font))
|
||||
pages.append(sorted(elements, key=lambda e: (-e[1], e[0])))
|
||||
return pages
|
||||
|
||||
def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op: dict[str, int]) -> dict[tuple[str, int], str]:
|
||||
"""Extract pseudocode for instructions. Returns {(name, opcode): pseudocode}."""
|
||||
# First pass: find all instruction headers across all pages
|
||||
all_instructions: list[tuple[int, float, str, int]] = [] # (page_idx, y, name, opcode)
|
||||
for page_idx, page in enumerate(pages):
|
||||
by_y: dict[int, list[tuple[float, str]]] = {}
|
||||
for x, y, t, _ in page:
|
||||
by_y.setdefault(round(y), []).append((x, t))
|
||||
for y, items in sorted(by_y.items(), reverse=True):
|
||||
left = [(x, t) for x, t in items if 55 < x < 65]
|
||||
right = [(x, t) for x, t in items if 535 < x < 550]
|
||||
if left and right and left[0][1] in name_to_op and right[0][1].isdigit():
|
||||
all_instructions.append((page_idx, y, left[0][1], int(right[0][1])))
|
||||
|
||||
# Second pass: extract pseudocode between consecutive instructions
|
||||
pcode: dict[tuple[str, int], str] = {}
|
||||
for i, (page_idx, y, name, opcode) in enumerate(all_instructions):
|
||||
if i + 1 < len(all_instructions):
|
||||
next_page, next_y = all_instructions[i + 1][0], all_instructions[i + 1][1]
|
||||
else:
|
||||
next_page, next_y = page_idx, 0
|
||||
# Collect F6 text from current position to next instruction (pseudocode is at x ≈ 69)
|
||||
lines: list[tuple[int, float, str]] = []
|
||||
for p in range(page_idx, next_page + 1):
|
||||
start_y = y if p == page_idx else 800
|
||||
end_y = next_y if p == next_page else 0
|
||||
lines.extend((p, y2, t) for x, y2, t, f in pages[p] if f in ('/F6.0', '/F7.0') and end_y < y2 < start_y and 60 < x < 80)
|
||||
if lines:
|
||||
sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
|
||||
# Stop at large Y gaps (>30) - indicates section break
|
||||
filtered = [sorted_lines[0]]
|
||||
for j in range(1, len(sorted_lines)):
|
||||
prev_page, prev_y, _ = sorted_lines[j-1]
|
||||
curr_page, curr_y, _ = sorted_lines[j]
|
||||
if curr_page == prev_page and prev_y - curr_y > 30: break
|
||||
if curr_page != prev_page and prev_y > 60 and curr_y < 730: break
|
||||
filtered.append(sorted_lines[j])
|
||||
pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
|
||||
if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
|
||||
return pcode
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Code generation
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def write_common(all_fmts: dict[str, int], all_op_types: set[str], path: pathlib.Path) -> None:
|
||||
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "from enum import Enum, auto", ""]
|
||||
lines.append("class ReprEnum(Enum):")
|
||||
lines.append(' """Enum with clean repr that roundtrips with eval()."""')
|
||||
lines.append(' def __repr__(self): return f"{type(self).__name__}.{self.name}"')
|
||||
lines.append("")
|
||||
lines.append("class Fmt(Enum):")
|
||||
for fmt in sorted(all_fmts.keys()): lines.append(f" {fmt} = auto()")
|
||||
lines.append("")
|
||||
lines.append("FMT_BITS = {")
|
||||
for fmt, bits in sorted(all_fmts.items()): lines.append(f" Fmt.{fmt}: {bits},")
|
||||
lines.append("}")
|
||||
lines.append("")
|
||||
lines.append("class OpType(Enum):")
|
||||
for ot in sorted(all_op_types): lines.append(f" {ot} = auto()")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_enum(enums, path):
|
||||
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
|
||||
"from tinygrad.runtime.autogen.amd.common import ReprEnum, Fmt, FMT_BITS, OpType # noqa: F401", ""]
|
||||
for name, ops in sorted(enums.items()):
|
||||
if not ops: continue
|
||||
suffix = "_E32" if name in ("VOP1", "VOP2", "VOPC") else "_E64" if name == "VOP3" else ""
|
||||
lines.append(f"class {name}(ReprEnum):" if name in ("HWREG", "MSG") else f"class {name}Op(ReprEnum):")
|
||||
aliases = []
|
||||
for op, mem in sorted(ops.items()):
|
||||
msuf = suffix if name != "VOP3" or op < 512 else ""
|
||||
lines.append(f" {mem}{msuf} = {op}")
|
||||
if msuf: aliases.append((mem, msuf))
|
||||
for mem, msuf in aliases: lines.append(f" {mem} = {mem}{msuf}")
|
||||
lines.append("")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_ins(encodings, enums, suffix_only_ops, types, arch, path):
|
||||
_VGPR_FIELDS = {"vdst", "vdstx", "vsrc0", "vsrc1", "vsrc2", "vsrc3", "vsrcx1", "vsrcy1", "vaddr", "vdata", "data", "data0", "data1", "addr", "vsrc"}
|
||||
_VARIANT_SUFFIXES = ("_LIT", "_DPP16", "_DPP8", "_SDWA_SDST", "_SDWA", "_MFMA")
|
||||
def get_base_fmt(fmt):
|
||||
for sfx in _VARIANT_SUFFIXES: fmt = fmt.replace(sfx, "")
|
||||
return fmt
|
||||
def field_def(name, hi, lo, fmt, enc_bits=None):
|
||||
bits = hi - lo + 1
|
||||
base_fmt = get_base_fmt(fmt)
|
||||
if name == "encoding" and enc_bits: return f"FixedBitField({hi}, {lo}, 0b{enc_bits})"
|
||||
if name == "op" and fmt not in ("DPP", "SDWA"): return f"EnumBitField({hi}, {lo}, {base_fmt}Op)"
|
||||
if name in ("opx", "opy"): return f"EnumBitField({hi}, {lo}, VOPDOp)"
|
||||
if name == "vdsty": return f"VDSTYField({hi}, {lo})"
|
||||
if name in _VGPR_FIELDS and bits == 8: return f"VGPRField({hi}, {lo})"
|
||||
if name == "sbase" and bits == 6: return f"SBaseField({hi}, {lo})"
|
||||
if name in ("srsrc", "ssamp") and bits == 5: return f"SRsrcField({hi}, {lo})"
|
||||
if name in ("sdst", "sdata") and bits == 7: return f"SGPRField({hi}, {lo})"
|
||||
if name in ("soffset", "saddr") and bits == 7: return f"SGPRField({hi}, {lo}, default=NULL)"
|
||||
if name.startswith("ssrc") and bits == 8: return f"SSrcField({hi}, {lo})"
|
||||
if name in ("saddr", "soffset") and bits == 8: return f"SSrcField({hi}, {lo}, default=NULL)"
|
||||
if name.startswith("src") and bits == 9: return f"SrcField({hi}, {lo})"
|
||||
# GLOBAL/SCRATCH: offset is 13-bit signed [12:0], FLAT: 12-bit unsigned (XML has 12-bit for all)
|
||||
if name == "offset" and base_fmt in ("GLOBAL", "SCRATCH"): return f"BitField(12, {lo})"
|
||||
if base_fmt == "VOP3P" and name == "opsel_hi": return f"BitField({hi}, {lo}, default=3)"
|
||||
if base_fmt == "VOP3P" and name == "opsel_hi2": return f"BitField({hi}, {lo}, default=1)"
|
||||
return f"BitField({hi}, {lo})"
|
||||
ORDER = ['encoding', 'op', 'opx', 'opy', 'vdst', 'vdstx', 'vdsty', 'sdst', 'vdata', 'sdata', 'addr', 'vaddr', 'data', 'data0', 'data1',
|
||||
'src0', 'srcx0', 'srcy0', 'vsrc0', 'ssrc0', 'src1', 'vsrc1', 'vsrcx1', 'vsrcy1', 'ssrc1', 'src2', 'vsrc2', 'src3', 'vsrc3',
|
||||
'saddr', 'sbase', 'srsrc', 'ssamp', 'soffset', 'offset', 'simm16', 'literal', 'en', 'target', 'attr', 'attr_chan',
|
||||
'omod', 'neg', 'neg_hi', 'abs', 'clmp', 'opsel', 'opsel_hi', 'waitexp', 'wait_va',
|
||||
'dmask', 'dim', 'seg', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe', 'unrm', 'done', 'row',
|
||||
'dpp', 'fi', 'bc', 'row_mask', 'bank_mask', 'src0_neg', 'src0_abs', 'src1_neg', 'src1_abs',
|
||||
'cbsz', 'abid', 'acc_cd', 'acc', 'blgp', 'lane_sel_0', 'lane_sel_1', 'lane_sel_2', 'lane_sel_3',
|
||||
'lane_sel_4', 'lane_sel_5', 'lane_sel_6', 'lane_sel_7', 'dst_sel', 'dst_unused', 'src0_sel', 'src1_sel']
|
||||
def sort_fields(fields): return sorted(fields, key=lambda f: (ORDER.index(f[0]) if f[0] in ORDER else 999, f[2]))
|
||||
|
||||
# Separate base encodings from variants
|
||||
base_encodings, variant_encodings = {}, {}
|
||||
for enc_name, data in encodings.items():
|
||||
base = get_base_fmt(enc_name)
|
||||
if base == enc_name: base_encodings[enc_name] = data
|
||||
else: variant_encodings[enc_name] = data
|
||||
|
||||
# Build sets of ops by their vdst type from operand metadata
|
||||
sdst_opcodes: dict[str, set[int]] = {} # ops where vdst is OPR_SREG (writes to SGPR)
|
||||
for fmt, ops in enums.items():
|
||||
for op, name in ops.items():
|
||||
op_types = types.get((name, fmt), {})
|
||||
vdst_type = op_types.get("vdst", (None, None, None))[2]
|
||||
if vdst_type == "OPR_SREG": sdst_opcodes.setdefault(fmt, set()).add(op)
|
||||
|
||||
# collect only the XxxOp enums that are actually referenced in this arch's instruction definitions
|
||||
enum_names = sorted(f"{k}Op" for k in enums if enums[k] and k not in ("HWREG", "MSG"))
|
||||
# also re-export HWREG/MSG enums (plain enums, not instruction format ops)
|
||||
enum_names += sorted(k for k in enums if k in ("HWREG", "MSG") and enums[k])
|
||||
# collect DSL field types actually used by scanning generated field definitions
|
||||
all_field_defs = " ".join(field_def(fn, hi, lo, enc, eb) for enc, (flds, eb) in encodings.items() for fn, hi, lo in flds)
|
||||
_ALL_DSL = ["BitField", "EnumBitField", "FixedBitField", "NULL", "SBaseField", "SGPRField", "SRsrcField",
|
||||
"SSrcField", "SrcField", "VDSTYField", "VGPRField"]
|
||||
dsl_names = ["Inst"] + [n for n in _ALL_DSL if n in all_field_defs]
|
||||
# also re-export register names so `from ins import *` still provides them to downstream users
|
||||
_DSL_REGS = ["s", "v", "src", "VCC_LO", "VCC_HI", "VCC", "EXEC_LO", "EXEC_HI", "EXEC", "NULL", "OFF", "M0",
|
||||
"SCC", "VCCZ", "EXECZ", "ttmp", "INV_2PI", "SDWA", "DPP", "DPP16", "LIT", "SRC_LDS_DIRECT"]
|
||||
dsl_reexport = sorted(set(dsl_names + _DSL_REGS))
|
||||
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit", "# ruff: noqa: E501,F401",
|
||||
f"from tinygrad.renderer.amd.dsl import {', '.join(dsl_reexport)}",
|
||||
f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "import functools", ""]
|
||||
|
||||
def fmt_allowed(op_enum: str, ops: set[int]) -> str:
|
||||
"""Format allowed ops as {EnumName.MEMBER, ...}."""
|
||||
names = [f"{op_enum}.{enums[op_enum.removesuffix('Op')][op]}" for op in sorted(ops)]
|
||||
return "{" + ", ".join(names) + "}"
|
||||
|
||||
# Generate base classes first
|
||||
for enc_name, (fields, enc_bits) in sorted(base_encodings.items()):
|
||||
all_ops = set(enums.get(enc_name, {}).keys())
|
||||
# Get suffix-only ops for this format (these can't be used in base class)
|
||||
base_suffix_ops = set().union(*(d.get(enc_name, set()) for d in suffix_only_ops.values()))
|
||||
# Exclude SDST ops from base class (they need VOP1_SDST/VOP3_SDST/VOP3B)
|
||||
base_allowed = all_ops - base_suffix_ops - sdst_opcodes.get(enc_name, set())
|
||||
# RDNA3 FLAT/GLOBAL/SCRATCH share encoding bits, differentiated by seg field
|
||||
# RDNA4 VFLAT/VGLOBAL/VSCRATCH have distinct encoding bits, no seg field needed
|
||||
has_seg_field = any(fn == "seg" for fn, _, _ in fields)
|
||||
if enc_name in ("FLAT", "VFLAT") and has_seg_field:
|
||||
prefix = "V" if enc_name == "VFLAT" else ""
|
||||
flat_variants = [(f"{prefix}FLAT", 0, f"{prefix}FLATOp"), (f"{prefix}GLOBAL", 2, f"{prefix}GLOBALOp"),
|
||||
(f"{prefix}SCRATCH", 1, f"{prefix}SCRATCHOp")]
|
||||
for cls, seg, op_enum in flat_variants:
|
||||
cls_ops = set(enums.get(cls, {}).keys())
|
||||
lines.append(f"class {cls}(Inst):")
|
||||
for fn, hi, lo in sort_fields(fields):
|
||||
if fn == "seg": lines.append(f" seg = FixedBitField({hi}, {lo}, {seg})")
|
||||
elif fn == "op": lines.append(f" op = EnumBitField({hi}, {lo}, {op_enum}, {fmt_allowed(op_enum, cls_ops)})")
|
||||
else: lines.append(f" {fn} = {field_def(fn, hi, lo, cls, enc_bits)}")
|
||||
lines.append("")
|
||||
elif enc_name not in ("FLAT_GLOBAL", "FLAT_SCRATCH", "FLAT_GLBL", "DPP", "SDWA"):
|
||||
lines.append(f"class {enc_name}(Inst):")
|
||||
for fn, hi, lo in sort_fields(fields):
|
||||
if fn == "op":
|
||||
base_fmt = get_base_fmt(enc_name)
|
||||
lines.append(f" op = EnumBitField({hi}, {lo}, {base_fmt}Op, {fmt_allowed(f'{base_fmt}Op', base_allowed)})")
|
||||
else:
|
||||
lines.append(f" {fn} = {field_def(fn, hi, lo, enc_name, enc_bits if fn == 'encoding' else None)}")
|
||||
lines.append("")
|
||||
|
||||
# Generate variant classes that inherit from base (only add extra fields)
|
||||
for enc_name, (fields, enc_bits) in sorted(variant_encodings.items()):
|
||||
base = get_base_fmt(enc_name)
|
||||
if base not in base_encodings: continue # skip if no base class
|
||||
base_fields = {f[0] for f in base_encodings[base][0]}
|
||||
extra_fields = [(fn, hi, lo) for fn, hi, lo in fields if fn not in base_fields]
|
||||
# Check if this is a suffix-only variant
|
||||
variant_suffix = next((sfx for sfx in _VARIANT_SUFFIXES if enc_name.endswith(sfx)), None)
|
||||
is_suffix_variant = variant_suffix in suffix_only_ops
|
||||
all_ops = set(enums.get(base, {}).keys())
|
||||
if extra_fields or is_suffix_variant:
|
||||
lines.append(f"class {enc_name}({base}):")
|
||||
op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
|
||||
# _LIT classes: override op to allow all opcodes (base excludes lit-only ops)
|
||||
# other classes override op to only suffix-only opcodes
|
||||
if op_field and is_suffix_variant:
|
||||
_, hi, lo = op_field
|
||||
allowed_ops = all_ops if variant_suffix == "_LIT" else suffix_only_ops[variant_suffix][base]
|
||||
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', allowed_ops)})")
|
||||
for fn, hi, lo in sort_fields(extra_fields):
|
||||
lines.append(f" {fn} = {field_def(fn, hi, lo, enc_name)}")
|
||||
lines.append("")
|
||||
|
||||
# SDST variants (special case - redefine vdst field type, restrict to SDST ops)
|
||||
for base, field_hi, field_lo in [("VOP1", 24, 17), ("VOP3", 7, 0)]:
|
||||
if base not in base_encodings: continue
|
||||
sdst_ops = sdst_opcodes.get(base, set())
|
||||
if not sdst_ops: continue
|
||||
# For VOP3, all ops < 256 (compare/cmpx ops) use SDST encoding
|
||||
all_base_ops = set(enums.get(base, {}).keys())
|
||||
if base == "VOP3": sdst_ops = sdst_ops | {op for op in all_base_ops if op < 256}
|
||||
op_field = next((f for f in base_encodings[base][0] if f[0] == "op"), None)
|
||||
lines.append(f"class {base}_SDST({base}):")
|
||||
if op_field:
|
||||
_, hi, lo = op_field
|
||||
lines.append(f" op = EnumBitField({hi}, {lo}, {base}Op, {fmt_allowed(f'{base}Op', sdst_ops)})")
|
||||
lines.append(f" vdst = SSrcField({field_hi}, {field_lo})")
|
||||
lines.append("")
|
||||
# SDST_LIT class (for literals with SDST destination) - same ops, just adds literal field
|
||||
lit_enc = variant_encodings.get(f"{base}_LIT")
|
||||
if lit_enc:
|
||||
lit_field = next((f for f in lit_enc[0] if f[0] == "literal"), None)
|
||||
if lit_field:
|
||||
lines.append(f"class {base}_SDST_LIT({base}_SDST):")
|
||||
lines.append(f" literal = BitField({lit_field[1]}, {lit_field[2]})")
|
||||
lines.append("")
|
||||
|
||||
# Instruction helpers
|
||||
lines.append("# instruction helpers")
|
||||
for fmt, ops in sorted(enums.items()):
|
||||
if fmt not in base_encodings and fmt not in ("GLOBAL", "SCRATCH", "VGLOBAL", "VSCRATCH"): continue
|
||||
suffix = "_E32" if fmt in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt == "VOP3" else ""
|
||||
op_to_suffix = {op:suffix for suffix,ops in suffix_only_ops.items() for op in ops.get(fmt, set())}
|
||||
fmt_sdst_ops = sdst_opcodes.get(fmt, set())
|
||||
for op, name in sorted(ops.items()):
|
||||
# ADDTID ops are in both FLAT and GLOBAL enums (for pcode); only generate helper for GLOBAL/VGLOBAL
|
||||
if "ADDTID" in name and fmt in ("FLAT", "VFLAT"): continue
|
||||
msuf = suffix if fmt != "VOP3" or op < 512 else ""
|
||||
# Determine class: SDST variants, suffix-specific variants (e.g., _MFMA, _LIT), or base
|
||||
if fmt == "VOP1" and op in fmt_sdst_ops: cls = "VOP1_SDST"
|
||||
elif fmt == "VOP3" and (op in fmt_sdst_ops or op < 256): cls = "VOP3_SDST"
|
||||
elif op_to_suffix.get(op): cls = f"{fmt}{op_to_suffix[op]}"
|
||||
else: cls = fmt
|
||||
lines.append(f"{name.lower()}{msuf.lower()} = functools.partial({cls}, {fmt}Op.{name}{msuf})")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_operands(types: dict, enums: dict, arch: str, path: pathlib.Path) -> None:
|
||||
valid = {(name, fmt) for fmt, ops in enums.items() for name in ops.values()}
|
||||
# only import enums that are actually used as keys in OPERANDS
|
||||
used_bases = {eb for (nm, eb) in types if (nm, eb) in valid}
|
||||
enum_names = sorted(f"{k}Op" for k in used_bases)
|
||||
lines: list[str] = ["# autogenerated from AMD ISA XML - do not edit",
|
||||
"from tinygrad.runtime.autogen.amd.common import Fmt, OpType",
|
||||
f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", ""]
|
||||
lines.append("# instruction operand info: {Op: {field: (Fmt, size_bits, OpType)}}")
|
||||
lines.append("OPERANDS = {")
|
||||
def fmt_val(v):
|
||||
fmt, size, otype = v
|
||||
return f"({f'Fmt.{fmt}' if fmt else 'None'}, {size}, {f'OpType.{otype}' if otype else 'None'})"
|
||||
for (name, enc_base), fields in sorted(types.items()):
|
||||
if (name, enc_base) not in valid: continue
|
||||
fstr = ", ".join(f'"{k}": {fmt_val(v)}' for k, v in sorted(fields.items()))
|
||||
lines.append(f' {enc_base}Op.{name}: {{{fstr}}},')
|
||||
lines.append("}")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
def write_pcode(pcode: dict[tuple[str, int], str], enums: dict[str, dict[int, str]], arch: str, path: pathlib.Path) -> None:
|
||||
"""Write str_pcode.py file from extracted pseudocode."""
|
||||
entries: list[tuple[str, str, int, str]] = []
|
||||
for fmt_name, ops in enums.items():
|
||||
member_suffix = "_E32" if fmt_name in ("VOP1", "VOP2", "VOPC") else "_E64" if fmt_name == "VOP3" else ""
|
||||
for opcode, name in ops.items():
|
||||
if (name, opcode) in pcode:
|
||||
msuf = member_suffix if fmt_name != "VOP3" or opcode < 512 else ""
|
||||
entries.append((f"{fmt_name}Op", f"{name}{msuf}", opcode, pcode[(name, opcode)]))
|
||||
enum_names = sorted(set(e[0] for e in entries))
|
||||
lines = ["# autogenerated from AMD ISA PDF - do not edit", "# ruff: noqa: E501",
|
||||
f"from tinygrad.runtime.autogen.amd.{arch}.enum import {', '.join(enum_names)}", "", "PCODE = {"]
|
||||
for enum_name, name, opcode, code in sorted(entries, key=lambda x: (x[0], x[2])):
|
||||
lines.append(f" {enum_name}.{name}: {code!r},")
|
||||
lines.append("}")
|
||||
with open(path, "w") as f: f.write("\n".join(lines))
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# Main
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
if __name__ == "__main__":
|
||||
all_fmts: dict[str, int] = {}
|
||||
all_op_types: set[str] = set()
|
||||
arch_data: dict[str, dict] = {}
|
||||
# First pass: parse XML for all architectures
|
||||
for arch, cfg in ARCHS.items():
|
||||
print(f"Parsing XML: {cfg['xml']} -> {arch}")
|
||||
encodings, enums, types, fmts, op_types_set, suffix_only_ops = parse_xml(cfg["xml"])
|
||||
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
|
||||
for fmt, fields in FIELD_FIXES.get(arch, {}).items():
|
||||
if fmt in encodings: encodings[fmt] = (encodings[fmt][0] + fields, encodings[fmt][1])
|
||||
arch_data[arch] = {"encodings": encodings, "enums": enums, "types": types, "suffix_only_ops": suffix_only_ops}
|
||||
for fmt, bits in fmts.items():
|
||||
assert fmt not in all_fmts or all_fmts[fmt] == bits, f"FMT_BITS mismatch for {fmt}: {all_fmts[fmt]} vs {bits}"
|
||||
all_fmts[fmt] = bits
|
||||
all_op_types.update(op_types_set)
|
||||
# Write common.py
|
||||
autogen_base = pathlib.Path(__file__).parents[2] / "runtime" / "autogen" / "amd"
|
||||
common_path = autogen_base / "common.py"
|
||||
write_common(all_fmts, all_op_types, common_path)
|
||||
print(f"Wrote common.py: {len(all_fmts)} formats, {len(all_op_types)} op types")
|
||||
# Write per-arch files from XML
|
||||
for arch, data in arch_data.items():
|
||||
base = autogen_base / arch
|
||||
write_enum(data["enums"], base / "enum.py")
|
||||
write_ins(data["encodings"], data["enums"], data["suffix_only_ops"], data["types"], arch, base / "ins.py")
|
||||
write_operands(data["types"], data["enums"], arch, base / "operands.py")
|
||||
print(f" {arch}: {len(data['encodings'])} encodings, {sum(len(v) for v in data['enums'].values())} instructions")
|
||||
# Second pass: parse PDFs and write pcode
|
||||
for arch, cfg in ARCHS.items():
|
||||
print(f"Parsing PDF: {arch}...")
|
||||
pages = extract_pdf_text(cfg["pdf"])
|
||||
name_to_op = {name: op for ops in arch_data[arch]["enums"].values() for op, name in ops.items()}
|
||||
pcode = extract_pcode(pages, name_to_op)
|
||||
base = autogen_base / arch
|
||||
write_pcode(pcode, arch_data[arch]["enums"], arch, base / "str_pcode.py")
|
||||
print(f" {arch}: {len(pcode)} pcode entries")
|
||||
745
tinygrad_repo/tinygrad/renderer/amd/sqtt.py
Normal file
745
tinygrad_repo/tinygrad/renderer/amd/sqtt.py
Normal file
@@ -0,0 +1,745 @@
|
||||
"""SQTT (SQ Thread Trace) packet encoder and decoder for AMD GPUs.
|
||||
|
||||
This module provides encoding and decoding of raw SQTT byte streams.
|
||||
The format is nibble-based with variable-width packets determined by a state machine.
|
||||
Uses BitField infrastructure from dsl.py, similar to GPU instruction encoding.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterator
|
||||
from enum import Enum
|
||||
from tinygrad.helpers import getenv, colored
|
||||
from tinygrad.renderer.amd.dsl import BitField, FixedBitField, Inst, bits
|
||||
from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm # same encoding as RDNA4
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# FIELD ENUMS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class MemSrc(Enum):
|
||||
LDS = 0
|
||||
LDS_ALT = 1
|
||||
VMEM = 2
|
||||
VMEM_ALT = 3
|
||||
|
||||
class AluSrc(Enum):
|
||||
NONE = 0
|
||||
SALU = 1
|
||||
VALU = 2
|
||||
VALU_SALU = 3
|
||||
|
||||
# construct other SIMD instruction operation types, name becomes OTHER_{category}_{cycles}
|
||||
def add_other_simd(cls:type[Enum], ranges:list[tuple[str, int, int, int]]) -> None:
|
||||
for category, start, end, base_cycle in ranges:
|
||||
for value in range(start, end + 1):
|
||||
cls._value2member_map_[value] = obj = object.__new__(cls)
|
||||
obj._value_ = value
|
||||
obj._name_ = f"OTHER_{category}_{value - start + base_cycle}"
|
||||
|
||||
class InstOp(Enum):
|
||||
"""SQTT instruction operation types for RDNA3 (gfx1100).
|
||||
|
||||
Memory ops appear in two ranges depending on which SIMD executes them:
|
||||
- 0x1x-0x2x range: ops on traced SIMD
|
||||
- 0x5x range: ops on other SIMD (OTHER_ prefix)
|
||||
|
||||
GLOBAL memory ops encoding depends on addressing mode AND size:
|
||||
- Loads: 0x21 (saddr=SGPR) or 0x22 (saddr=NULL), all sizes same
|
||||
- Stores: base + size_offset, where VADDR is shifted +1 from SADDR
|
||||
SADDR: 0x24(32) 0x25(64) 0x26(96) 0x27(128)
|
||||
VADDR: 0x25(32) 0x26(64) 0x27(96) 0x28(128)
|
||||
|
||||
OTHER_ range follows same pattern but values overlap differently.
|
||||
"""
|
||||
SALU = 0x0
|
||||
SMEM_RD = 0x1
|
||||
JUMP = 0x3 # branch taken
|
||||
JUMP_NO = 0x4 # branch not taken
|
||||
CALL = 0x5 # s_call_b64
|
||||
MESSAGE = 0x9
|
||||
VALUT_4 = 0xb # transcendental: exp, log, rcp, sqrt, sin, cos
|
||||
VALUB_2 = 0xd # 64-bit shifts: lshl, lshr, ashr
|
||||
VALUB_4 = 0xe # 64-bit multiply-add
|
||||
VALUB_16 = 0xf # 64-bit: add, mul, fma, rcp, sqrt, rounding, frexp, div helpers
|
||||
VINTERP = 0x12 # interpolation: v_interp_p10_f32, v_interp_p2_f32
|
||||
BARRIER = 0x13
|
||||
|
||||
# FLAT memory ops on traced SIMD (0x1x range)
|
||||
FLAT_RD_2 = 0x1c
|
||||
FLAT_WR_3 = 0x1d
|
||||
FLAT_WR_4 = 0x1e
|
||||
FLAT_WR_5 = 0x1f
|
||||
FLAT_WR_6 = 0x20
|
||||
|
||||
# GLOBAL memory ops on traced SIMD (0x2x range)
|
||||
SGMEM_RD_1 = 0x21 # saddr=SGPR, all sizes
|
||||
SGMEM_RD_2 = 0x22 # saddr=NULL, all sizes
|
||||
SGMEM_WR_2 = 0x24 # saddr=SGPR, 32-bit
|
||||
SGMEM_WR_3 = 0x25 # saddr=SGPR 64 or saddr=NULL 32
|
||||
SGMEM_WR_4 = 0x26 # saddr=SGPR 96 or saddr=NULL 64
|
||||
SGMEM_WR_5 = 0x27 # saddr=SGPR 128 or saddr=NULL 96
|
||||
SGMEM_WR_6 = 0x28 # saddr=NULL, 128-bit
|
||||
|
||||
# LDS ops on traced SIMD
|
||||
LDS_RD = 0x29
|
||||
LDS_WR_1 = 0x2a # ds_append, ds_consume, ds_store_addtid_b32
|
||||
LDS_WR_2 = 0x2b
|
||||
LDS_WR_3 = 0x2c
|
||||
LDS_WR_4 = 0x2d
|
||||
LDS_WR_5 = 0x2e
|
||||
|
||||
# EXEC-modifying ops (0x7x range)
|
||||
SALU_WR_EXEC = 0x72 # s_*_saveexec_b32/b64
|
||||
VALU1_WR_EXEC = 0x73 # v_cmpx_*
|
||||
# Memory ops on other SIMD (0x5x range)
|
||||
add_other_simd(InstOp, [("LDS", 0x50, 0x54, 1), ("FLAT", 0x55, 0x59, 2), ("VMEM", 0x5a, 0x66, 1)])
|
||||
|
||||
class InstOpRDNA4(Enum):
|
||||
"""SQTT instruction operation types for RDNA4 (gfx1200). Different encoding from RDNA3."""
|
||||
SALU = 0x0
|
||||
SMEM = 0x1
|
||||
SMEM_WR = 0x2
|
||||
JUMP = 0x3
|
||||
JUMP_NO = 0x4
|
||||
CALL = 0x5
|
||||
SALU_NO_EXEC = 0x7
|
||||
MESSAGE = 0x9
|
||||
VALU_1 = 0xa
|
||||
VALUT_4 = 0xb
|
||||
VALUB_1 = 0xc
|
||||
VALUB_2 = 0xd
|
||||
VALUB_4 = 0xe
|
||||
VALUB_16 = 0xf
|
||||
VINTERP = 0x12
|
||||
BARRIER_WAIT = 0x13
|
||||
FLAT_RD_2 = 0x1c
|
||||
FLAT_WR_3 = 0x1d
|
||||
FLAT_WR_4 = 0x1e
|
||||
FLAT_WR_5 = 0x1f
|
||||
FLAT_WR_6 = 0x20
|
||||
VMEM_RD_1 = 0x21
|
||||
VMEM_RD_2 = 0x22
|
||||
VMEM_WR_1 = 0x23
|
||||
VMEM_WR_2 = 0x24
|
||||
VMEM_WR_3 = 0x25
|
||||
VMEM_WR_4 = 0x26
|
||||
VMEM_WR_5 = 0x27
|
||||
VMEM_WR_6 = 0x28
|
||||
LDS_RD = 0x29
|
||||
LDS_WR_1 = 0x2a
|
||||
LDS_WR_2 = 0x2b
|
||||
LDS_WR_3 = 0x2c
|
||||
LDS_WR_4 = 0x2d
|
||||
LDS_WR_5 = 0x2e
|
||||
BUF_RD_1 = 0x2f
|
||||
BUF_RD_2 = 0x30
|
||||
BUF_WR_1 = 0x31
|
||||
BUF_WR_2 = 0x32
|
||||
BUF_WR_3 = 0x33
|
||||
BUF_WR_4 = 0x34
|
||||
BUF_WR_5 = 0x35
|
||||
BUF_WR_6 = 0x36
|
||||
LDS_DIR_LOAD = 0x6e
|
||||
LDS_PARAM_LOAD = 0x6f
|
||||
SALU_WR_EXEC = 0x72
|
||||
VALU1_WR_EXEC = 0x73
|
||||
VALU_WR_EXEC_2 = 0x74
|
||||
OTHER_LDS_6 = 0x77
|
||||
OTHER_LDS_10 = 0x78
|
||||
BARRIER_SIGNAL = 0x7a
|
||||
DYN_VGPR = 0x87
|
||||
BARRIER_JOIN = 0x8a
|
||||
WMMA_8 = 0x8c
|
||||
WMMA_16 = 0x8d
|
||||
WMMA_32 = 0x8e
|
||||
WMMA_64 = 0x8f
|
||||
VALU_DPFP = 0x92
|
||||
SALU_FLOAT_3 = 0x98
|
||||
VALU_SCL_TRANS = 0x99
|
||||
SALU_2 = 0x9b
|
||||
SALU_5 = 0x9c
|
||||
add_other_simd(InstOpRDNA4, [("LDS", 0x50, 0x54, 1), ("FLAT", 0x55, 0x59, 2), ("VMEM", 0xbc, 0xdd, 1)])
|
||||
|
||||
class InstOpCDNA(Enum):
|
||||
SMEM_RD = 0
|
||||
SALU_32 = 1
|
||||
VMEM_RD = 2
|
||||
VMEM_WR = 3
|
||||
FLAT_WR = 4
|
||||
VALU_32 = 5
|
||||
LDS = 6
|
||||
PC = 7
|
||||
JUMP = 12
|
||||
NEXT = 13
|
||||
FLAT_RD = 14
|
||||
OTHER_MSG = 15
|
||||
SMEM_WR = 16
|
||||
SALU_64 = 17
|
||||
VALU_64 = 18
|
||||
VALU_MAI = 28
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PACKET TYPE BASE CLASS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class PacketType:
|
||||
"""Base class for SQTT packet types."""
|
||||
encoding: FixedBitField
|
||||
_raw: int
|
||||
_time: int
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
cls._fields = {k: v for k, v in cls.__dict__.items() if isinstance(v, BitField)} # type: ignore[attr-defined]
|
||||
cls._size_nibbles = ((max((f.hi for f in cls._fields.values()), default=0) + 4) // 4) # type: ignore[attr-defined]
|
||||
|
||||
@classmethod
|
||||
def from_raw(cls, raw: int, time: int = 0):
|
||||
inst = object.__new__(cls)
|
||||
inst._raw, inst._time = raw, time
|
||||
return inst
|
||||
|
||||
def __repr__(self) -> str:
|
||||
fields_str = ", ".join(f"{k}={getattr(self, k)}" for k in self._fields if not k.startswith('_') and k != 'encoding') # type: ignore[attr-defined]
|
||||
return f"{self.__class__.__name__}({fields_str})"
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# TS PACKET TYPE DEFINITIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class TS_DELTA_S8_W3(PacketType):
|
||||
encoding = bits[6:0] == 0b0100001
|
||||
delta = bits[10:8]
|
||||
_padding = bits[71:11]
|
||||
|
||||
class TS_DELTA_S5_W3(PacketType):
|
||||
encoding = bits[4:0] == 0b00110
|
||||
delta = bits[7:5]
|
||||
_padding = bits[51:8]
|
||||
|
||||
class TS_DELTA_S5_W3_RDNA4(PacketType): # Layout 4: 52->56 bits
|
||||
encoding = bits[4:0] == 0b00110
|
||||
delta = bits[9:7]
|
||||
_padding = bits[55:10]
|
||||
|
||||
class TS_DELTA_SHORT(PacketType):
|
||||
encoding = bits[3:0] == 0b1000
|
||||
delta = bits[7:4]
|
||||
|
||||
class TS_DELTA_OR_MARK(PacketType):
|
||||
encoding = bits[6:0] == 0b0000001
|
||||
delta = bits[47:12]
|
||||
pl = bits[8:8]
|
||||
rt = bits[9:9]
|
||||
@property
|
||||
def is_marker(self) -> bool: return bool(self.rt and not self.pl)
|
||||
|
||||
class TS_DELTA_OR_MARK_RDNA4(TS_DELTA_OR_MARK):
|
||||
delta = bits[63:12]
|
||||
rt = bits[7:7]
|
||||
pl = bits[8:8]
|
||||
tl = bits[9:9]
|
||||
|
||||
class TS_DELTA_S5_W2(PacketType):
|
||||
encoding = bits[4:0] == 0b11100
|
||||
delta = bits[6:5]
|
||||
_padding = bits[47:7]
|
||||
|
||||
class TS_DELTA_S5_W2_RDNA4(PacketType): # Layout 4: 48->40 bits
|
||||
encoding = bits[4:0] == 0b11100
|
||||
delta = bits[6:5]
|
||||
_padding = bits[39:7]
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PACKET TYPE DEFINITIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class VALUINST(PacketType): # exclude: 1 << 2
|
||||
encoding = bits[2:0] == 0b011
|
||||
delta = bits[5:3]
|
||||
flag = bits[6:6]
|
||||
wave = bits[11:7]
|
||||
|
||||
class VMEMEXEC(PacketType): # exclude: 1 << 0
|
||||
encoding = bits[3:0] == 0b1111
|
||||
delta = bits[5:4]
|
||||
src = bits[7:6].enum(MemSrc)
|
||||
|
||||
class ALUEXEC(PacketType): # exclude: 1 << 1
|
||||
encoding = bits[3:0] == 0b1110
|
||||
delta = bits[5:4]
|
||||
src = bits[7:6].enum(AluSrc)
|
||||
|
||||
class IMMEDIATE(PacketType): # exclude: 1 << 5
|
||||
encoding = bits[3:0] == 0b1101
|
||||
delta = bits[6:4]
|
||||
wave = bits[11:7]
|
||||
|
||||
class IMMEDIATE_MASK(PacketType): # exclude: 1 << 5
|
||||
encoding = bits[4:0] == 0b00100
|
||||
delta = bits[7:5]
|
||||
mask = bits[23:8]
|
||||
|
||||
class WAVERDY(PacketType): # exclude: 1 << 3
|
||||
encoding = bits[4:0] == 0b10100
|
||||
delta = bits[7:5]
|
||||
mask = bits[23:8]
|
||||
|
||||
class WAVEEND(PacketType): # exclude: 1 << 4
|
||||
encoding = bits[4:0] == 0b10101
|
||||
delta = bits[7:5]
|
||||
sa = bits[8:8]
|
||||
simd = bits[10:9]
|
||||
wgp = bits[13:11]
|
||||
wave = bits[19:15]
|
||||
@property
|
||||
def cu(self) -> int: return self.wgp | (self.sa << 3)
|
||||
|
||||
class WAVEEND_RDNA4(PacketType):
|
||||
encoding = bits[4:0] == 0b10101
|
||||
delta = bits[7:5]
|
||||
sa = bits[8:8]
|
||||
simd = bits[10:9]
|
||||
wgp = bits[14:11]
|
||||
wave = bits[19:15]
|
||||
@property
|
||||
def cu(self) -> int: return self.wgp | (self.sa << 4)
|
||||
|
||||
class WAVESTART(PacketType): # exclude: 1 << 4
|
||||
encoding = bits[4:0] == 0b01100
|
||||
delta = bits[6:5]
|
||||
sa = bits[7:7]
|
||||
simd = bits[9:8]
|
||||
wgp = bits[12:10]
|
||||
wave = bits[17:13]
|
||||
id7 = bits[31:18]
|
||||
@property
|
||||
def cu(self) -> int: return self.wgp | (self.sa << 3)
|
||||
|
||||
class WAVESTART_RDNA4(PacketType): # Layout 4: wgp is 4 bits, wave shifted to bits 15-19
|
||||
encoding = bits[4:0] == 0b01100
|
||||
delta = bits[6:5]
|
||||
sa = bits[7:7]
|
||||
simd = bits[9:8]
|
||||
wgp = bits[13:10]
|
||||
wave = bits[19:15]
|
||||
id7 = bits[31:20]
|
||||
@property
|
||||
def cu(self) -> int: return self.wgp | (self.sa << 4)
|
||||
|
||||
class WAVEALLOC(PacketType): # exclude: 1 << 10
|
||||
encoding = bits[4:0] == 0b00101
|
||||
delta = bits[7:5]
|
||||
_padding = bits[19:8]
|
||||
|
||||
class WAVEALLOC_RDNA4(PacketType): # Layout 4: 20->24 bits
|
||||
encoding = bits[4:0] == 0b00101
|
||||
delta = bits[7:5]
|
||||
_padding = bits[23:8]
|
||||
|
||||
class PERF(PacketType): # exclude: 1 << 11
|
||||
encoding = bits[4:0] == 0b10110
|
||||
delta = bits[7:5]
|
||||
arg = bits[27:8]
|
||||
|
||||
class PERF_RDNA4(PacketType): # Layout 4: 28->32 bits
|
||||
encoding = bits[4:0] == 0b10110
|
||||
delta = bits[9:7]
|
||||
arg = bits[31:10]
|
||||
|
||||
class NOP(PacketType):
|
||||
encoding = bits[3:0] == 0b0000
|
||||
delta = None # type: ignore
|
||||
_padding = bits[3:0]
|
||||
|
||||
class TS_WAVE_STATE(PacketType):
|
||||
encoding = bits[6:0] == 0b1010001
|
||||
delta = bits[15:7]
|
||||
coarse = bits[23:16]
|
||||
@property
|
||||
def wave_interest(self) -> bool: return bool(self.coarse & 1)
|
||||
@property
|
||||
def terminate_all(self) -> bool: return bool(self.coarse & 8)
|
||||
|
||||
class EVENT(PacketType): # exclude: 1 << 7
|
||||
encoding = bits[7:0] == 0b01100001
|
||||
delta = bits[10:8]
|
||||
event = bits[23:11]
|
||||
|
||||
class EVENT_BIG(PacketType):
|
||||
encoding = bits[7:0] == 0b11100001
|
||||
delta = bits[10:8]
|
||||
event = bits[31:11]
|
||||
|
||||
class REG(PacketType):
|
||||
encoding = bits[3:0] == 0b1001
|
||||
delta = bits[6:4]
|
||||
slot = bits[9:7]
|
||||
hi_byte = bits[15:8]
|
||||
subop = bits[31:16]
|
||||
val32 = bits[63:32]
|
||||
@property
|
||||
def is_config(self) -> bool: return bool(self.hi_byte & 0x80)
|
||||
|
||||
class SNAPSHOT(PacketType):
|
||||
encoding = bits[6:0] == 0b1110001
|
||||
delta = bits[9:7]
|
||||
snap = bits[63:10]
|
||||
|
||||
class LAYOUT_HEADER(PacketType):
|
||||
encoding = bits[6:0] == 0b0010001
|
||||
delta = None # type: ignore
|
||||
layout = bits[12:7]
|
||||
simd = bits[14:13]
|
||||
group = bits[17:15]
|
||||
sel_a = bits[31:28]
|
||||
sel_b = bits[36:33]
|
||||
flag4 = bits[59:59]
|
||||
_padding = bits[63:60]
|
||||
|
||||
class INST(PacketType):
|
||||
encoding = bits[2:0] == 0b010
|
||||
delta = bits[6:4]
|
||||
flag1 = bits[3:3]
|
||||
flag2 = bits[7:7]
|
||||
wave = bits[12:8]
|
||||
op = bits[19:13].enum(InstOp)
|
||||
|
||||
class INST_RDNA4(PacketType): # Layout 4: different delta position and InstOp encoding
|
||||
encoding = bits[2:0] == 0b010
|
||||
delta = bits[5:3]
|
||||
w64h = bits[6:6]
|
||||
wave = bits[11:7]
|
||||
op = bits[19:12].enum(InstOpRDNA4)
|
||||
|
||||
class UTILCTR(PacketType):
|
||||
encoding = bits[6:0] == 0b0110001
|
||||
delta = bits[8:7]
|
||||
ctr = bits[47:9]
|
||||
|
||||
# Packet types with rocprof type IDs as keys
|
||||
PACKET_TYPES_RDNA3: dict[int, type[PacketType]] = {
|
||||
1: VALUINST, 2: VMEMEXEC, 3: ALUEXEC, 4: IMMEDIATE, 5: IMMEDIATE_MASK, 6: WAVERDY, 7: TS_DELTA_S8_W3, 8: WAVEEND,
|
||||
9: WAVESTART, 10: TS_DELTA_S5_W2, 11: WAVEALLOC, 12: TS_DELTA_S5_W3, 13: PERF, 14: UTILCTR, 15: TS_DELTA_SHORT,
|
||||
16: NOP, 17: TS_WAVE_STATE, 18: EVENT, 19: EVENT_BIG, 20: REG, 21: SNAPSHOT, 22: TS_DELTA_OR_MARK, 23: LAYOUT_HEADER, 24: INST,
|
||||
}
|
||||
PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
|
||||
**PACKET_TYPES_RDNA3,
|
||||
8: WAVEEND_RDNA4, 9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
|
||||
12: TS_DELTA_S5_W3_RDNA4, 13: PERF_RDNA4, 22: TS_DELTA_OR_MARK_RDNA4, 24: INST_RDNA4,
|
||||
}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# CDNA PACKET TYPE DEFINITIONS
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
class CDNA_MISC(PacketType):
|
||||
"""pkt_fmt=0: 16-bit (Misc)"""
|
||||
encoding = bits[3:0] == 0
|
||||
delta = bits[11:4]
|
||||
sh = bits[12:12]
|
||||
misc_type = bits[15:13]
|
||||
|
||||
class CDNA_TIMESTAMP(PacketType):
|
||||
"""pkt_fmt=1: 64-bit timestamp packet (case 0x0)"""
|
||||
encoding = bits[3:0] == 1
|
||||
_reserved = bits[15:4]
|
||||
timestamp = bits[63:16] # stored as (data_word >> 0x10) in low 46 bits of local_58
|
||||
|
||||
class CDNA_REG(PacketType):
|
||||
"""pkt_fmt=2: 64-bit (Reg)"""
|
||||
encoding = bits[3:0] == 2
|
||||
pipe = bits[6:5]
|
||||
_me_raw = bits[8:7]
|
||||
_reserved = bits[15:9]
|
||||
regaddr = bits[31:16]
|
||||
regdata = bits[63:32]
|
||||
|
||||
class CDNA_WAVESTART(PacketType):
|
||||
"""type 3: 32-bit wave start (Wave/group_id)"""
|
||||
encoding = bits[3:0] == 3
|
||||
sh = bits[5:5]
|
||||
cu = bits[9:6]
|
||||
wave = bits[13:10]
|
||||
simd = bits[15:14]
|
||||
pipe = bits[17:16]
|
||||
me = bits[19:18]
|
||||
_reserved = bits[21:20]
|
||||
count = bits[28:22]
|
||||
_padding = bits[31:29]
|
||||
|
||||
class CDNA_WAVEALLOC(PacketType):
|
||||
"""pkt_fmt=4: 16-bit (Wave)"""
|
||||
encoding = bits[3:0] == 4
|
||||
sh = bits[5:5]
|
||||
cu = bits[9:6]
|
||||
wave = bits[13:10]
|
||||
simd = bits[15:14]
|
||||
|
||||
class CDNA_REG_CS(PacketType):
|
||||
"""type 5: 48-bit register CS write (RegCs)"""
|
||||
encoding = bits[3:0] == 5
|
||||
pipe = bits[6:5]
|
||||
_me_raw = bits[8:7]
|
||||
regaddr = bits[15:9]
|
||||
regdata = bits[47:16]
|
||||
|
||||
class CDNA_WAVEEND(PacketType):
|
||||
"""type 6: 16-bit wave end (group_id)"""
|
||||
encoding = bits[3:0] == 6
|
||||
sh = bits[5:5]
|
||||
cu = bits[9:6]
|
||||
wave = bits[13:10]
|
||||
simd = bits[15:14]
|
||||
|
||||
class CDNA_INST(PacketType):
|
||||
"""pkt_fmt=10: 16-bit (MsgInst)"""
|
||||
encoding = bits[3:0] == 10
|
||||
wave = bits[8:5]
|
||||
simd = bits[10:9]
|
||||
op = bits[15:11].enum(InstOpCDNA)
|
||||
|
||||
class CDNA_INST_PC(PacketType):
|
||||
"""pkt_fmt=11: 64-bit (MsgInstPc)"""
|
||||
encoding = bits[3:0] == 11
|
||||
wave = bits[8:5]
|
||||
simd = bits[10:9]
|
||||
_reserved = bits[14:11]
|
||||
err = bits[15:15]
|
||||
pc = bits[63:16]
|
||||
|
||||
class CDNA_ISSUE(PacketType):
|
||||
"""pkt_fmt=13: 32-bit (Issue)"""
|
||||
encoding = bits[3:0] == 13
|
||||
simd = bits[6:5]
|
||||
_gap = bits[7:7]
|
||||
inst0 = bits[9:8]
|
||||
inst1 = bits[11:10]
|
||||
inst2 = bits[13:12]
|
||||
inst3 = bits[15:14]
|
||||
inst4 = bits[17:16]
|
||||
inst5 = bits[19:18]
|
||||
inst6 = bits[21:20]
|
||||
inst7 = bits[23:22]
|
||||
inst8 = bits[25:24]
|
||||
inst9 = bits[27:26]
|
||||
_padding = bits[31:28]
|
||||
|
||||
class CDNA_PERF(PacketType):
|
||||
"""pkt_fmt=14: 64-bit (MsgPerf)"""
|
||||
encoding = bits[3:0] == 14
|
||||
sh = bits[5:5]
|
||||
cu = bits[9:6]
|
||||
cntr_bank = bits[11:10]
|
||||
cntr0 = bits[24:12]
|
||||
cntr1 = bits[37:25]
|
||||
cntr2 = bits[50:38]
|
||||
cntr3 = bits[63:51]
|
||||
|
||||
class CDNA_EVENT(PacketType):
|
||||
"""pkt_fmt=7: 16-bit"""
|
||||
encoding = bits[3:0] == 7
|
||||
_reserved = bits[15:4]
|
||||
|
||||
class CDNA_EVENT_CS(PacketType):
|
||||
"""pkt_fmt=8: 16-bit"""
|
||||
encoding = bits[3:0] == 8
|
||||
_reserved = bits[15:4]
|
||||
|
||||
class CDNA_EVENT_GFX1(PacketType):
|
||||
"""pkt_fmt=9: 16-bit"""
|
||||
encoding = bits[3:0] == 9
|
||||
_reserved = bits[15:4]
|
||||
|
||||
class CDNA_USERDATA(PacketType):
|
||||
"""pkt_fmt=12: 48-bit (UserData)"""
|
||||
encoding = bits[3:0] == 12
|
||||
sh = bits[5:5]
|
||||
cu = bits[9:6]
|
||||
wave = bits[13:10]
|
||||
simd = bits[15:14]
|
||||
data = bits[47:16]
|
||||
|
||||
class CDNA_REG_CS_PRIV(PacketType):
|
||||
"""pkt_fmt=15: 48-bit (RegCs)"""
|
||||
encoding = bits[3:0] == 15
|
||||
pipe = bits[6:5]
|
||||
_me_raw = bits[8:7]
|
||||
regaddr = bits[15:9]
|
||||
regdata = bits[47:16]
|
||||
|
||||
PACKET_TYPES_CDNA: dict[int, type[PacketType]] = {
|
||||
0: CDNA_MISC, 1: CDNA_TIMESTAMP, 2: CDNA_REG, 3: CDNA_WAVESTART, 4: CDNA_WAVEALLOC, 5: CDNA_REG_CS, 6: CDNA_WAVEEND,
|
||||
7: CDNA_EVENT, 8: CDNA_EVENT_CS, 9: CDNA_EVENT_GFX1, 10: CDNA_INST, 11: CDNA_INST_PC, 12: CDNA_USERDATA,
|
||||
13: CDNA_ISSUE, 14: CDNA_PERF, 15: CDNA_REG_CS_PRIV, 16: LAYOUT_HEADER,
|
||||
}
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# DECODER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _build_decode_tables(packet_types: dict[int, type[PacketType]]) -> tuple[dict[int, tuple], bytes]:
|
||||
# Build state table: byte -> opcode. Sort by mask specificity (more bits first), NOP last
|
||||
sorted_types = sorted(packet_types.items(), key=lambda x: (-bin(x[1].encoding.mask).count('1'), x[0] == 16))
|
||||
state_table = bytes(next((op for op, cls in sorted_types if (b & cls.encoding.mask) == cls.encoding.default), 16) for b in range(256))
|
||||
# Build decode info: opcode -> (pkt_cls, nib_count, delta_lo, delta_mask, special_case)
|
||||
# special_case: 0=none, 1=TS_DELTA_OR_MARK (check is_marker), 2=TS_DELTA_SHORT (add 4), 3=CDNA_MISC (*4), 4=CDNA_TIMESTAMP (absolute)
|
||||
_special = {TS_DELTA_OR_MARK: 1, TS_DELTA_OR_MARK_RDNA4: 1, TS_DELTA_SHORT: 2, CDNA_MISC: 3, CDNA_TIMESTAMP: 4}
|
||||
decode_info = {}
|
||||
for opcode, pkt_cls in packet_types.items():
|
||||
delta_field = getattr(pkt_cls, 'delta', None)
|
||||
special = _special.get(pkt_cls, 0)
|
||||
decode_info[opcode] = (pkt_cls, pkt_cls._size_nibbles, delta_field.lo if delta_field else 0, delta_field.mask if delta_field else 0, special) # type: ignore[attr-defined]
|
||||
return decode_info, state_table
|
||||
|
||||
_DECODE_INFO_RDNA3, _STATE_TABLE_RDNA3 = _build_decode_tables(PACKET_TYPES_RDNA3)
|
||||
_DECODE_INFO_RDNA4, _STATE_TABLE_RDNA4 = _build_decode_tables(PACKET_TYPES_RDNA4)
|
||||
_DECODE_INFO_CDNA, _STATE_TABLE_CDNA = _build_decode_tables(PACKET_TYPES_CDNA)
|
||||
|
||||
def decode(data: bytes) -> Iterator[PacketType]:
|
||||
"""Decode raw SQTT blob, yielding packet instances. Auto-detects RDNA (layout 3/4) vs CDNA."""
|
||||
n, reg, pos, nib_off, nib_count, time, ts_offset = len(data), 0, 0, 0, 16, 0, None
|
||||
decode_info, state_table = _DECODE_INFO_RDNA3, _STATE_TABLE_RDNA3 # start RDNA3, auto-detect switches if needed
|
||||
|
||||
while pos + ((nib_count + nib_off + 1) >> 1) <= n:
|
||||
need = nib_count - nib_off
|
||||
# 1. if unaligned, read high nibble to align
|
||||
if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
|
||||
# 2. read all full bytes at once
|
||||
if (byte_count := need >> 1):
|
||||
read_bytes = min(byte_count, 8)
|
||||
chunk = int.from_bytes(data[pos:pos + read_bytes], 'little')
|
||||
reg, pos = (reg >> (read_bytes * 8)) | (chunk << (64 - read_bytes * 8)), pos + byte_count
|
||||
# 3. if odd, read low nibble
|
||||
if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
|
||||
|
||||
opcode = state_table[reg & 0xFF]
|
||||
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
|
||||
delta = (reg >> delta_lo) & delta_mask
|
||||
if special == 1: # TS_DELTA_OR_MARK
|
||||
pkt = pkt_cls.from_raw(reg, 0) # create packet to check is_marker
|
||||
if pkt.is_marker: delta = 0
|
||||
elif special == 2: delta += 4 # TS_DELTA_SHORT
|
||||
elif special == 3: delta *= 4 # CDNA_DELTA
|
||||
elif special == 4: # CDNA_TIMESTAMP (absolute timestamp anchoring)
|
||||
if (reg >> 4) & 0xfff == 0: # unk_0 == 0 means absolute timestamp
|
||||
abs_ts = reg >> 16
|
||||
if ts_offset is None: ts_offset = abs_ts - time
|
||||
else: time = ((abs_ts - ts_offset) & ~3) - 4
|
||||
delta = 0
|
||||
time += delta
|
||||
pkt = pkt_cls.from_raw(reg, time)
|
||||
# auto-detect: first packet is always LAYOUT_HEADER (RDNA layout 3/4) or misdetected (CDNA)
|
||||
if pkt_cls is LAYOUT_HEADER:
|
||||
if pkt.layout == 4: decode_info, state_table = _DECODE_INFO_RDNA4, _STATE_TABLE_RDNA4
|
||||
elif pkt.layout != 3: # not a real LAYOUT_HEADER — switch to CDNA and re-decode first packet
|
||||
decode_info, state_table = _DECODE_INFO_CDNA, _STATE_TABLE_CDNA
|
||||
opcode = state_table[reg & 0xFF]
|
||||
pkt_cls, nib_count, delta_lo, delta_mask, special = decode_info[opcode]
|
||||
if special == 4 and (reg >> 4) & 0xfff == 0: # CDNA_TIMESTAMP absolute
|
||||
ts_offset = (reg >> 16) - time
|
||||
pkt = pkt_cls.from_raw(reg, time)
|
||||
yield pkt
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# MAPPER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InstructionInfo:
|
||||
pc: int
|
||||
wave: int
|
||||
inst: Inst
|
||||
|
||||
def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, InstructionInfo|None]]:
|
||||
"""maps SQTT packets to instructions, yields (packet, instruction_info or None)"""
|
||||
# map pcs to insts
|
||||
from tinygrad.viz.serve import amd_decode
|
||||
pc_map = amd_decode(lib, target)
|
||||
wave_pc:dict[int, int] = {}
|
||||
# only processing packets on one [CU, SIMD] unit
|
||||
def simd_select(p) -> bool: return getattr(p, "cu", 0) == 0 and getattr(p, "simd", 0) == 0
|
||||
for p in decode(data):
|
||||
if not simd_select(p): continue
|
||||
if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART)):
|
||||
assert p.wave not in wave_pc, "only one inflight wave per unit"
|
||||
wave_pc[p.wave] = next(iter(pc_map))
|
||||
elif isinstance(p, (WAVEEND, WAVEEND_RDNA4)):
|
||||
pc = wave_pc.pop(p.wave)
|
||||
yield (p, InstructionInfo(pc, p.wave, s_endpgm()))
|
||||
elif isinstance(p, IMMEDIATE_MASK):
|
||||
# immediate mask may yield multiple times per packet
|
||||
for wave in range(16):
|
||||
if p.mask & (1 << wave):
|
||||
inst = pc_map[pc:=wave_pc[wave]]
|
||||
wave_pc[wave] += inst.size()
|
||||
yield (p, InstructionInfo(pc, wave, inst))
|
||||
# map INST events on this SIMD to the program counter, we know the waves
|
||||
elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)) and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_")):
|
||||
inst = pc_map[pc:=wave_pc[p.wave]]
|
||||
# s_delay_alu, s_wait_alu and s_barrier_wait instructions are skipped
|
||||
while (inst_op:=getattr(inst, 'op_name', '')) in {"S_DELAY_ALU", "S_WAIT_ALU", "S_BARRIER_WAIT"}:
|
||||
wave_pc[p.wave] += inst.size()
|
||||
inst = pc_map[pc:=wave_pc[p.wave]]
|
||||
# assert branch always has a JUMP packet
|
||||
if "BRANCH" in inst_op and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("JUMP")):
|
||||
raise AssertionError(f"{inst_op} can only be followed by JUMP, got {p}")
|
||||
# JUMP handling
|
||||
if isinstance(p, (INST, INST_RDNA4)) and p.op in {InstOp.JUMP, InstOpRDNA4.JUMP}:
|
||||
x = getattr(inst, 'simm16') & 0xffff
|
||||
wave_pc[p.wave] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
|
||||
else:
|
||||
wave_pc[p.wave] += inst.size()
|
||||
yield (p, InstructionInfo(pc, p.wave, inst))
|
||||
# for all other packets (VMEMEXEC, ALUEXEC, OTHER_ INST, etc.), yield with None
|
||||
else: yield (p, None)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# PRINTER
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
PACKET_COLORS = {
|
||||
"INST": "WHITE", "VALUINST": "BLACK", "VMEMEXEC": "yellow", "ALUEXEC": "yellow",
|
||||
"IMMEDIATE": "YELLOW", "IMMEDIATE_MASK": "YELLOW", "WAVERDY": "cyan", "WAVEALLOC": "cyan",
|
||||
"WAVEEND": "blue", "WAVESTART": "blue", "PERF": "magenta", "EVENT": "red", "EVENT_BIG": "red",
|
||||
"REG": "green", "LAYOUT_HEADER": "white", "SNAPSHOT": "white", "UTILCTR": "green",
|
||||
}
|
||||
|
||||
def format_packet(p) -> str:
|
||||
name = type(p).__name__
|
||||
if isinstance(p, (INST, INST_RDNA4)):
|
||||
op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
|
||||
fields = f"wave={p.wave} op={op_name}" + ((" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "") if isinstance(p, INST) else "")
|
||||
elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
|
||||
elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
|
||||
elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
|
||||
elif isinstance(p, (WAVESTART, WAVESTART_RDNA4, WAVEEND, WAVEEND_RDNA4)): fields = f"wave={p.wave} simd={p.simd} cu={p.cu}"
|
||||
elif hasattr(p, '_fields'):
|
||||
filt = {'delta', 'encoding'} if not isinstance(p, (TS_DELTA_OR_MARK, TS_DELTA_OR_MARK_RDNA4)) else {'encoding'}
|
||||
fields = " ".join(f"{k}=0x{getattr(p, k):x}" if k in {'snap', 'val32'} else f"{k}={getattr(p, k)}"
|
||||
for k in p._fields if not k.startswith('_') and k not in filt)
|
||||
else: fields = ""
|
||||
return f"{p._time:8}: {colored(f'{name:18}', PACKET_COLORS.get(name.replace('_RDNA4', ''), 'white'))} {fields}"
|
||||
|
||||
def print_packets(packets) -> None:
|
||||
skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK",
|
||||
"TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"} if not getenv("NOSKIP") else {"NOP"}
|
||||
for data in packets:
|
||||
p, inst = data if isinstance(data, tuple) else (data, None)
|
||||
if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p), f"inst={inst.inst}" if inst is not None else '')
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys, pickle
|
||||
from tinygrad.helpers import temp
|
||||
with open(temp("profile.pkl", append_user=True) if len(sys.argv) < 2 else sys.argv[1], "rb") as f:
|
||||
data = pickle.load(f)
|
||||
prg_events = {e.tag: e for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
|
||||
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
|
||||
dev_targets = {e.device:f"gfx{e.props['gfx_target_version']//1000}" for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.props}
|
||||
evt_num = getenv("SQTT_EVENT", -1)
|
||||
for i, event in enumerate(sqtt_events):
|
||||
prg = prg_events.get(event.kern)
|
||||
print(f"=== event {i} {prg.name if prg is not None else ''} ===")
|
||||
if evt_num == -1 or i == evt_num:
|
||||
print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))
|
||||
print("\n")
|
||||
Reference in New Issue
Block a user