forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Prebuilt Release @ ab07000
This commit is contained in:
IQ.Lvbs history cleanup
90
tinygrad_repo/tinygrad/renderer/__init__.py
Normal file
90
tinygrad_repo/tinygrad/renderer/__init__.py
Normal file
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
from typing import Callable, cast
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.helpers import prod, Target, EMULATED_DTYPES
|
||||
from tinygrad.uop.ops import Ops, UOp, sint, ssimplify, smin, GroupOp, PatternMatcher
|
||||
from tinygrad.dtype import AddrSpace, PtrDType, DType, dtypes
|
||||
from tinygrad.codegen.opt.tc import TensorCore
|
||||
from tinygrad.device import Compiler
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Estimates:
|
||||
# number of FLOPS used in the Kernel
|
||||
ops:sint = 0
|
||||
# bytes accessed in loads and stores
|
||||
lds:sint = 0
|
||||
# total bytes accessed, counting only once for bytes that are accessed multiple times
|
||||
mem:sint = 0
|
||||
def __add__(self, o:Estimates): return Estimates(self.ops + o.ops, self.lds + o.lds, self.mem + o.mem)
|
||||
def simplify(self): return Estimates(ssimplify(self.ops), ssimplify(self.lds), ssimplify(self.mem))
|
||||
@staticmethod
|
||||
def from_uops(uops:tuple[UOp, ...], ignore_indexing=False) -> Estimates:
|
||||
flops: sint = 0
|
||||
lds: sint = 0
|
||||
mem: dict[tuple[UOp, Ops], sint] = {}
|
||||
mults: sint = 1
|
||||
mult_stack: list[sint] = []
|
||||
dont_count: set[UOp] = set()
|
||||
if ignore_indexing:
|
||||
def range_gate(x): return x.op is not Ops.RANGE
|
||||
for u in uops:
|
||||
if u.op in {Ops.LOAD, Ops.STORE}:
|
||||
# if u.src[0] is INDEX, we have to include the buffer since it might be an AFTER
|
||||
dont_count = dont_count.union((UOp.sink(*u.src[0].src[1:]) if u.src[0].op is Ops.INDEX else u.src[0]).toposort(range_gate))
|
||||
# TODO: is this correct? this all needs to be cleaned up
|
||||
if len(u.src) > 2: dont_count = dont_count.union(u.src[2].toposort())
|
||||
elif u.op is Ops.IF:
|
||||
dont_count = dont_count.union(u.src[0].toposort())
|
||||
for u in uops:
|
||||
if u.op in {Ops.LOAD, Ops.STORE}:
|
||||
buf = u
|
||||
while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
|
||||
if buf.op is Ops.PARAM:
|
||||
# u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
|
||||
accessed = mem.get((buf, u.op), 0) + u.src[0].dtype.base.itemsize * mults
|
||||
mem[(buf, u.op)] = smin(accessed, buf.ptrdtype.nbytes()) if buf.ptrdtype.size != -1 else accessed
|
||||
if u.op is Ops.RANGE:
|
||||
mult_stack.append(mults)
|
||||
mults *= cast(sint, u.src[0].ssimplify())
|
||||
# SPECIAL are already counted in mults
|
||||
mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults
|
||||
elif u.op is Ops.END: mults = mult_stack.pop(-1)
|
||||
elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
|
||||
elif u.op is Ops.DEFINE_VAR and u.arg[0] == 'core_id': mults *= u.arg[2] + 1
|
||||
elif u.op is Ops.LOAD and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
|
||||
lds += u.dtype.itemsize * mults
|
||||
elif u.op is Ops.STORE and (not isinstance(u.src[0].dtype, PtrDType) or u.src[0].dtype.addrspace != AddrSpace.REG):
|
||||
lds += u.src[1].dtype.itemsize * mults
|
||||
elif u.op in GroupOp.ALU and u not in dont_count: flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.dtype.count
|
||||
elif u.op is Ops.WMMA and u not in dont_count: flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
|
||||
return Estimates(flops, lds, sum(mem.values()))
|
||||
|
||||
class Renderer:
|
||||
target: Target
|
||||
suffix: str = ""
|
||||
# TODO: make this generic with a list of supported types
|
||||
supports_float4: bool = True
|
||||
has_local: bool = True
|
||||
has_threads: bool = False
|
||||
has_shared: bool = True
|
||||
has_aux: bool = False # additional program info, eg. image shapes
|
||||
# NOTE: these two should be in (x,y,z) order to match the max_sizes argument in get_grouped_dims
|
||||
global_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
|
||||
local_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now
|
||||
global_prod_max: tuple[int, ...]|None = None
|
||||
shared_max: int = 32768
|
||||
tensor_cores: list[TensorCore] = []
|
||||
pre_matcher: PatternMatcher|None = None
|
||||
extra_matcher: PatternMatcher|None = None
|
||||
code_for_op: dict[Ops, Callable] = {}
|
||||
|
||||
compiler: Compiler = Compiler()
|
||||
|
||||
def __init__(self, target:Target): self.target = target
|
||||
def __reduce__(self): return self.__class__, (self.target,)
|
||||
def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")
|
||||
def asm(self, prg:UOp, lin:UOp) -> bytes: raise NotImplementedError("needs an assembler")
|
||||
def aux(self, uops:list[UOp]) -> dict: raise NotImplementedError("needs aux")
|
||||
def supported_dtypes(self) -> set[DType]:
|
||||
# double can't be bitcast to anything without long support
|
||||
return set(dtypes.all) - {dtypes.weakint} - ({dtypes.double} if dtypes.long in EMULATED_DTYPES.tolist(dtypes) else set())
|
||||
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")
|
||||
597
tinygrad_repo/tinygrad/renderer/cstyle.py
Normal file
597
tinygrad_repo/tinygrad/renderer/cstyle.py
Normal file
@@ -0,0 +1,597 @@
|
||||
from typing import Literal, Callable, cast
|
||||
import math, sys, struct
|
||||
from collections import defaultdict, Counter
|
||||
from tinygrad.codegen.opt import tc
|
||||
from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters
|
||||
from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT, IMAGE, FLOAT16
|
||||
from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate, float_to_bf16
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.codegen.late.devectorizer import no_vectorized_alu
|
||||
|
||||
|
||||
base_rewrite = PatternMatcher([
|
||||
(UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
|
||||
(UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
|
||||
(UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
|
||||
(UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
|
||||
# r method accesses
|
||||
(UPat(Ops.RANGE, name="x"),
|
||||
lambda ctx,x: f"for ({ctx.render_dtype(x.dtype)} {ctx[x]} = 0; {ctx[x]} < {ctx[x.src[0]]}; {ctx[x]}++) {{"),
|
||||
(UPat(Ops.STACK, name="x"),
|
||||
lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_dtype(x.dtype))}" + \
|
||||
f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x:
|
||||
f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_dtype(x.dtype)})" if x.dtype.count > 1 and not isinstance(x.dtype, PtrDType) else None),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, ctx[x.src[0]])})"),
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x:
|
||||
f"__builtin_bit_cast({ctx.render_dtype(x.dtype)}, ({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
|
||||
(UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
|
||||
(UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
|
||||
(UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"),
|
||||
# const
|
||||
(UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x.dtype, ctx.infinity)})"),
|
||||
(UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x.dtype, f'-{ctx.infinity}')})"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, ctx.nan)})" if math.isnan(x.arg) else None),
|
||||
(UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.arg}f"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.arg}ll"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}ull"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}u"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.arg else "0"),
|
||||
# consts are rendered to larger type and casted
|
||||
(UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, f'{x.arg}f')})"),
|
||||
(UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, f'{x.arg}u')})"),
|
||||
(UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x.dtype, str(x.arg))})"),
|
||||
# default const render
|
||||
(UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
|
||||
# new load/store
|
||||
(UPat.var("buf").index(UPat.var('idx')), lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
|
||||
(UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
|
||||
(UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
|
||||
(UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
|
||||
# alu/gep
|
||||
# TODO: look for left-associative
|
||||
(UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](
|
||||
*([strip_parens(ctx[v]) if v.op == x.op and x.op in {Ops.ADD, Ops.MUL, Ops.XOR, Ops.OR, Ops.AND} else ctx[v] for v in x.src]), x.dtype)),
|
||||
(UPat(Ops.GEP, name="x"), lambda ctx,x: ctx[x.src[0]] + \
|
||||
(f"[{x.arg[0]}]" if x.src[0].dtype.count > ctx.gep_arr_threshold else f".{'xyzwabcd'[x.arg[0]]}")),
|
||||
# custom passes through with format
|
||||
(UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg.format(*[ctx[y] for y in x.src])),
|
||||
])
|
||||
|
||||
extra_pm = PatternMatcher([
|
||||
# devectorize any bools
|
||||
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.INDEX), dtype=dtypes.bool, name="alu"), no_vectorized_alu),
|
||||
# CAST (from bool) can't be vectorized
|
||||
(UPat(Ops.CAST, src=(UPat(dtype=dtypes.bool),), name="alu"), no_vectorized_alu),
|
||||
# WHERE can't be vectorized
|
||||
(UPat(Ops.WHERE, name="alu"), no_vectorized_alu),
|
||||
])
|
||||
|
||||
def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
|
||||
patterns = PatternMatcher([
|
||||
(UPat(Ops.WHERE, src=(UPat.var("b"), UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))),
|
||||
lambda b,x,y: UOp(Ops.WHERE, dtype=dtypes.float, src=(b,x.cast(dtypes.float),y.cast(dtypes.float))).cast(x.dtype)),
|
||||
(UPat(GroupOp.ALU, dtype=dts, name="x"),
|
||||
lambda x: UOp(x.op, dtypes.float, tuple(vv.cast(dtypes.float) for vv in x.src), x.arg).cast(x.dtype)),
|
||||
(UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))),
|
||||
lambda alu,x,y: UOp(alu.op, dtypes.bool, (x.cast(dtypes.float), y.cast(dtypes.float)), alu.arg))])
|
||||
if casting:
|
||||
# add float intermediate casting
|
||||
patterns += PatternMatcher([
|
||||
(UPat(Ops.CAST, dts, (UPat.var("x"),), name="y"), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=dtypes.float else None),
|
||||
(UPat(Ops.CAST, name="x", src=(UPat.var("y", dts),)), lambda x,y: y.cast(dtypes.float).cast(x.dtype) if x.dtype!=dtypes.float else None)])
|
||||
return patterns
|
||||
|
||||
def cast_float_to_bf16(x: UOp) -> UOp:
|
||||
assert x.dtype == dtypes.float, "cast float -> bf16 must start with float"
|
||||
x = x.bitcast(dtypes.uint)
|
||||
x = (-x & 0x7f800000).ne(0).where(x + ((x >> 16) & 1) + 0x7fff, (x & 0xffff).ne(0).where((x | 0x10000), x))
|
||||
return (x >> 16).cast(dtypes.ushort).bitcast(dtypes.bfloat16)
|
||||
|
||||
# manual bfloat16 casting patterns (shared between LLVM, Clang, and AMD renderers to avoid compiler intrinsics)
|
||||
pm_manual_bf16_cast = PatternMatcher([
|
||||
(UPat(Ops.CAST, dtypes.float, (UPat.var("x", dtypes.bfloat16),)),
|
||||
lambda x: (x.bitcast(dtypes.ushort).cast(dtypes.uint)<<16).bitcast(dtypes.float)),
|
||||
(UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat.var("x", dtype=dtypes.float),)), cast_float_to_bf16),
|
||||
])
|
||||
|
||||
def uops_to_dtypes(uops:list[UOp]) -> list[DType]: return dedup(u.dtype for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType)))
|
||||
|
||||
# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
|
||||
def wmma_args(uops:list[UOp]):
|
||||
return dedup((uop.arg[0], uop.arg[1], uop.arg[2], uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA)
|
||||
|
||||
class CStyleLanguage(Renderer):
|
||||
kernel_typedef: str = "void"
|
||||
buffer_prefix: str = ""
|
||||
buffer_suffix: str = ""
|
||||
smem_align: str = ""
|
||||
smem_prefix: str = ""
|
||||
smem_prefix_for_cast: bool = True
|
||||
arg_int_prefix: str = "const int"
|
||||
barrier: str = ""
|
||||
code_for_workitem: dict[Literal["g", "l", "i"], Callable] = {}
|
||||
extra_args: list[str] = []
|
||||
float4: str|None = None
|
||||
float4_style: tuple[str, str] = ('(', ')')
|
||||
gep_arr_threshold: int = 4
|
||||
type_map: dict[DType, str] = {}
|
||||
infinity: str = "INFINITY"
|
||||
nan: str = "NAN"
|
||||
code_for_op: dict = {
|
||||
Ops.SQRT: lambda x,dtype: f"sqrt({x})", Ops.RECIPROCAL: lambda x,dtype: f"(1/{x})", Ops.NEG: lambda x,dtype: f"-{x}",
|
||||
Ops.EXP2: lambda x,dtype: f"exp2({x})", Ops.LOG2: lambda x,dtype: f"log2({x})", Ops.SIN: lambda x,dtype: f"sin({x})",
|
||||
Ops.TRUNC: lambda x,dtype: f"trunc({x})",
|
||||
Ops.AND: lambda a,b,dtype: f"({a}&{b})", Ops.XOR: lambda a,b,dtype: f"({a}^{b})", Ops.OR: lambda a,b,dtype: f"({a}|{b})",
|
||||
Ops.ADD: lambda a,b,dtype: f"({a}+{b})", Ops.SUB: lambda a,b,dtype: f"({a}-{b})", Ops.MUL: lambda a,b,dtype: f"({a}*{b})",
|
||||
Ops.CMOD: lambda a,b,dtype: f"({a}%{b})", Ops.CDIV: lambda a,b,dtype: f"({a}/{b})", Ops.CMPNE: lambda a,b,dtype: f"({a}!={b})",
|
||||
Ops.SHR: lambda a,b,dtype: f"({a}>>{b})", Ops.SHL: lambda a,b,dtype: f"({a}<<{b})", Ops.CMPLT: lambda a,b,dtype: f"({a}<{b})",
|
||||
Ops.WHERE: lambda a,b,c,dtype: f"({a}?{b}:{c})", Ops.CMPEQ: lambda a,b,dtype: f"({a}=={b})"}
|
||||
|
||||
string_rewrite = base_rewrite
|
||||
extra_matcher = extra_pm
|
||||
|
||||
def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
|
||||
tmp = ""
|
||||
if any(isinstance(dtype, ImageDType) for _,(dtype,_) in bufs):
|
||||
tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
|
||||
buftypes = [(name, self.render_dtype(dtype, mutable)+self.buffer_suffix if isinstance(dtype, (ImageDType, PtrDType)) else
|
||||
self.arg_int_prefix if dtype == dtypes.int else None) for name,(dtype,mutable) in bufs]
|
||||
local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
|
||||
launch_bounds = prod([d.vmax for d in local_dims])
|
||||
prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] +
|
||||
[', '.join([f'{t} {name}' for name,t in buftypes] + self.extra_args)] +
|
||||
[") {\n" + tmp] + ['\n'.join(kernel), "\n}"])
|
||||
return prg if prefix is None else "\n".join(prefix)+f"\n{prg}"
|
||||
|
||||
def render_cast(self, dt:DType, val: str) -> str: return f"({self.render_dtype(dt)})({val})"
|
||||
def render_dtype(self, dt:DType, mutable=True) -> str:
|
||||
if isinstance(dt, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t"
|
||||
if isinstance(dt, PtrDType):
|
||||
prefix = ""
|
||||
if dt.addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix
|
||||
if dt.addrspace == AddrSpace.GLOBAL: prefix = self.buffer_prefix
|
||||
return prefix + self.render_dtype(dt.base) + "*"
|
||||
if dt.count > 1: return self.type_map.get(scalar:=dt.scalar(), scalar.name).replace(" ", "_") + str(dt.count)
|
||||
return self.type_map.get(scalar:=dt.scalar(), scalar.name)
|
||||
|
||||
def __getitem__(self, key): return self.r[key] # hacky helper
|
||||
def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[DType,bool]]]]:
|
||||
r: dict[UOp, str] = {}
|
||||
self.r = r
|
||||
|
||||
child_count = Counter(v for ru in uops for v in ru.src)
|
||||
# find which PARAMs are stored to with a single toposort
|
||||
writable_params = {u for u in UOp.sink(*[u.src[0] for u in uops if u.op is Ops.STORE]).toposort(lambda u: u.op != Ops.END) if u.op is Ops.PARAM}
|
||||
bufs: dict[UOp, tuple[str, tuple[DType, bool]]] = {}
|
||||
kernel = []
|
||||
depth = 1
|
||||
c: defaultdict[str, int] = defaultdict(int)
|
||||
name = "test"
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP}: continue
|
||||
if u.op is Ops.AFTER:
|
||||
r[u] = r[u.src[0]]
|
||||
continue
|
||||
if u.op is Ops.SINK:
|
||||
if u.arg is not None: name = u.arg.function_name
|
||||
continue
|
||||
if u.op in (Ops.PARAM, Ops.DEFINE_VAR):
|
||||
if u.op is not Ops.PARAM: r[u] = u.arg[0]
|
||||
elif isinstance(u.dtype, ImageDType): r[u] = f"data{u.arg.slot}_{u.dtype.shape[0]}x{u.dtype.shape[1]}"
|
||||
else: r[u] = f"data{u.arg.slot}_{sz}" if (sz:=u.max_numel()) > 0 else f"data{u.arg.slot}"
|
||||
bufs[u] = (r[u], (u.dtype, u in writable_params))
|
||||
continue
|
||||
|
||||
# naming
|
||||
prefix = None
|
||||
if u.op is Ops.SPECIAL: r[u] = u.arg
|
||||
elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u)
|
||||
else:
|
||||
prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
|
||||
Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.STACK: "cast",
|
||||
Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
|
||||
r[u] = f"{prefix}{c[prefix]}"
|
||||
|
||||
l = cast(str, self.string_rewrite.rewrite(u, ctx=self))
|
||||
assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}"
|
||||
|
||||
if u.op in {Ops.ENDIF, Ops.END}: depth -= 1
|
||||
if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \
|
||||
(u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \
|
||||
(u.op is Ops.CAST and isinstance(u.dtype, PtrDType)) or \
|
||||
(u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
|
||||
r[u] = l
|
||||
else:
|
||||
if u.op not in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG} and u.dtype != dtypes.void:
|
||||
l = f"{self.render_dtype(u.dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
|
||||
kernel.append(" "*depth + l)
|
||||
if prefix: c[prefix] += 1 # if it was used, increment
|
||||
if u.op in {Ops.IF, Ops.RANGE}: depth += 1
|
||||
del self.r
|
||||
|
||||
# NOTE: this relies on bufs dict preserving order
|
||||
return (name, kernel, list(bufs.values()))
|
||||
def render(self, uops:list[UOp]) -> str: return self.render_kernel(*self._render(uops), uops)
|
||||
|
||||
class ClangRenderer(CStyleLanguage):
|
||||
float4 = "(float4)"
|
||||
float4_style = ('{', '}')
|
||||
gep_arr_threshold = 0
|
||||
has_local = False
|
||||
has_threads = bool(getenv("THREADS", 1))
|
||||
global_max = (CPU_COUNT.value, 0, 0)
|
||||
infinity = "__builtin_inff()"
|
||||
nan = '__builtin_nanf("")'
|
||||
|
||||
# language options
|
||||
buffer_suffix = " restrict"
|
||||
type_map = {dtypes.bool:"_Bool", dtypes.half:"__fp16"}
|
||||
code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC, Ops.RECIPROCAL]}),
|
||||
Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
|
||||
Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})",
|
||||
Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"}
|
||||
|
||||
# LLVM legalizes double => half/bf16 cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
|
||||
# there is also no native bfl16 <-> fp16 conversion on those CPUs
|
||||
extra_matcher = PatternMatcher([(UPat.var("x", dtypes.float64).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
|
||||
(UPat.var("x", dtypes.float64).cast(dtypes.bfloat16), lambda x: x.cast(dtypes.float32).cast(dtypes.bfloat16)),
|
||||
(UPat.var("x", dtypes.bfloat16).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16)),
|
||||
(UPat((Ops.SQRT, Ops.TRUNC), name="alu"), no_vectorized_alu)]) + create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast + \
|
||||
CStyleLanguage.extra_matcher
|
||||
|
||||
if sys.platform == 'win32':
|
||||
kernel_typedef = "__attribute__((ms_abi)) void"
|
||||
def render_vector_prefix(self, dt:DType) -> str:
|
||||
# round (down) to power of two (this is actually the default clang behavior)
|
||||
alignment = 2**int(math.log2(dt.itemsize)) if getenv("ALIGNED", 1) and not dtypes.is_bool(dt) else 1
|
||||
return f"typedef {self.render_dtype(dt.scalar())} {self.render_dtype(dt)} __attribute__((aligned({alignment}),ext_vector_type({dt.count})));"
|
||||
|
||||
def _render_defines(self, uops) -> list[str]:
|
||||
prefix = [self.render_vector_prefix(dt) for dt in uops_to_dtypes(uops) if dt.count > 1]
|
||||
# https://github.com/corsix/amx
|
||||
for name, (N, M, _), dtype_in, _, _, _, _, _ in wmma_args(uops):
|
||||
prefix += [
|
||||
'#define AMX_SET(imm5) __asm("nop\\nnop\\nnop\\n.word (0x201000+(%0<<5)+%1)" : : "i"(17), "i"(imm5) : "memory")',
|
||||
'#define AMX(op, gpr, btf) __asm(".word (0x201000+(%0 << 5)+0%1-((0%1>>4)*6))" : : "i"(op), "r"((unsigned long long)(gpr)+(btf)) : "memory")',
|
||||
]
|
||||
# 'static' in C roughly means that function symbol isn't exported. LLVM puts those symbols at the end of object file which allows Clang JIT
|
||||
# to just jump at the start of a shellcode without having to deal with symbols or trampolines at all. This is better than having to inline
|
||||
# wmma function every time it is called or wasting complexity on a symbol parsing and a memory page on trampoline.
|
||||
out, dt1, dt2 = self.render_dtype(dtype_in.vec(N*N)), self.render_dtype(dtype_in.vec(N)), self.render_dtype(dtype_in.vec(M))
|
||||
prefix += [f"""static {out} __{name}({dt1} data1, {dt2} data2, {out} data0){{
|
||||
AMX_SET(0);\n for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(4, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }}
|
||||
AMX(0, (int *)(&data2), 0ull<<62); AMX(1, (int *)(&data1), 0ull<<62); AMX(12, 0, 0ull);
|
||||
for(int ridx0 = 0; ridx0 < 16; ridx0++){{ AMX(5, (int *)(&data0), 0ull<<62 | (ridx0*4ull)<<56 | ridx0*64ull); }}
|
||||
AMX_SET(1);\n return data0;\n}}"""]
|
||||
return prefix
|
||||
def _render_body(self, function_name, kernel, bufs, uops, pref=None) -> str: return super().render_kernel(function_name, kernel, bufs, uops, pref)
|
||||
def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[DType,bool]]]) -> str: return ""
|
||||
|
||||
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
|
||||
defines = '\n'.join(self._render_defines(uops))
|
||||
return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs)
|
||||
|
||||
def supported_dtypes(self):
|
||||
return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or self.target.arch.startswith(("x86", "arm"))) and d not in dtypes.fp8s}
|
||||
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_cpu import ClangCompiler
|
||||
if "AMX" in target.arch: self.tensor_cores = tc.amx
|
||||
self.compiler = ClangCompiler([x for x in target.arch.split(",") if x != "AMX"])
|
||||
|
||||
class OpenCLRenderer(CStyleLanguage):
|
||||
has_aux = True
|
||||
|
||||
# language options
|
||||
kernel_typedef = "__kernel void"
|
||||
buffer_prefix = "__global "
|
||||
smem_align = "__attribute__ ((aligned (16))) "
|
||||
smem_prefix = "__local "
|
||||
barrier = "barrier(CLK_LOCAL_MEM_FENCE);"
|
||||
float4 = "(float4)"
|
||||
code_for_workitem = {"g": lambda x: f"get_group_id({x})", "l": lambda x: f"get_local_id({x})", "i": lambda x: f"get_global_id({x})"}
|
||||
type_map = { dtypes.int8: "char", dtypes.uint8: "uchar", dtypes.uint32: "uint", dtypes.uint16: "ushort", dtypes.uint64: "ulong",
|
||||
dtypes.bfloat16: "ushort" }
|
||||
extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast + extra_pm
|
||||
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
|
||||
# bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
|
||||
(UPat(Ops.CONST, dtypes.bfloat16, name="x"),
|
||||
lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"),
|
||||
# load/store image (OpenCL)
|
||||
(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"),
|
||||
(UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))),
|
||||
lambda ctx,buf,idx_y,idx_x,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]})):{ctx[var]})"),
|
||||
(UPat(Ops.LOAD, dtype=dtypes.float.vec(4), src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)),
|
||||
lambda ctx,buf,idx_y,idx_x: f"read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]}))"),
|
||||
(UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),
|
||||
UPat.var("var", dtypes.float.vec(4)))),
|
||||
lambda ctx,buf,idx_y,idx_x,var: f"write_imagef({ctx[buf]}, (int2)({ctx[idx_x]},{ctx[idx_y]}), {ctx[var]});"),
|
||||
]) + base_rewrite
|
||||
|
||||
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
|
||||
if any(uop.dtype.base == dtypes.half for uop in uops): prefix = (["#pragma OPENCL EXTENSION cl_khr_fp16 : enable"] + (prefix or []))
|
||||
return super().render_kernel(function_name, kernel, bufs, uops, prefix)
|
||||
|
||||
def aux(self, uops:list[UOp]):
|
||||
arg_dtypes:list[list[tuple[int, DType]]] = []
|
||||
for i,u in enumerate(u for u in uops if u.op is Ops.PARAM):
|
||||
while len(arg_dtypes) <= u.arg.slot: arg_dtypes.append([])
|
||||
arg_dtypes[u.arg.slot].append((i, u.dtype))
|
||||
return tuple(tuple(a) for a in arg_dtypes),
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes()
|
||||
if (d != dtypes.half or "cl_khr_fp16" in self.target.arch) and
|
||||
(d != dtypes.double or "cl_khr_fp64" in self.target.arch) and d not in dtypes.fp8s}
|
||||
|
||||
class IntelRenderer(OpenCLRenderer):
|
||||
suffix, kernel_typedef = "INTEL", "__attribute__((intel_reqd_sub_group_size(8)))\n" + "__kernel void"
|
||||
tensor_cores = tc.intel
|
||||
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat.var('x', dtype=dtypes.float),)), lambda ctx,x: f"intel_convert_bfloat16_as_ushort({ctx[x]})"),
|
||||
(UPat(Ops.CAST, dtype=dtypes.float, src=(UPat.var('x', dtype=dtypes.bfloat16),)), lambda ctx,x: f"intel_convert_as_bfloat16_float({ctx[x]})"),
|
||||
]) + OpenCLRenderer.string_rewrite
|
||||
|
||||
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
|
||||
prefix = []
|
||||
for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops):
|
||||
dt_in = ("ushort", "bf16") if dtype_in == dtypes.bfloat16 else (dtype_in.name, "f16")
|
||||
prefix.append(f"""{dtype_out.name}8 __{name}({dt_in[0]}16 a, {dt_in[0]}16 b, {dtype_out.name}8 c) {{
|
||||
return intel_sub_group_{dt_in[1]}_{dt_in[1]}_matrix_mad_k16(as_int8(a), as_int8(b), c);\n}}""")
|
||||
return super().render_kernel(function_name, kernel, bufs, uops, prefix or None)
|
||||
|
||||
class MetalRenderer(CStyleLanguage):
|
||||
shared_max = 32768
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.ops_metal import MetalCompiler
|
||||
self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if target.arch.startswith("Apple") and int(target.arch[5:]) >= 7 else []
|
||||
|
||||
# language options
|
||||
kernel_typedef = "kernel void"
|
||||
buffer_prefix = "device "
|
||||
smem_prefix = "threadgroup __attribute__((aligned(16))) "
|
||||
arg_int_prefix = "constant int&"
|
||||
barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);"
|
||||
float4 = "float4"
|
||||
code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"}
|
||||
# uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]`
|
||||
extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]']
|
||||
type_map = {dtypes.bfloat16: "bfloat"}
|
||||
|
||||
# precise::sin
|
||||
code_for_op = {**CStyleLanguage.code_for_op, Ops.SIN: lambda x,dtype: f"precise::sin({x})"}
|
||||
|
||||
# upcast to float32 all the ops that don't support bfloat16
|
||||
extra_matcher = PatternMatcher([
|
||||
# NOTE: this is copied from PTX
|
||||
(UPat((Ops.SQRT, Ops.EXP2, Ops.LOG2, Ops.SIN), dtype=dtypes.bfloat16, name="x"),
|
||||
lambda x: (UOp(x.op, dtypes.float, tuple(vv.cast(dtypes.float) for vv in x.src), x.arg).cast(dtypes.bfloat16))),
|
||||
]) + extra_pm
|
||||
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
|
||||
]) + base_rewrite
|
||||
|
||||
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
|
||||
prefix = ["#include <metal_stdlib>","using namespace metal;"]
|
||||
deduped_wmma_args = dedup([(name, dtype_in, dtype_out) for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops)])
|
||||
for name, dtype_in, dtype_out in deduped_wmma_args: prefix.append(
|
||||
f"""{(dstr_out:=self.render_dtype(dtype_out.vec(2)))} __{name}({(dstr_in:=self.render_dtype(dtype_in.vec(2)))} a, {dstr_in} b, {dstr_out} c){{
|
||||
simdgroup_{self.render_dtype(dtype_in)}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(dtype_out)}8x8 mat_c;
|
||||
mat_a.thread_elements()[0] = a[0]; mat_b.thread_elements()[0] = b[0]; mat_c.thread_elements()[0] = c[0];
|
||||
mat_a.thread_elements()[1] = a[1]; mat_b.thread_elements()[1] = b[1]; mat_c.thread_elements()[1] = c[1];
|
||||
simdgroup_multiply_accumulate(mat_c, mat_a, mat_b, mat_c);\n return {dstr_out}(mat_c.thread_elements()[0], mat_c.thread_elements()[1]);\n}}""")
|
||||
return super().render_kernel(function_name, kernel, bufs, uops, prefix)
|
||||
|
||||
def supported_dtypes(self):
|
||||
return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or ((arch:=self.target.arch).startswith("Apple") and int(arch[5:]) >= 6))
|
||||
and d not in dtypes.fp8s+(dtypes.double,)}
|
||||
|
||||
_nms = list("xyzwabcdefghijkl") + [f'v{i}' for i in range(16, 32)]
|
||||
|
||||
class CUDARenderer(CStyleLanguage):
|
||||
global_max = (2147483647, 65535, 65535)
|
||||
local_max = (1024, 1024, 64)
|
||||
shared_max = 49152
|
||||
|
||||
def __init__(self, target:Target, use_nvcc=False):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_cuda import NVRTCCompiler, NVCCCompiler
|
||||
iface, dev, arch = target.interface, target.device, target.arch
|
||||
self.compiler = (NVCCCompiler if use_nvcc else NVRTCCompiler)(arch, ptx=iface.startswith("MOCK") or dev == "CUDA", cache_key=dev.lower())
|
||||
self.tensor_cores = tc.get_cuda(arch)
|
||||
|
||||
# language options
|
||||
# https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
|
||||
kernel_typedef = "extern \"C\" __global__ void __launch_bounds__({launch_bounds})"
|
||||
smem_prefix = "__shared__ __align__(16) "
|
||||
smem_prefix_for_cast = False
|
||||
barrier = "__syncthreads();"
|
||||
float4 = "make_float4"
|
||||
gep_arr_threshold = 8
|
||||
code_for_workitem = {"g": lambda x: f"blockIdx.{chr(120+int(x))}", "l": lambda x: f"threadIdx.{chr(120+int(x))}",
|
||||
"i": lambda x: f"(blockIdx.{chr(120+int(x))}*blockDim.{chr(120+int(x))}+threadIdx.{chr(120+int(x))})"}
|
||||
code_for_op = { **CStyleLanguage.code_for_op,
|
||||
Ops.TRUNC: lambda x,dtype: f"htrunc({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"trunc({x})",
|
||||
Ops.SIN: lambda x,dtype: f"hsin({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sin({x})",
|
||||
Ops.LOG2: lambda x,dtype: f"hlog2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"log2({x})",
|
||||
Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})",
|
||||
Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})",
|
||||
Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" }
|
||||
type_map = {dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"}
|
||||
extra_matcher = create_non_native_float_pats(dtypes.fp8s, casting=False) + PatternMatcher([
|
||||
(UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None),
|
||||
]) + extra_pm
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"tg_bitcast<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
|
||||
]) + base_rewrite
|
||||
|
||||
def render_vector_prefix(self, dt:DType) -> str:
|
||||
vec, scal = self.render_dtype(dt), self.render_dtype(dt.scalar()),
|
||||
elems, header = ', '.join(_nms[:dt.count]), ', '.join([f"{scal} {x}" for x in _nms[:dt.count]])
|
||||
return f"struct __align__({dt.itemsize}) {vec} {{ {scal} {elems}; }}; __device__ {vec} make_{vec}({header}) {{ {vec} r={{{elems}}}; return r; }}"
|
||||
|
||||
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
|
||||
# TODO: why is dtypes.bfloat16.name == "__bf16"? would be easier not override dtypes.name
|
||||
prefix = ["#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))",
|
||||
"template <class T, class F> __device__ __forceinline__ T tg_bitcast(F v) { union U { F f; T t; }; U u; u.f = v; return u.t; }"]
|
||||
used_dtypes = uops_to_dtypes(uops)
|
||||
if any(dt.scalar() in dtypes.fp8s for dt in used_dtypes): prefix.append("#include <cuda_fp8.h>")
|
||||
if any(dt.scalar() == dtypes.half for dt in used_dtypes): prefix.append("#include <cuda_fp16.h>")
|
||||
if any(dt.scalar() == dtypes.bfloat16 for dt in used_dtypes): prefix.append("#include <cuda_bf16.h>")
|
||||
prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if (dt.count in (4,8) and dt.scalar() in {dtypes.half, dtypes.bfloat16})
|
||||
or (dt.count in (2,4,8,16) and dt.scalar() in dtypes.fp8s)]
|
||||
dt_map_in = { dtypes.float: "tf32", dtypes.half: "f16", dtypes.bfloat16: "bf16", dtypes.fp8e4m3: "e4m3", dtypes.fp8e5m2: "e5m2" }
|
||||
dt_map_out = { dtypes.float: "f32", dtypes.half: "f16" }
|
||||
for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes, _ in wmma_args(uops):
|
||||
upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
|
||||
wmma_dtypes = [self.render_dtype(dtype.vec(size)) for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)]
|
||||
n_operands = [size*dtype.itemsize//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] # 4 => CUDA reg size in bytes
|
||||
operands = [f"%{i}" for i in range(sum(n_operands))]
|
||||
|
||||
# mma operands => {c}, {a}, {b}, {c}
|
||||
prefix.append(f"""__device__ {wmma_dtypes[2]} __{name}({wmma_dtypes[0]} a, {wmma_dtypes[1]} b, {wmma_dtypes[2]} c){{
|
||||
int *a_pk = (int *)(&a), *b_pk = (int *)(&b), *c_pk = (int *)(&c);
|
||||
asm("mma.sync.aligned.m{M}n{N}k{K}.row.col.{dt_map_out[dtype_out]}.{dt_map_in[dtype_in]}.{dt_map_in[dtype_in]}.{dt_map_out[dtype_out]}"
|
||||
"{{{", ".join(operands[:n_operands[2]])}}}, {{{", ".join(operands[n_operands[2]:n_operands[2]+n_operands[0]])}}},"
|
||||
"{{{", ".join(operands[-n_operands[1]:])}}}, {{{", ".join(operands[:n_operands[2]])}}};"
|
||||
: {", ".join([f'"+r"(c_pk[{i}])' for i in range(n_operands[2])])}
|
||||
: {", ".join([f'"r"(a_pk[{i}])' for i in range(n_operands[0])])}, {", ".join([f'"r"(b_pk[{i}])' for i in range(n_operands[1])])});
|
||||
return c;\n}}""")
|
||||
|
||||
return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix)
|
||||
|
||||
def supported_dtypes(self):
|
||||
ver = int(self.target.arch[3:])
|
||||
return {d for d in super().supported_dtypes() if (d != dtypes.half or ver >= 53) and (d != dtypes.bfloat16 or ver >= 80)
|
||||
and (d not in dtypes.fp8_ocp or ver >= 89) and d not in dtypes.fp8_fnuz}
|
||||
|
||||
class NVCCRenderer(CUDARenderer):
|
||||
def __init__(self, target:Target): super().__init__(target, use_nvcc=True)
|
||||
|
||||
def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype.scalar())
|
||||
def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})"
|
||||
|
||||
class HIPRenderer(CStyleLanguage):
|
||||
shared_max = 65536
|
||||
# NOTE: this is only really needed on gfx12, even though gfx11 reports the same limitation
|
||||
global_max = (2147483647, 65535, 65535)
|
||||
global_prod_max = (0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF)
|
||||
|
||||
@staticmethod
|
||||
def is_cdna(arch): return arch.split(":")[0] in {"gfx942", "gfx950"}
|
||||
@staticmethod
|
||||
def is_cdna4(arch): return arch.split(":")[0] == "gfx950"
|
||||
def __init__(self, target:Target, use_hipcc=False): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler
|
||||
self.compiler, self.tensor_cores = (HIPCCCompiler if use_hipcc else HIPCompiler)(target.arch), tc.get_amd(target.arch)
|
||||
if not self.is_cdna4(target.arch): self.extra_matcher += pm_manual_bf16_cast + extra_pm
|
||||
if self.is_cdna(target.arch):
|
||||
self.string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
|
||||
f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[1][2] == 128 else None),
|
||||
(UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
|
||||
(UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) else None),
|
||||
(UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"),
|
||||
(UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"),
|
||||
(UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.arg}f, {fp8_index(x.dtype)})"),
|
||||
(UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",),
|
||||
lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"),
|
||||
(UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",),
|
||||
lambda ctx,x,y: f"__builtin_amdgcn_cvt_f32_{('fp8', 'bf8')[fp8_index(y.dtype)]}((unsigned int){ctx[x.src[0]]}, 0)"),
|
||||
]) + base_rewrite
|
||||
|
||||
# https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size
|
||||
# NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters
|
||||
kernel_typedef = 'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {launch_bounds})))'
|
||||
code_for_workitem = {"g": lambda x: f"__ockl_get_group_id({x})", "l": lambda x: f"__ockl_get_local_id({x})",
|
||||
"i": lambda x: f"(__ockl_get_group_id({x})*__ockl_get_local_size({x})+__ockl_get_local_id({x}))"}
|
||||
code_for_op = {**CStyleLanguage.code_for_op, Ops.TRUNC: _ocml("trunc"), Ops.SIN: _ocml("sin"),
|
||||
Ops.LOG2: _ocml("log2"), Ops.EXP2: _ocml("exp2"), Ops.SQRT: _ocml("sqrt")}
|
||||
smem_prefix = "__attribute__((shared, aligned(16)))"
|
||||
smem_prefix_for_cast: bool = False
|
||||
barrier = '__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");' + '__builtin_amdgcn_s_barrier();' + \
|
||||
'__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");'
|
||||
float4 = "make_float4"
|
||||
type_map = {dtypes.bfloat16: "hip_bfloat16", dtypes.fp8e4m3: "hip_fp8", dtypes.fp8e5m2: "hip_bf8"}
|
||||
extra_matcher = create_non_native_float_pats((dtypes.bfloat16, *dtypes.fp8s)) + PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(4)),
|
||||
lambda x: UOp(Ops.WMMA, x.dtype, (x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
|
||||
x.src[2]), (*x.arg,)) if x.src[0].dtype in (dtypes.fp8e4m3.vec(8), dtypes.fp8e5m2.vec(8)) else None),
|
||||
# bfloat16 constant casting
|
||||
(UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))),
|
||||
])
|
||||
|
||||
def asm(self, prg:UOp, lin:UOp) -> bytes:
|
||||
from tinygrad.renderer.amd.elf import assemble_linear
|
||||
return assemble_linear(prg, lin, self.target.arch)
|
||||
|
||||
def render_vector_prefix(self, dtype:DType) -> str:
|
||||
vec, scal = self.render_dtype(dtype), self.render_dtype(dtype.scalar())
|
||||
return f"typedef {scal} {vec} __attribute__((ext_vector_type({dtype.count})));\nstatic inline __attribute__((device)) "+ \
|
||||
f"{vec} make_{vec}({', '.join([f'{scal} {x}' for x in _nms[:dtype.count]])}) {{ return {{ {', '.join(_nms[:dtype.count])} }}; }}"
|
||||
|
||||
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
|
||||
prefix, ockl = [], []
|
||||
type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
|
||||
used_dtypes = uops_to_dtypes(uops)
|
||||
if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops):
|
||||
prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
|
||||
if any(u.op is Ops.SPECIAL for u in uops):
|
||||
prefix.append("typedef long unsigned int size_t;")
|
||||
ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
|
||||
ocml_ops = {Ops.EXP2: ("exp2", "pure"), Ops.LOG2: ("log2", "pure"), Ops.SQRT: ("sqrt", "const"), Ops.SIN: ("sin", ""), Ops.TRUNC: ("trunc", "")}
|
||||
ocml = [(f"__ocml_{ocml_ops[op][0]}_f{dt.bitsize}", dt.name, dt.name, ocml_ops[op][1])
|
||||
for op, dt in dedup((u.op, u.dtype.scalar()) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
|
||||
if any(dt.scalar() == dtypes.bfloat16 for dt in used_dtypes):
|
||||
prefix.append(f"typedef {'__bf16' if self.is_cdna4(self.target.arch) else 'unsigned short'} hip_bfloat16;")
|
||||
if any(dt.scalar() == dtypes.half for dt in used_dtypes): prefix.append("#define half _Float16")
|
||||
if any(dt.scalar() in dtypes.fp8s for dt in used_dtypes):
|
||||
prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"]
|
||||
if any((u.op is Ops.CAST and u.dtype in dtypes.fp8s and u.src[0].dtype == dtypes.float) or
|
||||
(u.op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops):
|
||||
prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {
|
||||
v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v;
|
||||
return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""")
|
||||
prefix += [f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml]
|
||||
prefix += [self.render_vector_prefix(dt) for dt in used_dtypes if dt.count > 1]
|
||||
|
||||
for name, (N, M, K), dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
|
||||
if self.is_cdna(self.target.arch):
|
||||
if (N, M, K) == (16, 16, 16): type_map[dtypes.bfloat16] = 'bf16_1k'
|
||||
elif (N, M, K) == (16, 16, 32): type_map = {**type_map, dtypes.bfloat16: "_bf16", dtypes.half: "_f16"}
|
||||
elif (N, M, K) == (16, 16, 128): type_map = {**type_map, dtypes.fp8e4m3: "_f8f6f4", dtypes.fp8e5m2: "_f8f6f4"}
|
||||
prefix.append(f"#define __{name} __builtin_amdgcn_mfma_{'scale_' if K == 128 else ''}f32_{N}x{M}x{K}{type_map[dtype_in]}")
|
||||
# #define __WMMA_16_16_16_half_half __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12
|
||||
elif self.tensor_cores == tc.amd_rdna4:
|
||||
prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12")
|
||||
elif dtype_out == dtypes.float:
|
||||
prefix.append(f"#define __{name} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if dtype_in == dtypes.half else 'bf16'}_w32")
|
||||
else: prefix.append(f"static inline __attribute__((device)) half8 __{name}"+"""(half16 a, half16 b, half8 c) {
|
||||
half16 c_frag = {}; half8 d; for (int n = 0; n < 8; n++) { c_frag[n*2] = c[n]; }
|
||||
c_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(a, b, c_frag, false);
|
||||
for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""")
|
||||
return super().render_kernel(function_name, kernel, bufs, uops, prefix)
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes()
|
||||
if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}
|
||||
|
||||
class HIPCCRenderer(HIPRenderer):
|
||||
def __init__(self, target:Target): super().__init__(target, use_hipcc=True)
|
||||
|
||||
class QCOMCLRenderer(OpenCLRenderer):
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_qcom import QCOMCompiler
|
||||
self.compiler = QCOMCompiler(target.arch)
|
||||
|
||||
# QCOM compiler is flaky with half
|
||||
def supported_dtypes(self):
|
||||
return {d for d in Renderer.supported_dtypes(self)
|
||||
if (d != dtypes.float16 or (bool(IMAGE) and bool(FLOAT16))) and d not in dtypes.fp8s+(dtypes.bfloat16,dtypes.double)}
|
||||
42
tinygrad_repo/tinygrad/renderer/isa/__init__.py
Normal file
42
tinygrad_repo/tinygrad/renderer/isa/__init__.py
Normal file
@@ -0,0 +1,42 @@
|
||||
from __future__ import annotations
|
||||
import itertools
|
||||
from dataclasses import dataclass, field
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.uop.ops import PatternMatcher, UOp, Ops, consumer_map_from_toposort
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Register:
|
||||
name: str
|
||||
index: int
|
||||
_cons: tuple[Register, ...] = field(default_factory=tuple)
|
||||
@property
|
||||
def cons(self): return self._cons or (self,)
|
||||
def __repr__(self): return self.name
|
||||
|
||||
class IselContext:
|
||||
def __init__(self, sink:UOp):
|
||||
self.uses = consumer_map_from_toposort(sink.toposort())
|
||||
self.reg_n = itertools.count()
|
||||
arg_order = {Ops.PARAM: 0, Ops.DEFINE_VAR: 1, Ops.SPECIAL: 2}
|
||||
self.func_args = sorted([u for u in self.uses if u.op in arg_order], key=lambda k: (arg_order[k.op], k.arg))
|
||||
|
||||
def vreg(self, cons:tuple[Register, ...]|Register):
|
||||
return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
|
||||
|
||||
@dataclass
|
||||
class PreRegAllocContext:
|
||||
lock: UOp|None = None
|
||||
clobbered: set[UOp] = field(default_factory=set)
|
||||
|
||||
class ISARenderer(Renderer):
|
||||
pre_isel_matcher: PatternMatcher
|
||||
isel_matcher: PatternMatcher
|
||||
pre_regalloc_matcher: PatternMatcher|None = None
|
||||
post_regalloc_matcher: PatternMatcher
|
||||
|
||||
def is_two_address(self, x:UOp) -> bool: return False
|
||||
def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific")
|
||||
def copy(self, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
|
||||
def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific")
|
||||
def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
|
||||
def asm_str(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific")
|
||||
905
tinygrad_repo/tinygrad/renderer/isa/x86.py
Normal file
905
tinygrad_repo/tinygrad/renderer/isa/x86.py
Normal file
@@ -0,0 +1,905 @@
|
||||
# flake8: noqa: E702
|
||||
# allow semicolons to put multiple ops on one line
|
||||
import sys, struct, functools
|
||||
from typing import cast
|
||||
from tinygrad.dtype import dtypes, PtrDType, DType, truncate, AddrSpace
|
||||
from tinygrad.uop import FastEnum, auto, Ops, GroupOp
|
||||
from tinygrad.uop.ops import UOp, UPat, PatternMatcher
|
||||
from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext
|
||||
from tinygrad.helpers import getenv, CPU_COUNT, unwrap, Target
|
||||
|
||||
# ***** X86 Ops *****
|
||||
|
||||
class X86Ops(FastEnum):
|
||||
# NOTE: X86Ops with i suffix are variants that take an immediate, m suffix are variants that can write to memory instead of read from
|
||||
# these aren't real instructions
|
||||
FRAME_INDEX = auto(); LABEL = auto()
|
||||
# index
|
||||
LEA = auto()
|
||||
# register / memory / immediate moves
|
||||
MOV = auto(); MOVm = auto(); MOVi = auto(); MOVABS = auto()
|
||||
VMOVSS = auto(); VMOVSD = auto(); VMOVUPS = auto()
|
||||
VMOVSSm = auto(); VMOVSDm = auto(); VMOVUPSm = auto()
|
||||
# casts
|
||||
MOVZX = auto(); MOVSX = auto(); MOVSXD = auto()
|
||||
VPMOVZXBW = auto(); VPMOVZXBD = auto(); VPMOVZXBQ = auto()
|
||||
VPMOVZXWD = auto(); VPMOVZXWQ = auto(); VPMOVZXDQ = auto()
|
||||
VPMOVSXBW = auto(); VPMOVSXBD = auto(); VPMOVSXBQ = auto()
|
||||
VPMOVSXWD = auto(); VPMOVSXWQ = auto(); VPMOVSXDQ = auto()
|
||||
VCVTDQ2PS = auto(); VCVTDQ2PD = auto(); VCVTTPS2DQ = auto(); VCVTTPD2DQ = auto()
|
||||
VCVTPH2PS = auto(); VCVTPS2PH = auto(); VCVTPS2PD = auto(); VCVTPD2PS = auto()
|
||||
VCVTSS2SD = auto(); VCVTSD2SS = auto(); VCVTSI2SS = auto(); VCVTSI2SD = auto()
|
||||
VCVTTSS2SI = auto(); VCVTTSD2SI = auto()
|
||||
# bitcasts
|
||||
VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto()
|
||||
# comparisons
|
||||
VUCOMISS = auto(); VUCOMISD = auto()
|
||||
VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto()
|
||||
VPCMPGTB = auto(); VPCMPGTW = auto(); VPCMPGTD = auto(); VPCMPGTQ = auto()
|
||||
VPCMPEQB = auto(); VPCMPEQW = auto(); VPCMPEQD = auto(); VPCMPEQQ = auto()
|
||||
SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto()
|
||||
# where
|
||||
CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto()
|
||||
VPBLENDVB = auto(); VBLENDVPS = auto(); VBLENDVPD = auto()
|
||||
# jumps
|
||||
JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto()
|
||||
# vectorize / gep
|
||||
VSHUFPS = auto(); VSHUFPD = auto(); VINSERTPS = auto(); VPSRLDQ = auto()
|
||||
VPEXTRB = auto(); VPEXTRW = auto(); VPEXTRD = auto(); VPEXTRQ = auto()
|
||||
VPINSRB = auto(); VPINSRW = auto(); VPINSRD = auto(); VPINSRQ = auto()
|
||||
VPBROADCASTB = auto(); VPBROADCASTW = auto(); VPBROADCASTD = auto(); VPBROADCASTQ = auto()
|
||||
VBROADCASTSS = auto()
|
||||
# int binary
|
||||
IDIV = auto(); DIV = auto()
|
||||
ADD = auto(); ADDi = auto(); SUB = auto(); SUBi = auto(); IMUL = auto(); IMULi = auto()
|
||||
AND = auto(); ANDi = auto(); XOR = auto(); XORi = auto(); OR = auto(); ORi = auto()
|
||||
SHL = auto(); SHLi = auto(); SHR = auto(); SHRi = auto(); SAR = auto(); SARi = auto(); CMP = auto(); CMPi = auto()
|
||||
# float unary (sometimes not unary)
|
||||
VROUNDSS = auto(); VROUNDSD = auto(); VROUNDPS = auto(); VROUNDPD = auto()
|
||||
VSQRTSS = auto(); VSQRTSD = auto(); VSQRTPS = auto(); VSQRTPD = auto()
|
||||
# float scalar / vector binary
|
||||
VADDSS = auto(); VADDSD = auto(); VADDPS = auto(); VADDPD = auto()
|
||||
VSUBSS = auto(); VSUBSD = auto(); VSUBPS = auto(); VSUBPD = auto()
|
||||
VMULSS = auto(); VMULSD = auto(); VMULPS = auto(); VMULPD = auto()
|
||||
VDIVSS = auto(); VDIVSD = auto(); VDIVPS = auto(); VDIVPD = auto()
|
||||
VMAXSS = auto(); VMAXSD = auto(); VMAXPS = auto(); VMAXPD = auto()
|
||||
VMINSS = auto(); VMINSD = auto(); VMINPS = auto(); VMINPD = auto()
|
||||
# int vector binary
|
||||
VPADDB = auto(); VPADDW = auto(); VPADDD = auto(); VPADDQ = auto()
|
||||
VPSUBB = auto(); VPSUBW = auto(); VPSUBD = auto(); VPSUBQ = auto()
|
||||
VPMULLW = auto(); VPMULLD = auto()
|
||||
# packed bitwise
|
||||
VPAND = auto(); VPOR = auto(); VPXOR = auto()
|
||||
# packed variable shifts
|
||||
VPSLLVD = auto(); VPSLLVQ = auto(); VPSRLVD = auto(); VPSRLVQ = auto(); VPSRAVD = auto()
|
||||
# fused multiply add
|
||||
VFMADD213SS = auto(); VFMADD213SD = auto(); VFMADD213PS = auto(); VFMADD213PD = auto()
|
||||
# return
|
||||
RET = auto()
|
||||
|
||||
class X86GroupOp:
|
||||
# X86Ops whose first src is also the destination
|
||||
TwoAddress = {X86Ops.ADD, X86Ops.ADDi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi, X86Ops.OR, X86Ops.ORi, X86Ops.IMUL,
|
||||
X86Ops.SUB, X86Ops.SUBi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi,
|
||||
X86Ops.IDIV, X86Ops.DIV, X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD,
|
||||
X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB}
|
||||
|
||||
# X86Ops whose first src can read from memory
|
||||
ReadMem1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ,
|
||||
X86Ops.VPMOVZXBW, X86Ops.VPMOVZXBD, X86Ops.VPMOVZXBQ, X86Ops.VPMOVZXWD, X86Ops.VPMOVZXWQ, X86Ops.VPMOVZXDQ,
|
||||
X86Ops.VPMOVSXBW, X86Ops.VPMOVSXBD, X86Ops.VPMOVSXBQ, X86Ops.VPMOVSXWD, X86Ops.VPMOVSXWQ, X86Ops.VPMOVSXDQ,
|
||||
X86Ops.VCVTDQ2PS, X86Ops.VCVTDQ2PD, X86Ops.VCVTTPS2DQ, X86Ops.VCVTTPD2DQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI,
|
||||
X86Ops.VCVTPH2PS, X86Ops.VCVTPS2PD, X86Ops.VCVTPD2PS, X86Ops.VROUNDPS, X86Ops.VROUNDPD, X86Ops.VSQRTPS, X86Ops.VSQRTPD,
|
||||
X86Ops.VPBROADCASTB, X86Ops.VPBROADCASTW, X86Ops.VPBROADCASTD, X86Ops.VPBROADCASTQ, X86Ops.VBROADCASTSS,
|
||||
X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA}
|
||||
|
||||
# X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first
|
||||
ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.SHL, X86Ops.SHR, X86Ops.SAR, X86Ops.IMUL, X86Ops.CMP,
|
||||
X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
|
||||
X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
|
||||
X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
|
||||
X86Ops.VPCMPEQB, X86Ops.VPCMPEQW, X86Ops.VPCMPEQD, X86Ops.VPCMPEQQ, X86Ops.VPBLENDVB, X86Ops.VBLENDVPS, X86Ops.VBLENDVPD,
|
||||
X86Ops.VPCMPGTB, X86Ops.VPCMPGTW, X86Ops.VPCMPGTD, X86Ops.VPCMPGTQ, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
|
||||
X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VSHUFPS, X86Ops.VINSERTPS,
|
||||
X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD,
|
||||
X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
|
||||
X86Ops.VMAXSS, X86Ops.VMAXSD, X86Ops.VMAXPS, X86Ops.VMAXPD, X86Ops.VMINSS, X86Ops.VMINSD, X86Ops.VMINPS, X86Ops.VMINPD,
|
||||
X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.VUCOMISS, X86Ops.VUCOMISD, X86Ops.IDIV, X86Ops.DIV,
|
||||
X86Ops.VSHUFPD}
|
||||
|
||||
# X86Ops whose third src can read from memory NOTE: these are TwoAddress so the third src is actually the second
|
||||
ReadMem3rd = {X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD}
|
||||
|
||||
# X86Ops that can write to memory
|
||||
WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
|
||||
X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHLi, X86Ops.SHRi, X86Ops.SARi, X86Ops.SETNE,
|
||||
X86Ops.SETE, X86Ops.SETL, X86Ops.SETB, X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
|
||||
|
||||
# X86Ops that read flags
|
||||
ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
|
||||
X86Ops.JE, X86Ops.JNE, X86Ops.JGE}
|
||||
|
||||
# X86Ops that write flags or can modify flags to undefined values
|
||||
WriteFlags = {X86Ops.CMP, X86Ops.CMPi, X86Ops.ADD, X86Ops.ADDi, X86Ops.SUB, X86Ops.SUBi, X86Ops.IMUL, X86Ops.IMULi, X86Ops.IDIV, X86Ops.DIV,
|
||||
X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi,
|
||||
X86Ops.OR, X86Ops.ORi, X86Ops.VUCOMISS, X86Ops.VUCOMISD}
|
||||
|
||||
# X86Ops whose first src is the rm field
|
||||
Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ}
|
||||
|
||||
# X86Ops whose second src is the rm field
|
||||
Rm2nd = ReadMem2nd | (ReadMem3rd & TwoAddress)
|
||||
|
||||
All = set(X86Ops)
|
||||
|
||||
# ***** X86 legalization *****
|
||||
|
||||
extra_matcher = PatternMatcher([
|
||||
# bool CMPNE is XOR, bool CMPEQ is XOR+XOR, bool CMPLT is XOR+AND
|
||||
(UPat.var('x', dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y),
|
||||
(UPat.var('x', dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True),
|
||||
(UPat.var('x', dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
|
||||
# can't cast from float16 to ints/float64 directly and vice versa
|
||||
(UPat.var("y", dtypes.float16).cast((dtypes.float64,)+dtypes.ints, name="x"), lambda y,x: y.cast(dtypes.float32).cast(x.dtype)),
|
||||
(UPat.var("y", (dtypes.float64,)+dtypes.ints).cast(dtypes.float16, name="x"), lambda y,x: y.cast(dtypes.float32).cast(x.dtype)),
|
||||
# can't cast from float to int8/16 directly and vice versa
|
||||
(UPat.var("y", dtypes.floats).cast(dtypes.int8s+dtypes.int16s, name="x"), lambda y,x: y.cast(dtypes.int32).cast(x.dtype)),
|
||||
(UPat.var("y", (dtypes.bool,)+dtypes.int8s+dtypes.int16s).cast(dtypes.floats, name="x"), lambda y,x: y.cast(dtypes.int32).cast(x.dtype)),
|
||||
# int/float casts only for signed int
|
||||
(UPat.var("y", dtypes.uint32).cast(dtypes.floats, name="x"), lambda y,x: y.cast(dtypes.int64).cast(x.dtype)),
|
||||
# casting uint64 to float requires special handling
|
||||
(UPat.var("y", dtypes.uint64).cast(dtypes.floats, name="x"), lambda y,x:
|
||||
(y >> 1).cast(dtypes.int64).cast(x.dtype) * 2 + (y & 1).cast(dtypes.int64).cast(x.dtype)),
|
||||
# no int8 mul or cmove, cast to int16
|
||||
(UPat.var("a", dtypes.int8s) * UPat.var("b"), lambda a,b: (a.cast(dtypes.int16) * b.cast(dtypes.int16)).cast(a.dtype)),
|
||||
(UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")),
|
||||
lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.dtype.count == 1 else None),
|
||||
# float16 alus are done in float32
|
||||
(UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op, dtypes.float.vec(x.dtype.count),
|
||||
tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
|
||||
(UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
|
||||
lambda x,a,b: UOp(x.op, x.dtype, (a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
|
||||
# no cmpne for packed ints, y != x => !(y==x)
|
||||
(UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
|
||||
lambda y,x,cmp: UOp(Ops.CMPEQ, cmp.dtype, (y,x))^True if y.dtype.count > 1 else None),
|
||||
# float where expects a mask
|
||||
(UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")),
|
||||
lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None),
|
||||
# rewrite -x -> 0 - x
|
||||
(UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, x.dtype, (x.const_like(0),) + x.src)),
|
||||
# TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction
|
||||
(UPat(Ops.CMOD, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x - y * x.alu(Ops.CDIV, y)),
|
||||
])
|
||||
|
||||
# ***** X86 pre instruction selection *****
|
||||
|
||||
def gated_load(ctx, base:UOp, idx:UOp, cast:UOp, alt:UOp, gate:UOp, x:UOp):
|
||||
local = UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(x.dtype.count, AddrSpace.LOCAL), arg=next(ctx))
|
||||
local_idx = local.index(UOp.const(dtypes.int32, 0), ptr=True)
|
||||
ptr = gate.where(base.index(idx, ptr=True), local_idx).after((local_idx if x.dtype.count == 1 else local).store(alt))
|
||||
return ptr.cast(cast.dtype).load(dtype=x.dtype)
|
||||
|
||||
def gated_store(base:UOp, idx:UOp, cast:UOp, gate:UOp, val:UOp):
|
||||
local = UOp(Ops.DEFINE_LOCAL, base.dtype.base.ptr(val.dtype.count, AddrSpace.LOCAL), arg=-1)
|
||||
ptr = gate.where(base.index(idx, ptr=True), local.index(UOp.const(dtypes.int32, 0), ptr=True))
|
||||
return ptr.cast(cast.dtype).store(val)
|
||||
|
||||
# these must be done in a separate matcher because they violate the spec
|
||||
pre_isel_matcher = PatternMatcher([
|
||||
# zero extending scalar 32bit int is a noop
|
||||
(UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.NOOP) if y.dtype.count == 1 else None),
|
||||
# cast between signed and unsigned int is a noop
|
||||
(UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"),
|
||||
lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize == y.dtype.itemsize else None),
|
||||
# cast to < scalar int is a noop
|
||||
(UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"),
|
||||
lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize < y.dtype.itemsize and y.dtype.count == 1 else None),
|
||||
# bitcasts between scalar floats and ints are real, rest are noops
|
||||
(UPat.var("y").bitcast().named("x"), lambda y,x: None if y.dtype in dtypes.floats and x.dtype in dtypes.ints or \
|
||||
y.dtype in dtypes.ints and x.dtype in dtypes.floats else x.replace(op=Ops.NOOP)),
|
||||
# noop of a noop is removed
|
||||
(UPat(Ops.NOOP, src=(UPat(Ops.NOOP),), name="x"), lambda x: x.replace(src=x.src[0].src)),
|
||||
# moving elements of a single register to another without shuffling is a noop
|
||||
(UPat(Ops.STACK, src=(UPat.var("y"),), allow_any_len=True, name="x"),
|
||||
lambda y,x: UOp(Ops.NOOP, x.dtype, y.src) if all(s.op is Ops.GEP and s.src == y.src and s.arg[0] == i for i,s in enumerate(x.src)) else None),
|
||||
# gated load/store become a conditional move on the index, the load/store are unconditional
|
||||
(UPat.var("base").index(UPat.var("idx")).or_casted(name="cast").load(UPat.var("alt"), UPat.var("gate"), name="x"), gated_load),
|
||||
(UPat.var("base").index(UPat.var("idx")).or_casted(name="cast").store(UPat.var("val"), UPat.var("gate")), gated_store),
|
||||
# TODO: remove this once we allow all flag producing ops in cmove
|
||||
# if gate in scalar int cmove is not a comparison need to add one to set the flag
|
||||
(UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
|
||||
lambda m,a,b: m.ne(0).where(a,b) if m.op not in GroupOp.Comparison and a.dtype.count == 1 else None),
|
||||
])
|
||||
|
||||
# ***** X86 registers *****
|
||||
|
||||
RAX = Register("rax", 0)
|
||||
RCX = Register("rcx", 1)
|
||||
RDX = Register("rdx", 2)
|
||||
RBX = Register("rbx", 3)
|
||||
RSP = Register("rsp", 4)
|
||||
RBP = Register("rbp", 5)
|
||||
RSI = Register("rsi", 6)
|
||||
RDI = Register("rdi", 7)
|
||||
GPR = (RAX, RCX, RDX, RBX, RSP, RBP, RSI, RDI) + tuple(Register(f"r{i}", i) for i in range(8, 16))
|
||||
XMM = tuple(Register(f"xmm{i}", i) for i in range(16))
|
||||
# gprs you can write to
|
||||
WGPR = tuple(r for r in GPR if r != RSP)
|
||||
|
||||
CALLEE_SAVED = (RBX, RBP, GPR[12], GPR[13], GPR[14], GPR[15]) + ((RSI, RDI) + XMM[6:16] if sys.platform == "win32" else ())
|
||||
|
||||
reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, "rdx": {4:"edx", 2:"dx", 1:"dl"}, "rbx": {4:"ebx", 2:"bx", 1:"bl"},
|
||||
"rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
|
||||
**{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}, **{f"xmm{i}": {64:f"zmm{i}", 32:f"ymm{i}"} for i in range(16)}}
|
||||
|
||||
# ***** X86 instruction selection *****
|
||||
# if s is used multiple times we don't fold
|
||||
def is_foldable(ctx:IselContext, x:UOp, s:UOp) -> bool: return len(ctx.uses[s]) == x.src.count(s) == 1
|
||||
def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.GEP else s
|
||||
def lane(x:UOp, i:int) -> int: return s.arg[0] if (s:=x.src[i]).op is Ops.GEP else 0
|
||||
def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
|
||||
def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.DEFINE_REG, dt, tag=None if reg is None else (reg,))
|
||||
def imm(dt:DType, v:int) -> UOp: return UOp.const(dt, truncate[dt](v)).rtag()
|
||||
def to_imm(c:UOp) -> UOp|None:
|
||||
if c.op is not Ops.CONST: return None
|
||||
if c.dtype is dtypes.int64: return imm(dtypes.int32, c.arg) if not c.overflows(dtypes.int32) else None
|
||||
if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.arg) if not c.overflows(dtypes.uint32) else None
|
||||
if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.arg)
|
||||
return None
|
||||
def cmp(x:UOp) -> UOp:
|
||||
if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void)
|
||||
if x.src[0].dtype is dtypes.float64: return x.ins(X86Ops.VUCOMISD, dtype=dtypes.void)
|
||||
return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i))
|
||||
def vcmp(x:UOp) -> UOp:
|
||||
v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
|
||||
if x.dtype.scalar() is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.dtype.count == 1 else X86Ops.VCMPPS, src=x.src + (v,))
|
||||
return x.ins(X86Ops.VCMPSD if x.dtype.count == 1 else X86Ops.VCMPPD, src=x.src + (v,))
|
||||
|
||||
# vshufps xmm2, xmm0, xmm1, imm
|
||||
# for 128 bit xmm2 selects its lower 2 32 bits from xmm0 and its upper 2 32 bits from xmm1 according to imm
|
||||
# for 256 bit ymm2 repeats the shuffle for its upper 128 bits selecting from the upper 128 bits of ymm0 and ymm1
|
||||
def vshufps(x:UOp) -> UOp|None:
|
||||
a, b = base(x, 0), base(x, 2)
|
||||
if not (a is base(x, 1) and b is base(x, 3)) or any(lane(x, i) > 3 for i in range(4)): return None
|
||||
if len(x.src) == 8:
|
||||
if not (a is base(x, 4) is base(x, 5) and b is base(x, 6) is base(x, 7)) or any(lane(x, i+4) != lane(x, i)+4 for i in range(4)): return None
|
||||
return x.ins(X86Ops.VSHUFPS, src=(a, b, imm(dtypes.uint8, sum(lane(x, i) << 2*i for i in range(4)))))
|
||||
|
||||
# vshufpd xmm2, xmm0, xmm1, imm
|
||||
# for 128 bit xmm2 selects its lower 64 bits from xmm0 and its upper 64 bits from xmm1 according to imm
|
||||
# for 256 bit ymm2 also selects its upper 128 bits from the upper 128 bits of ymm0 and ymm1 following the same constraint
|
||||
def vshufpd(x:UOp) -> UOp|None:
|
||||
a, b = base(x, 0), base(x, 1)
|
||||
if lane(x, 0) > 1 or lane(x, 1) > 1: return None
|
||||
if len(x.src) == 4 and not (a is base(x, 2) and b is base(x, 3) and lane(x, 2) > 1 and lane(x, 3) > 1): return None
|
||||
return x.ins(X86Ops.VSHUFPD, src=(a, b, imm(dtypes.uint8, sum(lane(x, i) << i for i in range(len(x.src))))))
|
||||
|
||||
# vinsertps xmm2, xmm0, xmm1, imm
|
||||
# inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2
|
||||
# this is the fallback slow case for when you can't match more a powerful shuffle
|
||||
def vinsertps(x:UOp) -> UOp:
|
||||
def _insert(ret:UOp, i:int) -> UOp:
|
||||
s, v = base(x, i), lane(x, i)
|
||||
# moving the 0th element into the 0th position does nothing
|
||||
return s if i == v == 0 else x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4)))
|
||||
return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype))
|
||||
|
||||
# vpinsq xmm2, xmm0, rax, imm
|
||||
# inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
|
||||
def vpins(x:UOp) -> UOp:
|
||||
op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.scalar().itemsize]
|
||||
return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype))
|
||||
|
||||
# vpbroadcastd xmm1, xmm0
|
||||
# inserts scalar int in xmm0 into all lanes of xmm1
|
||||
def vpbroadcast(ctx:IselContext, x:UOp, y:UOp) -> UOp:
|
||||
n = x.ins({1: X86Ops.VPBROADCASTB, 2: X86Ops.VPBROADCASTW, 4: X86Ops.VPBROADCASTD, 8: X86Ops.VPBROADCASTQ}[y.dtype.itemsize], src=(y,))
|
||||
if y.op is Ops.LOAD and len(y.src) == 1 and is_foldable(ctx, n, y): return n
|
||||
# if there isn't a load we can fold we need to move y from gpr to xmm
|
||||
# this is hacky but required because int.vec(1) isn't supported
|
||||
y = y if y.dtype.itemsize > 1 else y.cast(dtypes.int16)
|
||||
return n.replace(src=(y.bitcast({2:dtypes.float16, 4:dtypes.float32, 8:dtypes.float64}[y.dtype.itemsize]),))
|
||||
|
||||
# we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
|
||||
def idiv(ctx:IselContext, x:UOp) -> UOp:
|
||||
op = X86Ops.DIV if x.dtype in dtypes.uints else X86Ops.IDIV
|
||||
# for >8bit need to zero/sign extend rax to rdx
|
||||
if x.dtype in dtypes.int8s: ext = []
|
||||
elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(imm(min(dtypes.uint32, x.dtype), 0),), tag=(RDX,))]
|
||||
else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], imm(dtypes.uint8, x.dtype.itemsize * 8 - 1)), tag=(RDX,))]
|
||||
# for 8bit need to zero/sign extend al to ah
|
||||
if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=X86Ops.MOVZX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
|
||||
elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=X86Ops.MOVSX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
|
||||
else: dividend = x.ins(X86Ops.MOV, src=(x.src[0],), tag=(RAX,))
|
||||
# divisor can't be in rax or rdx
|
||||
divisor = x.ins(X86Ops.MOV, src=(x.src[1],), tag=tuple(r for r in WGPR if r not in (RAX, RDX)))
|
||||
# for >8bit both rax and rdx are written to
|
||||
defs = (ctx.vreg(RAX),) if x.dtype in dtypes.int8s else (ctx.vreg(RAX), ctx.vreg(RDX))
|
||||
idiv = x.ins(op, src=(dividend, divisor) + tuple(ext), tag=defs)
|
||||
# this move "cleanses" the register constraints (rax/rdx) of idiv as that only applies on definition and not on the uses of idiv
|
||||
return x.ins(X86Ops.MOV, src=(idiv,))
|
||||
|
||||
def fold_address(x:UOp) -> tuple[UOp, UOp, UOp]:
|
||||
def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v)
|
||||
def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v
|
||||
if x.op is not Ops.INDEX: return (x, UOp(Ops.NOOP), _disp(0))
|
||||
base, idx = x.src
|
||||
disp_scale = base.dtype.itemsize if isinstance(base.dtype, PtrDType) else 1
|
||||
if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].arg * disp_scale))
|
||||
if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.arg * disp_scale))
|
||||
return (base, _cast(idx), _disp(0))
|
||||
|
||||
def abi(ctx:IselContext, x:UOp) -> UOp|None:
|
||||
if isinstance(x.tag, tuple): return None
|
||||
i = ctx.func_args.index(x)
|
||||
def _stack_arg(disp:int): return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=X86Ops.FRAME_INDEX, dtype=dtypes.int32, tag=disp))
|
||||
if sys.platform == "win32": src = (x.replace(tag=((RCX, RDX, GPR[8], GPR[9])[i],)),) if i < 4 else _stack_arg((i-3)*8+32)
|
||||
else: src = (x.replace(tag=((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i],)),) if i < 6 else _stack_arg((i-5)*8)
|
||||
# this move "cleanses" the abi register constraint
|
||||
return x.ins(X86Ops.MOV, src=src)
|
||||
|
||||
def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
|
||||
# real registers
|
||||
if x.op is Ops.DEFINE_REG and x.tag is not None: return None
|
||||
# this is an immediate
|
||||
if x.arg is X86Ops.FRAME_INDEX: return None
|
||||
# no register definition
|
||||
if x.dtype is dtypes.void: return None
|
||||
# already allocated vregs
|
||||
if isinstance(x.tag, tuple) and x.tag[0]._cons: return None
|
||||
# allocate vreg definitions
|
||||
defs = []
|
||||
if isinstance(x.tag, tuple): defs = [ctx.vreg(x.tag)]
|
||||
elif x.dtype in dtypes.ints+(dtypes.bool,) or isinstance(x.dtype, PtrDType): defs = [ctx.vreg(WGPR)]
|
||||
elif x.dtype in dtypes.floats or x.dtype.count > 1: defs = [ctx.vreg(XMM)]
|
||||
# TODO: add this once the scheduler can track register pressure
|
||||
# if x.arg in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS))
|
||||
return x.replace(tag=tuple(defs))
|
||||
|
||||
dts = dtypes.ints + (dtypes.bool, dtypes.float16, dtypes.float32, dtypes.float64)
|
||||
dt_16bit = tuple(dt.vec(l) for dt in dts for l in [2,1] if l*dt.itemsize == 2 and dt not in dtypes.int16s)
|
||||
dt_32bit = tuple(dt.vec(l) for dt in dts for l in [4,2,1] if l*dt.itemsize == 4 and dt not in dtypes.int32s)
|
||||
dt_64bit = tuple(dt.vec(l) for dt in dts for l in [8,4,2,1] if l*dt.itemsize == 8 and dt not in dtypes.int64s)
|
||||
dt_128bit = tuple(dt.vec(l) for dt in dts for l in [16,8,4,2,1] if l*dt.itemsize == 16)
|
||||
|
||||
isel_matcher = PatternMatcher([
|
||||
# **** Op -> Op ****
|
||||
# cast to pointer is a noop
|
||||
(UPat.var("y").cast(name="x"), lambda y,x: y if isinstance(x.dtype, PtrDType) or y.dtype == dtypes.void else None),
|
||||
# float gep(0) is a noop as it just moves the 0th element from one xmm register to another
|
||||
# this is done here to not interfere with shuffles
|
||||
(UPat(dtype=dtypes.floats).gep(0, name="x"), lambda x: x.replace(op=Ops.NOOP, arg=None)),
|
||||
# range is lowered to acc, cmp, jmp after regalloc
|
||||
(UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])),
|
||||
(UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
|
||||
# **** Op -> X86Op ****
|
||||
# add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc
|
||||
# so regalloc builds the prologue/epilogue naturally
|
||||
(UPat(Ops.SINK, name="x"), lambda x:
|
||||
x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64.vec(2), r) for r in CALLEE_SAVED)),)) \
|
||||
if not x.src or x.src[0].arg is not X86Ops.RET else None),
|
||||
# function abi constraints
|
||||
(UPat((Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL), name="x"), abi),
|
||||
# these are treated the same for now
|
||||
(UPat(Ops.DEFINE_REG, name="x"), lambda x:
|
||||
x.replace(op=Ops.DEFINE_LOCAL, dtype=x.dtype.base.ptr(x.dtype.size, AddrSpace.LOCAL)) if isinstance(x.arg, int) else None),
|
||||
# constants that can't be immediates, move them to registers
|
||||
(UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.arg),)) if not x.tag else None),
|
||||
(UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),)) if not x.tag else None),
|
||||
(UPat.cvar("x", dtypes.floats), lambda x:
|
||||
UOp.const(dt:=to_int(x.dtype), struct.unpack(dt.fmt, struct.pack(x.dtype.fmt, x.arg))[0]).bitcast(x.dtype) if not x.tag else None),
|
||||
# TODO: these should use a.maximum(b) / a.minimum(b)
|
||||
((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float32), UPat.var("a")), lambda a,b:
|
||||
a.ins(X86Ops.VMAXSS if a.dtype.count == 1 else X86Ops.VMAXPS, src=(a, b))),
|
||||
((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float64), UPat.var("a")), lambda a,b:
|
||||
a.ins(X86Ops.VMAXSD if a.dtype.count == 1 else X86Ops.VMAXPD, src=(a, b))),
|
||||
((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda a,b:
|
||||
a.ins(X86Ops.VMINSS if a.dtype.count == 1 else X86Ops.VMINPS, src=(a, b))),
|
||||
((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda a,b:
|
||||
a.ins(X86Ops.VMINSD if a.dtype.count == 1 else X86Ops.VMINPD, src=(a, b))),
|
||||
# conditional moves that use masks NOTE: these currently assume a mask producing cmp exists
|
||||
(UPat.var("m").where(UPat.var("a", dtypes.ints), UPat.var("b")), lambda m,a,b:
|
||||
a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.dtype.count > 1 else None),
|
||||
(UPat.var("m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b:
|
||||
a.ins(X86Ops.VBLENDVPS, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
|
||||
(UPat.var("m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b:
|
||||
a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
|
||||
# in this case we have a mask producing comparison whose user expects a bool, so we convert to bool
|
||||
(UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x:
|
||||
x.replace(dtype=y.dtype).bitcast(to_int(y.dtype)).bitwise_and(1).f(Ops.NOOP, dtype=dtypes.bool)),
|
||||
# conditional moves that use flags
|
||||
(UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b:
|
||||
a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))),
|
||||
(UPat(Ops.CMPLT, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVB, src=(b, a, cmp(m)))),
|
||||
(UPat(Ops.CMPEQ, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVE, src=(b, a, cmp(m)))),
|
||||
(UPat(Ops.CMPNE, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVNE, src=(b, a, cmp(m)))),
|
||||
# jumps, use flags
|
||||
(UPat(Ops.IF, src=(UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.uints), UPat()), name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JB, src=(cmp(y),))),
|
||||
(UPat(Ops.IF, src=(UPat(Ops.CMPLT, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JL, src=(cmp(y),))),
|
||||
(UPat(Ops.IF, src=(UPat(Ops.CMPEQ, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JE, src=(cmp(y),))),
|
||||
(UPat(Ops.IF, src=(UPat(Ops.CMPNE, name="y"),), name="x"), lambda y,x: x.ins(X86Ops.JNE, src=(cmp(y),))),
|
||||
# comparisons whose user doesn't use the flag, move flag result to register
|
||||
(UPat(Ops.CMPLT, dtypes.bool, (UPat(dtype=dtypes.uints), UPat()), name="x"), lambda x: x.ins(X86Ops.SETB, src=(cmp(x),))),
|
||||
(UPat(Ops.CMPLT, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETL, src=(cmp(x),))),
|
||||
(UPat(Ops.CMPEQ, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))),
|
||||
(UPat(Ops.CMPNE, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))),
|
||||
# comparisons that produce masks (these aren't bool dtype)
|
||||
(UPat(GroupOp.Comparison, src=(UPat(dtype=(dtypes.float32, dtypes.float64)), UPat()), name="x"), vcmp),
|
||||
(UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int8s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQB)),
|
||||
(UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int16s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQW)),
|
||||
(UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int32s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQD)),
|
||||
(UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int64s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQQ)),
|
||||
(UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int8s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTB, src=(b, a))),
|
||||
(UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int16s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTW, src=(b, a))),
|
||||
(UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int32s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTD, src=(b, a))),
|
||||
(UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int64s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTQ, src=(b, a))),
|
||||
# float unary
|
||||
(UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.dtype.count == 1 else x.ins(X86Ops.VSQRTPS)),
|
||||
(UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.dtype.count == 1 else x.ins(X86Ops.VSQRTPD)),
|
||||
(UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x:
|
||||
x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else x.ins(X86Ops.VROUNDPS, src=(y, imm(dtypes.uint8, 3)))),
|
||||
(UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x:
|
||||
x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))),
|
||||
# shufles
|
||||
(UPat.var("y", dtypes.float32).broadcast(name="x"), lambda y,x: x.ins(X86Ops.VBROADCASTSS, src=(y,))),
|
||||
# for float16 we route the srcs through gprs unless we can fold them, this is suboptimal for values in xmms, in that case we want vpunpcklwd
|
||||
(UPat(Ops.STACK, dtypes.float16, name="x"), lambda ctx,x:
|
||||
vpins(x.replace(src=tuple(s if s.op is Ops.LOAD and is_foldable(ctx, x, s) else s.bitcast(dtypes.int16) for s in x.src)))),
|
||||
(UPat(Ops.STACK, (dtypes.float32.vec(4), dtypes.float32.vec(8)), name="x"), vshufps),
|
||||
(UPat(Ops.STACK, (dtypes.float64.vec(2), dtypes.float64.vec(4)), name="x"), vshufpd),
|
||||
(UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
|
||||
(UPat.var("y", dtypes.ints+(dtypes.bool,)).broadcast(name="x"), vpbroadcast),
|
||||
(UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins),
|
||||
# gep
|
||||
(UPat.var("y", dtypes.int8s+(dtypes.bool,)).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, x.arg[0])))),
|
||||
(UPat.var("y", dtypes.int16s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, x.arg[0])))),
|
||||
(UPat.var("y", dtypes.int32s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, x.arg[0])))),
|
||||
(UPat.var("y", dtypes.int64s).gep(name="x"), lambda y,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, x.arg[0])))),
|
||||
(UPat.var("y", dtypes.floats).gep(name="x"), lambda y,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, x.arg[0] * x.dtype.itemsize)))),
|
||||
# fused multiply add
|
||||
((UPat(Ops.MUL, dtypes.float32, name="a") + UPat.var("b")).named("c"), lambda ctx,a,b,c:
|
||||
a.ins(X86Ops.VFMADD213SS if a.dtype.count == 1 else X86Ops.VFMADD213PS, src=(*a.src, b)) if is_foldable(ctx, c, a) else None),
|
||||
((UPat(Ops.MUL, dtypes.float64, name="a") + UPat.var("b")).named("c"), lambda ctx,a,b,c:
|
||||
a.ins(X86Ops.VFMADD213SD if a.dtype.count == 1 else X86Ops.VFMADD213PD, src=(*a.src, b)) if is_foldable(ctx, c, a) else None),
|
||||
# packed bitwise
|
||||
((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.dtype.count > 1 else None),
|
||||
((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.dtype.count > 1 else None),
|
||||
((UPat() ^ UPat()).named("x"), lambda x: x.ins(X86Ops.VPXOR) if x.dtype.count > 1 else None),
|
||||
# packed int binary
|
||||
((UPat(dtype=dtypes.int32s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVD) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.int64s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVQ) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.uint32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVD) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.uint64) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVQ) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.int32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRAVD) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.int8s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDB) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.int16s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDW) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.int32s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDD) if x.dtype.count > 1 else None),
|
||||
((UPat(dtype=dtypes.int64s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDQ) if x.dtype.count > 1 else None),
|
||||
(UPat(Ops.SUB, dtypes.int8s, name="x"), lambda x: x.ins(X86Ops.VPSUBB) if x.dtype.count > 1 else None),
|
||||
(UPat(Ops.SUB, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPSUBW) if x.dtype.count > 1 else None),
|
||||
(UPat(Ops.SUB, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPSUBD) if x.dtype.count > 1 else None),
|
||||
(UPat(Ops.SUB, dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPSUBQ) if x.dtype.count > 1 else None),
|
||||
(UPat(Ops.MUL, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMULLW) if x.dtype.count > 1 else None),
|
||||
(UPat(Ops.MUL, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMULLD) if x.dtype.count > 1 else None),
|
||||
# scalar int binary
|
||||
((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv),
|
||||
# scalar int binary with immediate
|
||||
(UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.arg)))),
|
||||
(UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.arg)))),
|
||||
(UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.arg)))),
|
||||
(UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
|
||||
(UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
|
||||
(UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
|
||||
(UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
|
||||
(UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
|
||||
(UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
|
||||
# scalar int binary with register
|
||||
(UPat.var("a", dtypes.ints) << UPat.var("b"), lambda a,b: a.ins(X86Ops.SHL, src=(a, b))),
|
||||
(UPat.var("a", dtypes.uints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SHR, src=(a, b))),
|
||||
(UPat.var("a", dtypes.sints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SAR, src=(a, b))),
|
||||
(UPat.var("a", dtypes.ints) + UPat.var("b"), lambda a,b: a.ins(X86Ops.ADD, src=(a, b))),
|
||||
(UPat.var("a", dtypes.ints) * UPat.var("b"), lambda a,b: a.ins(X86Ops.IMUL, src=(a, b))),
|
||||
(UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.var("b"), lambda a,b: a.ins(X86Ops.AND, src=(a, b))),
|
||||
(UPat.var("a", dtypes.ints+(dtypes.bool,)) | UPat.var("b"), lambda a,b: a.ins(X86Ops.OR, src=(a, b))),
|
||||
(UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.var("b"), lambda a,b: a.ins(X86Ops.XOR, src=(a, b))),
|
||||
(UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.var("b"))), lambda a,b: a.ins(X86Ops.SUB, src=(a, b))),
|
||||
# float binary
|
||||
((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS if x.dtype.count == 1 else X86Ops.VADDPS)),
|
||||
((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD if x.dtype.count == 1 else X86Ops.VADDPD)),
|
||||
((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS if x.dtype.count == 1 else X86Ops.VMULPS)),
|
||||
((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD if x.dtype.count == 1 else X86Ops.VMULPD)),
|
||||
(UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS if x.dtype.count == 1 else X86Ops.VSUBPS)),
|
||||
(UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD if x.dtype.count == 1 else X86Ops.VSUBPD)),
|
||||
(UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS if x.dtype.count == 1 else X86Ops.VDIVPS)),
|
||||
(UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD if x.dtype.count == 1 else X86Ops.VDIVPD)),
|
||||
# casts
|
||||
(UPat(dtype=dtypes.int32).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PS) if x.dtype.count > 1 else None),
|
||||
(UPat(dtype=dtypes.int32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PD) if x.dtype.count > 1 else None),
|
||||
(UPat(dtype=dtypes.float32).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPS2DQ) if x.dtype.count > 1 else None),
|
||||
(UPat(dtype=dtypes.float64).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPD2DQ) if x.dtype.count > 1 else None),
|
||||
(UPat(dtype=dtypes.float32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PD) if x.dtype.count > 1 else None),
|
||||
(UPat(dtype=dtypes.float64).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPD2PS) if x.dtype.count > 1 else None),
|
||||
(UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (imm(dtypes.uint8, 4),))),
|
||||
(UPat(dtype=dtypes.float16).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPH2PS)),
|
||||
(UPat(dtype=dtypes.float32).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSS2SI)),
|
||||
(UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)),
|
||||
(UPat.var("y", dtypes.float32).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSS2SD, src=(y, y))),
|
||||
(UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
|
||||
(UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
|
||||
(UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
|
||||
(UPat(dtype=dtypes.uints+(dtypes.bool,)).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVZX) if x.dtype.count == 1 else None),
|
||||
(UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.dtype.count == 1 else None),
|
||||
(UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVSX) if x.dtype.count == 1 else None),
|
||||
(UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)),
|
||||
(UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)),
|
||||
(UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),
|
||||
(UPat(dtype=dtypes.uint16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWD)),
|
||||
(UPat(dtype=dtypes.uint16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWQ)),
|
||||
(UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXDQ)),
|
||||
(UPat(dtype=dtypes.int8).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBW)),
|
||||
(UPat(dtype=dtypes.int8).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBD)),
|
||||
(UPat(dtype=dtypes.int8).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBQ)),
|
||||
(UPat(dtype=dtypes.int16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWD)),
|
||||
(UPat(dtype=dtypes.int16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWQ)),
|
||||
(UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)),
|
||||
# bitcasts
|
||||
(UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))),
|
||||
(UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), vpins),
|
||||
(UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.ins(X86Ops.VMOVD)),
|
||||
(UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.ins(X86Ops.VMOVQ)),
|
||||
(UPat(dtype=dtypes.float32).bitcast(dtypes.int32s).named("x"), lambda x: x.ins(X86Ops.VMOVDm)),
|
||||
(UPat(dtype=dtypes.float64).bitcast(dtypes.int64s).named("x"), lambda x: x.ins(X86Ops.VMOVQm)),
|
||||
# index
|
||||
(UPat(Ops.INDEX, name="x"), lambda x: x.ins(X86Ops.LEA, src=fold_address(x))),
|
||||
# TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
|
||||
# copy, load, store
|
||||
# NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted
|
||||
(UPat(Ops.COPY, dt_128bit, name="x"), lambda x: x.ins(X86Ops.VMOVUPS)),
|
||||
(UPat(Ops.COPY, dt_64bit, name="x"), lambda x: x.ins(X86Ops.VMOVSD)),
|
||||
(UPat(Ops.COPY, dt_32bit+dt_16bit, name="x"), lambda x: x.ins(X86Ops.VMOVSS)),
|
||||
(UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV)),
|
||||
(UPat(Ops.LOAD, dt_128bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVUPS, src=fold_address(a))),
|
||||
(UPat(Ops.LOAD, dt_64bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVSD, src=fold_address(a))),
|
||||
(UPat(Ops.LOAD, dt_32bit, src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.VMOVSS, src=fold_address(a))),
|
||||
(UPat(Ops.LOAD, dt_16bit, src=(UPat(name="a"),), name="x"), lambda x,a:
|
||||
x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),))),
|
||||
(UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a: x.ins(X86Ops.MOV, src=fold_address(a))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_128bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVUPSm, src=fold_address(a) + (b,))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_64bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSDm, src=fold_address(a) + (b,))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_32bit), name="x"), lambda a,b,x: x.ins(X86Ops.VMOVSSm, src=fold_address(a) + (b,))),
|
||||
(UPat.var("a").store(UPat.var("b", dt_16bit), name="x"), lambda a,b,x: x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0)))),
|
||||
(UPat.var("a").store(UPat.var("b", dtypes.ints+(dtypes.bool,)), name="x"), lambda a,b,x:
|
||||
x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))),
|
||||
# **** X86Op -> X86Op ****
|
||||
# fold loads into X86Ops that allow it, if beneficial
|
||||
(UPat(Ops.INS, src=(UPat(Ops.LOAD, src=(UPat(name="a"),), name="y"),), allow_any_len=True, name="x"), lambda ctx,y,a,x:
|
||||
x.replace(src=fold_address(a) + x.src[1:]) if x.arg in X86GroupOp.ReadMem1st and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(), UPat(Ops.LOAD, src=(UPat(name="a"),), name="y")), allow_any_len=True, name="x"), lambda ctx,y,a,x:
|
||||
x.replace(src=x.src[:1] + fold_address(a) + x.src[2:]) if x.arg in X86GroupOp.ReadMem2nd and is_foldable(ctx, x, y) else None),
|
||||
(UPat(Ops.INS, src=(UPat(), UPat(), UPat(Ops.LOAD, src=(UPat(name="a"),), name="y")), allow_any_len=True, name="x"), lambda ctx,y,a,x:
|
||||
x.replace(src=x.src[:2] + fold_address(a) + x.src[3:]) if x.arg in X86GroupOp.ReadMem3rd and is_foldable(ctx, x, y) else None),
|
||||
# allocate virtual registers
|
||||
(UPat((Ops.INS, Ops.DEFINE_REG, Ops.DEFINE_LOCAL), name="x"), alloc_vregs),
|
||||
])
|
||||
|
||||
# ***** pre register allocation *****
|
||||
# this handles flag clobbers. Unfortunately x86 doesn't have a good way to store/restore the flag register (then regalloc would handle it)
|
||||
# so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional,
|
||||
# regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack
|
||||
def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
|
||||
flag_def = x if x.arg in X86GroupOp.WriteFlags or x.op in (Ops.RANGE, Ops.END) else x.src[-1] if x.arg in X86GroupOp.ReadFlags else None
|
||||
if flag_def is None: return None
|
||||
if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock)
|
||||
ctx.lock = flag_def
|
||||
if flag_def not in ctx.clobbered: return None
|
||||
ctx.clobbered.remove(flag_def)
|
||||
return (x, [flag_def, x])
|
||||
|
||||
pre_regalloc_matcher = PatternMatcher([
|
||||
(UPat((Ops.INS, Ops.RANGE, Ops.END), name="x"), flag_rematerialize),
|
||||
])
|
||||
|
||||
# ***** post register allocation *****
|
||||
# TODO: control flow should be overhauled so that this isn't necessary
|
||||
def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
|
||||
loop_label = "_".join(str(i) for i in x.arg[:-1])
|
||||
acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
|
||||
label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}")
|
||||
cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0]))
|
||||
jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
|
||||
ctx.loop_label[acc] = loop_label
|
||||
return (acc, [acc, label, cmp, jump_out])
|
||||
|
||||
# final rewrite to match the isa spec
|
||||
post_regalloc_matcher = PatternMatcher([
|
||||
# rewrite FRAME_INDEX to IMM now that the stack size is known
|
||||
(UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])),
|
||||
# rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
|
||||
(UPat(Ops.RANGE, name="x"), lambda ctx,x: lower_range(ctx, x)),
|
||||
# rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target
|
||||
(UPat(Ops.END, name="x"), lambda ctx,x: (jmp:=UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}"),
|
||||
[x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)), jmp, UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")])),
|
||||
# rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
|
||||
(UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
|
||||
[ctx.ren.copy(x.src[0], x.reg), nx] if x.reg != x.src[0].reg else [nx]) if x.arg in X86GroupOp.TwoAddress else None),
|
||||
])
|
||||
|
||||
# ***** X86 instruction encoding *****
|
||||
|
||||
def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> bytes|None:
|
||||
def _encode(reg_uop:UOp|None, rm_uop:UOp, idx_uop:UOp|None=None, disp_uop:UOp|None=None, vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes:
|
||||
nonlocal reg, opc
|
||||
# get the encoding values of the different fields
|
||||
reg = cast(int, cast(Register, reg_uop.reg).index if reg_uop is not None else reg)
|
||||
rm = cast(Register, rm_uop.reg).index
|
||||
idx = cast(Register, idx_uop.reg).index if idx_uop is not None and idx_uop.reg is not None else 4
|
||||
rm_sz = 8 if isinstance(rm_uop.dtype, PtrDType) and disp_uop is None else rm_uop.dtype.itemsize
|
||||
reg_sz = (reg_uop.dtype.itemsize if not isinstance(reg_uop.dtype, PtrDType) else 8) if reg_uop is not None else 0
|
||||
sz = reg_sz or rm_sz
|
||||
|
||||
# encode instruction
|
||||
inst = bytes([])
|
||||
assert 0 <= reg <= 15 and 0 <= idx <= 15 and 0 <= rm <= 15
|
||||
# r extends reg field, x extends index field, b extends rm or base field
|
||||
r, _x, b = reg >> 3, idx >> 3, rm >> 3
|
||||
if sel: # VEX bytes
|
||||
vvvv = cast(Register, vvvv_uop.reg).index if vvvv_uop is not None else 0
|
||||
l = (max(reg_sz, rm_sz) > 16) & 0b1
|
||||
if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
|
||||
else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
|
||||
else: # optional PREFIX and REX bytes
|
||||
# PREFIX byte signaling 16 bit variant of instruction
|
||||
if sz == 2: inst += bytes([0x66])
|
||||
# bit signaling 64 bit variant of instruction
|
||||
w = sz == 8
|
||||
# REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed
|
||||
if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2): inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b])
|
||||
# legacy 8bit opcode is 1 less than 16-64bit variants
|
||||
if (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}: opc -= 1
|
||||
# OPCODE byte
|
||||
inst += opc.to_bytes((opc.bit_length() + 7) // 8, 'big')
|
||||
# MODRM byte
|
||||
# now we only care about the lower 3 bits
|
||||
idx, rm, reg = idx & 0b111, rm & 0b111, reg & 0b111
|
||||
# 0b00 -- signals memory access with no displacement
|
||||
# 0b01 -- signals memory access with 8bit displacement
|
||||
# 0b10 -- signals memory access with 32bit displacement
|
||||
# 0b11 -- signals no memory access
|
||||
if disp_uop is not None:
|
||||
assert disp_uop.op is Ops.CONST, "displacement must be a constant"
|
||||
assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int"
|
||||
# rbp/r13 always require a displacement
|
||||
if disp_uop.arg != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
|
||||
else: mod = 0b00
|
||||
else: mod = 0b11
|
||||
# x 0b0 and idx 0b100 means rsp which means no index exists
|
||||
# rm 0b100 (rsp/r12) signals a sib byte is required, rm then is encoded in the base field of SIB
|
||||
_rm = rm if idx == 0b100 and _x == 0b0 else 0b100
|
||||
inst += bytes([mod << 6 | reg << 3 | _rm])
|
||||
# SIB byte
|
||||
if _rm == 0b100 and mod != 0b11:
|
||||
scale = {1: 0b00, 2: 0b01, 4: 0b10, 8: 0b11}[1 if idx == 0b100 and _x == 0b0 else rm_sz]
|
||||
inst += bytes([scale << 6 | idx << 3 | rm])
|
||||
# DISP byte
|
||||
if mod == 0b01 or mod == 0b10:
|
||||
assert disp_uop is not None
|
||||
inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.arg)
|
||||
# IMM byte
|
||||
if imm_uop is not None:
|
||||
if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.arg)
|
||||
elif isinstance(imm_uop.reg, Register): inst += bytes([(imm_uop.reg.index & 0b1111) << 4 | 0b0000])
|
||||
return inst
|
||||
|
||||
# get the encoding structure of the uop
|
||||
# when a uop writes to memory it takes the form of a store, dtype is void, no definition
|
||||
address:tuple[UOp|None, ...]
|
||||
if x.arg in X86GroupOp.WriteMem:
|
||||
if len(x.src) > 3: address, rest = x.src[:3], x.src[3:]
|
||||
else: address, rest = (x, None, None), x.src
|
||||
return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *rest[:1]))
|
||||
|
||||
if x.arg in X86GroupOp.Rm1st:
|
||||
if len(x.src) > 2: address, rest = x.src[:3], x.src[3:]
|
||||
else: address, rest = (x.src[0], None, None), x.src[1:]
|
||||
imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
|
||||
return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop))
|
||||
|
||||
if x.arg in X86GroupOp.Rm2nd:
|
||||
if len(x.src) > 3: address, rest = x.src[1:4], x.src[:1] + x.src[4:]
|
||||
else: address, rest = (x.src[1], None, None), x.src[:1] + x.src[2:]
|
||||
# cmp/vucomiss reg, rm don't define a new register
|
||||
return _encode(x, *address, *rest) if x.dtype is not dtypes.void else _encode(rest[0], *address)
|
||||
|
||||
return None
|
||||
|
||||
# https://www.felixcloutier.com/x86/
|
||||
# legacy version -> VEX version
|
||||
# prefix field: None -> 0 | 66 -> 1 | F3 -> 2 | F2 -> 3
|
||||
# opcode map select: 0F -> 1 | 0F38 -> 2 | 0F3A -> 3
|
||||
encodings = {
|
||||
# moves
|
||||
X86Ops.MOVABS: lambda x:
|
||||
bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | x.reg.index >> 3, 0xB8 + (x.reg.index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].arg),
|
||||
X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
|
||||
X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
|
||||
X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
|
||||
X86Ops.VMOVSD: lambda x: encode(x, 0x10, pp=3, sel=1), X86Ops.VMOVSDm: lambda x: encode(x, 0x11, pp=3, sel=1),
|
||||
X86Ops.VMOVUPS: lambda x: encode(x, 0x10, pp=0, sel=1), X86Ops.VMOVUPSm: lambda x: encode(x, 0x11, pp=0, sel=1),
|
||||
X86Ops.VMOVD: lambda x: encode(x, 0x6E, pp=1, sel=1), X86Ops.VMOVQ: lambda x: encode(x, 0x6E, pp=1, sel=1, we=1),
|
||||
X86Ops.VMOVDm: lambda x: encode(x, 0x7E, pp=1, sel=1), X86Ops.VMOVQm: lambda x: encode(x, 0x7E, pp=1, sel=1, we=1),
|
||||
# casts
|
||||
X86Ops.MOVZX: lambda x: encode(x, 0x0FB7),
|
||||
X86Ops.MOVSX: lambda x: encode(x, 0x0FBF), X86Ops.MOVSXD: lambda x: encode(x, 0x63),
|
||||
X86Ops.VPMOVZXBW: lambda x: encode(x, 0x30, pp=1, sel=2), X86Ops.VPMOVZXBD: lambda x: encode(x, 0x31, pp=1, sel=2),
|
||||
X86Ops.VPMOVZXBQ: lambda x: encode(x, 0x32, pp=1, sel=2), X86Ops.VPMOVZXWD: lambda x: encode(x, 0x33, pp=1, sel=2),
|
||||
X86Ops.VPMOVZXWQ: lambda x: encode(x, 0x34, pp=1, sel=2), X86Ops.VPMOVZXDQ: lambda x: encode(x, 0x35, pp=1, sel=2),
|
||||
X86Ops.VPMOVSXBW: lambda x: encode(x, 0x20, pp=1, sel=2), X86Ops.VPMOVSXBD: lambda x: encode(x, 0x21, pp=1, sel=2),
|
||||
X86Ops.VPMOVSXBQ: lambda x: encode(x, 0x22, pp=1, sel=2), X86Ops.VPMOVSXWD: lambda x: encode(x, 0x23, pp=1, sel=2),
|
||||
X86Ops.VPMOVSXWQ: lambda x: encode(x, 0x24, pp=1, sel=2), X86Ops.VPMOVSXDQ: lambda x: encode(x, 0x25, pp=1, sel=2),
|
||||
X86Ops.VCVTSS2SD: lambda x: encode(x, 0x5A, pp=2, sel=1), X86Ops.VCVTSD2SS: lambda x: encode(x, 0x5A, pp=3, sel=1),
|
||||
X86Ops.VCVTPH2PS: lambda x: encode(x, 0x13, pp=1, sel=2), X86Ops.VCVTPS2PH: lambda x: encode(x, 0x1D, pp=1, sel=3),
|
||||
X86Ops.VCVTDQ2PS: lambda x: encode(x, 0x5B, pp=0, sel=1), X86Ops.VCVTDQ2PD: lambda x: encode(x, 0xE6, pp=2, sel=1),
|
||||
X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1),
|
||||
X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1),
|
||||
X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=x.src[1].dtype.itemsize == 8),
|
||||
X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=x.src[1].dtype.itemsize == 8),
|
||||
X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8),
|
||||
X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8),
|
||||
# int division
|
||||
X86Ops.IDIV: lambda x: encode(x, 0xF7, reg=7), X86Ops.DIV: lambda x: encode(x, 0xF7, reg=6),
|
||||
# scalar int binary
|
||||
X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4),
|
||||
X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7),
|
||||
X86Ops.ADD: lambda x: encode(x, 0x03), X86Ops.ADDi: lambda x: encode(x, 0x81, reg=0),
|
||||
X86Ops.SUB: lambda x: encode(x, 0x2B), X86Ops.SUBi: lambda x: encode(x, 0x81, reg=5),
|
||||
X86Ops.AND: lambda x: encode(x, 0x23), X86Ops.ANDi: lambda x: encode(x, 0x81, reg=4),
|
||||
X86Ops.XOR: lambda x: encode(x, 0x33), X86Ops.XORi: lambda x: encode(x, 0x81, reg=6),
|
||||
X86Ops.OR: lambda x: encode(x, 0x0B), X86Ops.ORi: lambda x: encode(x, 0x81, reg=1),
|
||||
X86Ops.CMP: lambda x: encode(x, 0x3B), X86Ops.CMPi: lambda x: encode(x, 0x81, reg=7),
|
||||
X86Ops.IMUL: lambda x: encode(x, 0x0FAF), X86Ops.IMULi: lambda x: encode(x, 0x69),
|
||||
X86Ops.SETB: lambda x: encode(x, 0x0F92, reg=0), X86Ops.SETL: lambda x: encode(x, 0x0F9C, reg=0),
|
||||
X86Ops.SETE: lambda x: encode(x, 0x0F94, reg=0), X86Ops.SETNE: lambda x: encode(x, 0x0F95, reg=0),
|
||||
# packed bitwise NOTE: only bitwise and packed
|
||||
X86Ops.VPAND: lambda x: encode(x, 0xDB, pp=1, sel=1), X86Ops.VPXOR: lambda x: encode(x, 0xEF, pp=1, sel=1),
|
||||
X86Ops.VPOR: lambda x: encode(x, 0xEB, pp=1, sel=1),
|
||||
# unary
|
||||
X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTPS: lambda x: encode(x, 0x51, pp=0, sel=1),
|
||||
X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1), X86Ops.VSQRTPD: lambda x: encode(x, 0x51, pp=1, sel=1),
|
||||
X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDPS: lambda x: encode(x, 0x08, pp=1, sel=3),
|
||||
X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3), X86Ops.VROUNDPD: lambda x: encode(x, 0x09, pp=1, sel=3),
|
||||
# packed int binary
|
||||
X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1),
|
||||
X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1),
|
||||
X86Ops.VPCMPGTB: lambda x: encode(x, 0x64, pp=1, sel=1), X86Ops.VPCMPGTW: lambda x: encode(x, 0x65, pp=1, sel=1),
|
||||
X86Ops.VPCMPGTD: lambda x: encode(x, 0x66, pp=1, sel=1), X86Ops.VPCMPGTQ: lambda x: encode(x, 0x37, pp=1, sel=2),
|
||||
X86Ops.VPCMPEQB: lambda x: encode(x, 0x74, pp=1, sel=1), X86Ops.VPCMPEQW: lambda x: encode(x, 0x75, pp=1, sel=1),
|
||||
X86Ops.VPCMPEQD: lambda x: encode(x, 0x76, pp=1, sel=1), X86Ops.VPCMPEQQ: lambda x: encode(x, 0x29, pp=1, sel=2),
|
||||
X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2),
|
||||
X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1),
|
||||
X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1),
|
||||
X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1),
|
||||
X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1),
|
||||
X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2),
|
||||
# float cmp
|
||||
X86Ops.VUCOMISS: lambda x: encode(x, 0x2E, pp=0, sel=1), X86Ops.VUCOMISD: lambda x: encode(x, 0x2E, pp=1, sel=1),
|
||||
# scalar / packed float binary
|
||||
X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1),
|
||||
X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1),
|
||||
X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBPS: lambda x: encode(x, 0x5C, pp=0, sel=1),
|
||||
X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1), X86Ops.VSUBPD: lambda x: encode(x, 0x5C, pp=1, sel=1),
|
||||
X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULPS: lambda x: encode(x, 0x59, pp=0, sel=1),
|
||||
X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1), X86Ops.VMULPD: lambda x: encode(x, 0x59, pp=1, sel=1),
|
||||
X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVPS: lambda x: encode(x, 0x5E, pp=0, sel=1),
|
||||
X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1), X86Ops.VDIVPD: lambda x: encode(x, 0x5E, pp=1, sel=1),
|
||||
X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPPS: lambda x: encode(x, 0xC2, pp=0, sel=1),
|
||||
X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1), X86Ops.VCMPPD: lambda x: encode(x, 0xC2, pp=1, sel=1),
|
||||
X86Ops.VMAXSS: lambda x: encode(x, 0x5F, pp=2, sel=1), X86Ops.VMAXPS: lambda x: encode(x, 0x5F, pp=0, sel=1),
|
||||
X86Ops.VMAXSD: lambda x: encode(x, 0x5F, pp=3, sel=1), X86Ops.VMAXPD: lambda x: encode(x, 0x5F, pp=1, sel=1),
|
||||
X86Ops.VMINSS: lambda x: encode(x, 0x5D, pp=2, sel=1), X86Ops.VMINPS: lambda x: encode(x, 0x5D, pp=0, sel=1),
|
||||
X86Ops.VMINSD: lambda x: encode(x, 0x5D, pp=3, sel=1), X86Ops.VMINPD: lambda x: encode(x, 0x5D, pp=1, sel=1),
|
||||
# ternary
|
||||
X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C),
|
||||
X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45),
|
||||
X86Ops.VFMADD213SS: lambda x: encode(x, 0xA9, pp=1, sel=2), X86Ops.VFMADD213SD: lambda x: encode(x, 0xA9, pp=1, sel=2, we=1),
|
||||
X86Ops.VFMADD213PS: lambda x: encode(x, 0xA8, pp=1, sel=2), X86Ops.VFMADD213PD: lambda x: encode(x, 0xA8, pp=1, sel=2, we=1),
|
||||
X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3),
|
||||
X86Ops.VPBLENDVB: lambda x: encode(x, 0x4C, pp=1, sel=3),
|
||||
# shuffles
|
||||
X86Ops.VPBROADCASTB: lambda x: encode(x, 0x78, pp=1, sel=2), X86Ops.VPBROADCASTW: lambda x: encode(x, 0x79, pp=1, sel=2),
|
||||
X86Ops.VPBROADCASTD: lambda x: encode(x, 0x58, pp=1, sel=2), X86Ops.VPBROADCASTQ: lambda x: encode(x, 0x59, pp=1, sel=2),
|
||||
X86Ops.VBROADCASTSS: lambda x: encode(x, 0x18, pp=1, sel=2), X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1),
|
||||
X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1),
|
||||
X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3), X86Ops.VPINSRQ: lambda x: encode(x, 0x22, pp=1, sel=3, we=1),
|
||||
X86Ops.VSHUFPS: lambda x: encode(x, 0xC6, pp=0, sel=1), X86Ops.VSHUFPD: lambda x: encode(x, 0xC6, pp=1, sel=1),
|
||||
X86Ops.VINSERTPS: lambda x: encode(x, 0x21, pp=1, sel=3),
|
||||
# extract
|
||||
X86Ops.VPEXTRB: lambda x: encode(x, 0x14, pp=1, sel=3), X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3),
|
||||
X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3), X86Ops.VPEXTRQ: lambda x: encode(x, 0x16, pp=1, sel=3, we=1),
|
||||
# jumps are encoded with a placeholder which gets patched later once the real offset is known
|
||||
X86Ops.JE: lambda x: bytes([0x0F, 0x84]) + int(0).to_bytes(4),
|
||||
X86Ops.JNE: lambda x: bytes([0x0F, 0x85]) + int(0).to_bytes(4),
|
||||
X86Ops.JL: lambda x: bytes([0x0F, 0x8C]) + int(0).to_bytes(4),
|
||||
X86Ops.JB: lambda x: bytes([0x0F, 0x82]) + int(0).to_bytes(4),
|
||||
X86Ops.JGE: lambda x: bytes([0x0F, 0x8D]) + int(0).to_bytes(4),
|
||||
X86Ops.JMP: lambda x: bytes([0xE9]) + int(0).to_bytes(4),
|
||||
X86Ops.RET: lambda x: bytes([0xC3]),
|
||||
}
|
||||
|
||||
class X86Renderer(ISARenderer):
|
||||
device = "CPU"
|
||||
has_local = False
|
||||
has_threads = bool(getenv("THREADS", 1))
|
||||
global_max = (CPU_COUNT.value, 0, 0)
|
||||
extra_matcher = extra_matcher
|
||||
pre_isel_matcher = pre_isel_matcher
|
||||
isel_matcher = isel_matcher
|
||||
pre_regalloc_matcher = pre_regalloc_matcher
|
||||
post_regalloc_matcher = post_regalloc_matcher
|
||||
code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_cpu import X86Compiler
|
||||
self.compiler = X86Compiler()
|
||||
def is_two_address(self, x:UOp) -> bool: return x.arg in X86GroupOp.TwoAddress
|
||||
def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
|
||||
# nasty hacks to deal with pointers TODO: rm pointers
|
||||
def copy(self, x:UOp, reg:Register):
|
||||
dt = dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype
|
||||
ret = isel_matcher.rewrite(UOp(Ops.COPY, dt, (x,), tag=reg))
|
||||
assert ret is not None
|
||||
return ret.replace(dtype=x.dtype)
|
||||
|
||||
def spill(self, disp:UOp, x:UOp) -> UOp:
|
||||
nx = x.replace(dtype=dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype)
|
||||
ret = isel_matcher.rewrite(self.stack_pointer().index(disp).store(nx))
|
||||
assert ret is not None
|
||||
return ret.replace(src=(s if s is not nx else x for s in ret.src))
|
||||
|
||||
def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
|
||||
ndt = dtypes.uint64 if isinstance(x.dtype, PtrDType) else x.dtype
|
||||
ret = isel_matcher.rewrite(self.stack_pointer().index(disp).load(dtype=ndt, tag=reg))
|
||||
assert ret is not None
|
||||
return ret.replace(dtype=x.dtype)
|
||||
|
||||
def asm_str(self, uops:list[UOp], function_name:str) -> str:
|
||||
def _format_op(x:UOp) -> str: return f" {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}"
|
||||
def _format_operands(x:UOp) -> str:
|
||||
def _format(src:tuple[UOp, ...]) -> list[str]:
|
||||
return [str(s.arg) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize if not isinstance(s.dtype, PtrDType) else 8, o) if \
|
||||
(o:=str(s.reg)) in reg_strs else o for s in src if s.reg is not None]
|
||||
def _mem_adress(base:UOp, idx:UOp, disp:UOp) -> list[str]:
|
||||
return [f"[{base.reg}" + (f" + {idx.reg}*{base.dtype.itemsize}" if idx.reg else "") + (f" + {disp.arg}" if disp.arg else "") + "]"]
|
||||
|
||||
if len(x.src) > 3 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:3]) + _format(x.src[3:])
|
||||
elif len(x.src) > 2 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:3]) + _format(x.src[3:])
|
||||
elif len(x.src) > 3 and x.arg in X86GroupOp.Rm2nd: ret = _format((x, x.src[0])) + _mem_adress(*x.src[1:4]) + _format(x.src[4:])
|
||||
else: ret = _format((x,) + x.src)
|
||||
return ", ".join(ret)
|
||||
|
||||
asm = [f".{function_name}:"]
|
||||
for u in uops:
|
||||
if u.op is not Ops.INS: continue
|
||||
if u.arg is X86Ops.LABEL: asm.append(f"{str(u.tag)}:")
|
||||
elif u.arg is X86Ops.RET: asm.append(_format_op(u))
|
||||
else: asm.append(_format_op(u) + " " + _format_operands(u))
|
||||
return "\n".join(asm)
|
||||
|
||||
def render(self, uops:list[UOp]) -> str:
|
||||
targets: dict[str, int] = {}
|
||||
jumps: dict[UOp, int] = {}
|
||||
binary = bytearray()
|
||||
for u in uops:
|
||||
if u.op is not Ops.INS: continue
|
||||
if u.arg is X86Ops.LABEL:
|
||||
targets[u.tag] = len(binary)
|
||||
continue
|
||||
if u.arg not in encodings or (l:=encodings[u.arg](u)) is None:
|
||||
raise RuntimeError(f"failed to encode {u.arg} with {u.dtype} srcs {[x.dtype for x in u.src]}")
|
||||
binary.extend(l)
|
||||
if u.arg in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary)
|
||||
# fixup jump targets now that encoding size is known
|
||||
for u in uops:
|
||||
if (t:=jumps.get(u)) is not None: binary[t-4:t] = (targets[u.tag] - t).to_bytes(4, 'little', signed=True)
|
||||
return binary.hex()
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes() if d not in dtypes.fp8s+(dtypes.bfloat16,)}
|
||||
300
tinygrad_repo/tinygrad/renderer/llvmir.py
Normal file
300
tinygrad_repo/tinygrad/renderer/llvmir.py
Normal file
@@ -0,0 +1,300 @@
|
||||
from typing import cast
|
||||
import math, struct, sys
|
||||
from tinygrad.codegen.opt import tc
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.renderer.cstyle import HIPRenderer, create_non_native_float_pats, pm_manual_bf16_cast
|
||||
from tinygrad.uop.decompositions import xexp2, xlog2
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
|
||||
from tinygrad.dtype import dtypes, float_to_fp8, DType, PtrDType, truncate
|
||||
from tinygrad.helpers import prod, Target, CPU_COUNT, getenv, OSX
|
||||
|
||||
def ldt(dt:DType):
|
||||
if dt.vcount > 1: return f"<{dt.vcount} x {ldt(dt.scalar())}>"
|
||||
if isinstance(dt, PtrDType): return ldt(dt.base) + "*"
|
||||
return {dtypes.void: "void", dtypes.bool: "i1", dtypes.int8: "i8", dtypes.int16: "i16", dtypes.int32: "i32", dtypes.int64: "i64",
|
||||
dtypes.uint8: "i8", dtypes.uint16: "i16", dtypes.uint32: "i32", dtypes.uint64: "i64", dtypes.fp8e4m3: "i8", dtypes.fp8e5m2: "i8",
|
||||
dtypes.float16: "half", dtypes.bfloat16: "bfloat", dtypes.float32: "float", dtypes.float64: "double"}[dt]
|
||||
|
||||
def lconst(x, dtype:DType):
|
||||
if dtype in dtypes.floats:
|
||||
if dtype in dtypes.fp8s: return float_to_fp8(x, dtype)
|
||||
if math.isinf(x) or math.isnan(x): return "0x%02X%02X%02X%02X%02X%02X%02X%02X" % tuple(struct.pack("d",x)[::-1])
|
||||
return truncate[dtype](x)
|
||||
return int(x)
|
||||
|
||||
def lcast(input_type:DType, output_type:DType):
|
||||
if dtypes.is_float(input_type):
|
||||
if dtypes.is_float(output_type): return 'fpext' if output_type.itemsize > input_type.itemsize else 'fptrunc'
|
||||
if dtypes.is_int(output_type): return 'fptoui' if dtypes.is_unsigned(output_type) else 'fptosi'
|
||||
if dtypes.is_unsigned(input_type) or dtypes.is_bool(input_type):
|
||||
if dtypes.is_float(output_type): return 'uitofp'
|
||||
if dtypes.is_int(output_type): return 'trunc' if output_type.itemsize < input_type.itemsize else 'zext'
|
||||
if dtypes.is_int(input_type):
|
||||
if dtypes.is_float(output_type): return 'sitofp'
|
||||
if dtypes.is_int(output_type): return 'trunc' if output_type.itemsize < input_type.itemsize else 'sext'
|
||||
raise NotImplementedError(f"cast from {input_type} -> {output_type} not implemented")
|
||||
|
||||
# https://github.com/corsix/amx
|
||||
def render_wmma_amx(ctx, wmma: UOp) -> str:
|
||||
def AMX(op, gpr): return f'call void asm sideeffect ".word (0x201000+($0<<5)+0$1-((0$1>>4)*6))", "i,r,~{{memory}}"(i32 {op}, i64 {gpr}) #0; AMX'
|
||||
|
||||
return "\n".join([
|
||||
*[f' store {ldt(src.dtype)} {ctx[src]}, {ldt(src.dtype.ptr())} {ctx[wmma]}_amx{i}, align {src.dtype.itemsize}' for i,src in enumerate(wmma.src)],
|
||||
f' call void asm sideeffect "nop\\0Anop\\0Anop\\0A.word ({0x201000 + (17 << 5) + 0})", "~{{memory}}"() #0; AMX set', # set
|
||||
*[f' {ctx[wmma]}_ld{i} = add i64 {ctx[wmma]}_ptr_amx2, {i*4<<56 | i*64}\n {AMX(4,f"{ctx[wmma]}_ld{i}")} ldz' for i in range(16)], # ldz
|
||||
f' {AMX(0, f"{ctx[wmma]}_ptr_amx1")} ldx\n {AMX(1, f"{ctx[wmma]}_ptr_amx0")} ldy\n {AMX(12, 0)} fma32', # ldx ldy fma
|
||||
*[f' {ctx[wmma]}_st{i} = add i64 {ctx[wmma]}_ptr_amx2, {i*4<<56 | i*64}\n {AMX(5,f"{ctx[wmma]}_st{i}")} stz' for i in range(16)], # stz
|
||||
f' call void asm sideeffect "nop\\0Anop\\0Anop\\0A.word ({0x201000 + (17 << 5) + 1})", "~{{memory}}"() #0; AMX clr', # clr
|
||||
f' {ctx[wmma]} = load {ldt(wmma.dtype)}, ptr {ctx[wmma]}_amx2, align {wmma.dtype.itemsize}'])
|
||||
|
||||
def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
|
||||
dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16",
|
||||
dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8"}
|
||||
# https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
|
||||
N,M,K = wmma.arg[1]
|
||||
if cdna:
|
||||
if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"})
|
||||
return f" {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
|
||||
f".{N}x{M}x{K}{dt_map[wmma.arg[2]]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
|
||||
# https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
|
||||
# example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
|
||||
return f" {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \
|
||||
f"{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + (", i1 false)" \
|
||||
if wmma.dtype.scalar() != dtypes.float else ")")
|
||||
|
||||
# llvm ops, lop[<dtype>][<op>]
|
||||
unsigned_lop = { Ops.ADD: "add", Ops.MUL: "mul", Ops.CDIV: "udiv", Ops.CMOD: "urem",
|
||||
Ops.CMPLT: "icmp ult", Ops.CMPNE: "icmp ne", Ops.CMPEQ: "icmp eq", Ops.OR: "or", Ops.AND: "and", Ops.XOR: "xor",
|
||||
Ops.SHL: "shl", Ops.SHR: "lshr",}
|
||||
signed_lop = {**unsigned_lop, Ops.ADD: "add nsw", Ops.CMPLT: "icmp slt", Ops.CDIV: "sdiv", Ops.CMOD: "srem", Ops.SHR: "ashr"}
|
||||
flags = " nsz arcp contract afn"
|
||||
float_lop = {Ops.ADD: "fadd"+flags, Ops.MUL: "fmul"+flags, Ops.CMPLT: f"fcmp{flags} olt",
|
||||
Ops.CMPNE: f"fcmp{flags} une", Ops.CMPEQ: f"fcmp{flags} oeq", Ops.FDIV: "fdiv"+flags}
|
||||
lop = {**{x:unsigned_lop for x in (dtypes.bool,)+dtypes.uints}, **{x:signed_lop for x in dtypes.sints}, **{x:float_lop for x in dtypes.floats}}
|
||||
|
||||
base_rewrite = PatternMatcher([
|
||||
# memory load/store
|
||||
(UPat(Ops.INDEX, name="x"), lambda ctx,x:
|
||||
f" {ctx[x]} = getelementptr inbounds {ldt(x.dtype.base)}, {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"),
|
||||
(UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"),
|
||||
lambda ctx,x,idx,alt,mask:
|
||||
f" br label {ctx[x]}_entry\n{ctx[x][1:]}_entry:\n"
|
||||
f" br i1 {ctx[mask]}, label {ctx[x]}_load, label {ctx[x]}_exit\n{ctx[x][1:]}_load:\n"
|
||||
f" {ctx[x]}_yes = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}\n"
|
||||
f" br label {ctx[x]}_exit\n{ctx[x][1:]}_exit:\n"
|
||||
f" {ctx[x]} = phi {ldt(x.dtype)} [{ctx[x]}_yes, {ctx[x]}_load], [{ctx[alt]}, {ctx[x]}_entry]"),
|
||||
(UPat.var('idx').load(name="x"), lambda ctx,x,idx: f" {ctx[x]} = load {ldt(x.dtype)}, {ldt(idx.dtype)} {ctx[idx]}"),
|
||||
(UPat.var('idx').store(UPat.var("var")), lambda ctx,idx,var: f" store {ldt(var.dtype)} {ctx[var]}, {ldt(idx.dtype)} {ctx[idx]}"),
|
||||
|
||||
# GEP/VECTORIZE/CAST for float4 support
|
||||
(UPat(Ops.GEP, name="x"), lambda ctx,x: f" {ctx[x]} = extractelement {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, i32 {x.arg[0]}"),
|
||||
(UPat(Ops.STACK, src=UPat.var('y'), name="x"), lambda ctx,x,y:
|
||||
f" {ctx[x]}_z = insertelement <1 x {ldt(y.dtype)}> poison, {ldt(y.dtype)} {ctx[y]}, i32 0\n"
|
||||
f" {ctx[x]} = shufflevector <1 x {ldt(y.dtype)}> {ctx[x]}_z, <1 x {ldt(y.dtype)}> poison, <{x.dtype.count} x i32> zeroinitializer"),
|
||||
(UPat(Ops.STACK, name="x"), lambda ctx,x: "\n".join([(f" {ctx[x]}_{i}" if i+1 != len(x.src) else f" {ctx[x]}")+
|
||||
f" = insertelement {ldt(x.dtype)} "+(f"{ctx[x]}_{i-1}" if i != 0 else "poison")+
|
||||
f", {ldt(u.dtype)} {ctx[u]}, i32 {i}" for i,u in enumerate(x.src)])),
|
||||
# unary/binary/ternary ops
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f" {ctx[x]} = bitcast {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: f" {ctx[x]} = {lcast(x.src[0].dtype, x.dtype)} {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
|
||||
(UPat(Ops.TRUNC, name="x"),
|
||||
lambda ctx,x: f" {ctx[x]} = call {ldt(x.dtype)} @llvm.trunc.{ldt(x.dtype.scalar())}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),
|
||||
(UPat(GroupOp.Binary, name="x"), lambda ctx,x:
|
||||
f" {ctx[x]} = {lop[x.src[0].dtype.scalar()][x.op]} {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ctx[x.src[1]]}"),
|
||||
(UPat(Ops.WHERE, name="x"), lambda ctx,x:
|
||||
f" {ctx[x]} = select {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[2].dtype)} {ctx[x.src[2]]}"),
|
||||
|
||||
# range
|
||||
(UPat(Ops.RANGE, name="r"), lambda ctx,r:
|
||||
f" br label %loop_entry_{range_str(r)}\n"
|
||||
f"loop_entry_{range_str(r)}:\n"
|
||||
f" br label %loop_latch_{range_str(r)}\n"
|
||||
f"loop_latch_{range_str(r)}:\n"
|
||||
f" {ctx[r]} = phi {ldt(r.dtype)} [ 0, %loop_entry_{range_str(r)} ], [ {ctx[r]}phi, %loop_footer_{range_str(r)} ]\n"
|
||||
f" {ctx[r]}phi = add {ldt(r.dtype)} {ctx[r]}, 1\n"
|
||||
f" {ctx[r]}cmp = icmp ult {ldt(r.dtype)} {ctx[r]}, {ctx[r.src[0]]}\n"
|
||||
f" br i1 {ctx[r]}cmp, label %loop_body_{range_str(r)}, label %loop_exit_{range_str(r)}\n"
|
||||
f"loop_body_{range_str(r)}:"),
|
||||
(UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, name="r"))), lambda r:
|
||||
f" br label %loop_footer_{range_str(r)}\n"
|
||||
f"loop_footer_{range_str(r)}:\n"
|
||||
f" br label %loop_latch_{range_str(r)}\n"
|
||||
f"loop_exit_{range_str(r)}:"),
|
||||
|
||||
# if
|
||||
(UPat(Ops.IF, name="x"), lambda ctx,x: f" br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"),
|
||||
(UPat(Ops.ENDIF, name="x"), lambda ctx,x: f" br label %ifskip_{ctx[x.src[0]][1:]}\nifskip_{ctx[x.src[0]][1:]}:"),
|
||||
|
||||
(UPat(Ops.BARRIER), lambda ctx: "")
|
||||
])
|
||||
|
||||
class LLVMRenderer(Renderer):
|
||||
supports_float4 = True
|
||||
abi: str | None
|
||||
string_rewrite: PatternMatcher
|
||||
code_for_op = {k:lambda:None for v in lop.values() for k in v.keys()}
|
||||
|
||||
extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
|
||||
def _render_fn(self, name:str, args:list[tuple[str,DType]], kernel:list[str], prefix:list[str]|None=None) -> str:
|
||||
# NOTE: CPUAllocator promises 0x20 alignment
|
||||
sargs = ", ".join([f"{ldt(dt)}{' noalias align 32' if isinstance(dt, PtrDType) else ''} {name}" for name,dt in args])
|
||||
return "\n".join((prefix or []) + [f"define{' ' + self.abi if self.abi else ''} void @{name}({sargs}) #0", "{"] + kernel + [" ret void\n}"])
|
||||
def _render_kernel(self, uops: list[UOp], prefix:list[str]|None=None) -> tuple[tuple[str, ...], str]:
|
||||
r: dict[UOp, str] = {}
|
||||
args: list[tuple[str, DType]] = []
|
||||
kernel: list[str] = []
|
||||
vc = -1
|
||||
|
||||
local_args: list[str] = []
|
||||
for u in uops:
|
||||
if self.tensor_cores == tc.amx and u.op is Ops.WMMA: # prealloc aux buffers as AMX can only load from memory
|
||||
vc += 1
|
||||
r[u] = f"%wmma{vc}"
|
||||
for i, dtype in enumerate(u.arg[2].vec(sz) for sz in [prod(size for _, size in upcast) for upcast in u.arg[6]]):
|
||||
kernel += [f" {r[u]}_amx{i} = alloca {ldt(dtype)}, align {dtype.itemsize}",
|
||||
f" {r[u]}_ptr_amx{i} = ptrtoint {ldt(dtype.ptr())} {r[u]}_amx{i} to i64"]
|
||||
|
||||
name = "test"
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP}: continue
|
||||
if u.op is Ops.AFTER:
|
||||
r[u] = r[u.src[0]]
|
||||
continue
|
||||
if u.op is Ops.SINK:
|
||||
if u.arg is not None: name = u.arg.function_name
|
||||
continue
|
||||
if u.op in (Ops.PARAM, Ops.DEFINE_VAR):
|
||||
r[u] = f"%data{u.arg.slot}" if u.op is Ops.PARAM else f"%{u.expr}"
|
||||
args.append((r[u], u.dtype))
|
||||
elif u.op in (Ops.DEFINE_LOCAL, Ops.DEFINE_REG):
|
||||
r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}"
|
||||
assert isinstance(u.dtype, PtrDType)
|
||||
if u.op is Ops.DEFINE_REG:
|
||||
kernel.append(f" {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}]")
|
||||
elif self.has_local:
|
||||
local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{u.dtype.size} x {ldt(u.dtype)}] undef, align 16")
|
||||
kernel.append(f" {r[u]} = addrspacecast [{u.dtype.size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{u.dtype.size} x {ldt(u.dtype)}]*")
|
||||
else:
|
||||
kernel.append(f" {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}], align 16")
|
||||
elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype)
|
||||
elif u.op is Ops.CAST and (ldt(u.dtype) == ldt(u.src[0].dtype) or isinstance(u.dtype, PtrDType)):
|
||||
r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast
|
||||
else:
|
||||
# if it's an assign target, it's already preallocated
|
||||
if u not in r:
|
||||
vc += 1
|
||||
r[u] = f"%v{vc}"
|
||||
|
||||
# do the rendering of the llvm ir code
|
||||
if (l:=self.string_rewrite.rewrite(u, ctx=r)) is None:
|
||||
raise RuntimeError(f"failed to render {u.op} with {u.dtype} srcs {[x.dtype for x in u.src]}")
|
||||
kernel.append(cast(str, l))
|
||||
return tuple(local_args), self._render_fn(name, args, kernel, prefix)
|
||||
|
||||
class CPULLVMRenderer(LLVMRenderer):
|
||||
has_local = False
|
||||
has_threads = bool(getenv("THREADS", 1))
|
||||
global_max = (CPU_COUNT.value, 0, 0)
|
||||
abi = 'win64cc' if sys.platform == 'win32' else None
|
||||
string_rewrite = base_rewrite + PatternMatcher([(UPat(Ops.WMMA, name="wmma"), render_wmma_amx)])
|
||||
def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
|
||||
def _render_footer(self, uops: list[UOp]) -> str: return 'attributes #0 = { alwaysinline nounwind "no-builtins" "no-trapping-math"="true" }'
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
|
||||
if "AMX" in target.arch: self.tensor_cores = tc.amx
|
||||
self.compiler = CPULLVMCompiler([x for x in target.arch.split(",") if x != "AMX"])
|
||||
|
||||
# FIXME: fp16 works on non-osx, but only if the cpu supports it
|
||||
def supported_dtypes(self):
|
||||
return {d for d in super().supported_dtypes() if
|
||||
(d != dtypes.bfloat16 or self.target.arch.startswith(("x86", "arm"))) and (d != dtypes.half or OSX) and d not in dtypes.fp8s}
|
||||
|
||||
barrier = 'fence syncscope("workgroup") release\ntail call void @llvm.amdgcn.s.barrier()\nfence syncscope("workgroup") acquire\n'
|
||||
code_for_workitem = {"g": lambda x: f"tail call i32 @llvm.amdgcn.workgroup.id.{chr(120+int(x))}()",
|
||||
"l": lambda x: f"tail call i32 @llvm.amdgcn.workitem.id.{chr(120+int(x))}()"}
|
||||
# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#llvm-ir-intrinsics
|
||||
llvm_intrinsics = {Ops.SQRT: "sqrt", Ops.LOG2: "log2", Ops.EXP2: "exp2"}
|
||||
class AMDLLVMRenderer(LLVMRenderer):
|
||||
has_local = True
|
||||
shared_max = HIPRenderer.shared_max
|
||||
global_max = HIPRenderer.global_max
|
||||
global_prod_max = HIPRenderer.global_prod_max
|
||||
abi = "amdgpu_kernel"
|
||||
code_for_op = {**LLVMRenderer.code_for_op, **{op: lambda: None for op in llvm_intrinsics}}
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.SPECIAL, name="x"), lambda ctx, x: f" {ctx[x]} = " + f"{ code_for_workitem[x.arg[0]](x.arg[-1])}; "),
|
||||
(UPat(tuple(llvm_intrinsics), name="x"),
|
||||
lambda ctx, x: f" {ctx[x]} = call {ldt(x.dtype)} @llvm.{llvm_intrinsics[x.op]}.{ldt(x.dtype.scalar())}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),
|
||||
(UPat(Ops.BARRIER), lambda ctx: barrier),
|
||||
(UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x:
|
||||
f" {ctx[x]} = call i8 @f32_to_fp8({ldt(x.src[0].dtype)} {ctx[x.src[0]]}, i1 {'1' if x.dtype == dtypes.fp8e5m2 else '0'})"),
|
||||
(UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), lambda ctx,x,y:
|
||||
f" {ctx[x.src[0]]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n"
|
||||
f" {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x.src[0]]}_i32, i32 0)"),
|
||||
]) + base_rewrite
|
||||
extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([
|
||||
(UPat(Ops.CAST, dtype=dtypes.half.vec(16), src=UPat.var("y", dtypes.half.vec(8))),
|
||||
lambda y: UOp(Ops.STACK, dtypes.half.vec(16), tuple(y.gep(i // 2) if i % 2 == 0 else UOp.const(dtypes.half, 0.0) for i in range(16)))),
|
||||
(UPat(Ops.CAST, dtype=dtypes.half.vec(8), src=UPat.var("y", dtypes.half.vec(16))),
|
||||
lambda y: UOp(Ops.STACK, dtypes.half.vec(8), tuple(y.gep(i * 2) for i in range(8)))),
|
||||
# amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
|
||||
(UPat(Ops.LOG2, dtype=dtypes.double, src=(UPat.var("d"),)), xlog2),
|
||||
(UPat(Ops.EXP2, dtype=dtypes.double, src=(UPat.var("d"),)), xexp2),
|
||||
])
|
||||
def asm(self, prg: UOp, lin: UOp) -> bytes:
|
||||
from tinygrad.renderer.amd.elf import assemble_linear
|
||||
return assemble_linear(prg, lin, self.target.arch)
|
||||
def render(self, uops: list[UOp]) -> str:
|
||||
prefix = ["""define i8 @f32_to_fp8(float %val, i1 %is_bf8) {
|
||||
entry: %ival = bitcast float %val to i32\n %exp = and i32 %ival, 2139095040\n %is_special = icmp eq i32 %exp, 2139095040
|
||||
br i1 %is_special, label %select_clip, label %clip
|
||||
clip: br i1 %is_bf8, label %bf8_clip, label %fp8_clip
|
||||
bf8_clip: %clamped_bf8 = call float @llvm.amdgcn.fmed3.f32(float %val, float 57344.0, float -57344.0)\n br label %select_clip
|
||||
fp8_clip: %clamped_fp8 = call float @llvm.amdgcn.fmed3.f32(float %val, float 448.0, float -448.0) \n br label %select_clip
|
||||
select_clip: %phi_val = phi float [%val, %entry], [%clamped_bf8, %bf8_clip], [%clamped_fp8, %fp8_clip]\n br i1 %is_bf8, label %do_bf8, label %do_fp8
|
||||
do_bf8: %packed_bf8 = call i32 @llvm.amdgcn.cvt.pk.bf8.f32(float %phi_val, float %phi_val, i32 0, i1 false)\n br label %exit
|
||||
do_fp8: %packed_fp8 = call i32 @llvm.amdgcn.cvt.pk.fp8.f32(float %phi_val, float %phi_val, i32 0, i1 false)\n br label %exit
|
||||
exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc = trunc i32 %packed to i8\n ret i8 %trunc
|
||||
}""".replace(": ", ":\n ")] if any(u.dtype in dtypes.fp8s for u in uops) else []
|
||||
return "\n".join((k:=self._render_kernel(uops, prefix))[0] + (k[1], self._render_footer(uops)))
|
||||
def _render_footer(self, uops: list[UOp]) -> str:
|
||||
# TODO: this is copied from cstyle
|
||||
local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
|
||||
requiredMaxThreadsPerBlock = prod([d.vmax for d in local_dims])
|
||||
attributes = ["alwaysinline", "nounwind", '"no-builtins"',
|
||||
f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
|
||||
return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
|
||||
self.compiler, self.tensor_cores, self.is_cdna = AMDLLVMCompiler(target.arch), tc.get_amd(target.arch), HIPRenderer.is_cdna(target.arch)
|
||||
self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
|
||||
if self.is_cdna:
|
||||
self.extra_matcher += PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(4)),
|
||||
lambda x: UOp(Ops.WMMA, dtypes.float.vec(4), (x.src[0].bitcast(dtypes.uint16.vec(4)), x.src[1].bitcast(dtypes.uint16.vec(4)),
|
||||
x.src[2]), (*x.arg,)) if x.src[0].dtype == dtypes.bfloat16.vec(4) else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(4)),
|
||||
lambda x: UOp(Ops.WMMA, dtypes.float.vec(4), (x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
|
||||
x.src[2]), (*x.arg,)) if x.src[0].dtype in (dtypes.fp8e4m3.vec(8), dtypes.fp8e5m2.vec(8)) else None),
|
||||
])
|
||||
if target.arch in {"gfx1100", "gfx1151"}:
|
||||
self.extra_matcher += PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.half.vec(8)),
|
||||
lambda x: UOp(Ops.WMMA, dtypes.half.vec(16), (x.src[0], x.src[1], x.src[2].cast(dtypes.half.vec(16))), (*x.arg,)).cast(dtypes.half.vec(8))),
|
||||
(UPat(Ops.WMMA, name="x"), lambda x: UOp(Ops.WMMA, x.dtype, (x.src[0].bitcast(dtypes.uint16.vec(16)), x.src[1].bitcast(dtypes.uint16.vec(16)),
|
||||
x.src[2]), x.arg) if x.src[0].dtype == dtypes.bfloat16.vec(16) else None),
|
||||
])
|
||||
if target.arch in {"gfx1200", "gfx1201"}:
|
||||
self.extra_matcher += PatternMatcher([
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16.vec(8)), lambda x: UOp(Ops.WMMA, dtypes.uint16.vec(8),
|
||||
(x.src[0].bitcast(dtypes.uint16.vec(8)), x.src[1].bitcast(dtypes.uint16.vec(8)), x.src[2].bitcast(dtypes.uint16.vec(8))), (*x.arg,))
|
||||
.bitcast(dtypes.bfloat16.vec(8)) if x.src[0].dtype == dtypes.bfloat16.vec(8) else None),
|
||||
(UPat(Ops.WMMA, name="x", dtype=dtypes.float.vec(8)),
|
||||
lambda x: UOp(Ops.WMMA, dtypes.float.vec(8), (x.src[0].bitcast(dtypes.uint16.vec(8)), x.src[1].bitcast(dtypes.uint16.vec(8)),
|
||||
x.src[2]), (*x.arg,)) if x.src[0].dtype == dtypes.bfloat16.vec(8) else None)
|
||||
])
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes()
|
||||
if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}
|
||||
305
tinygrad_repo/tinygrad/renderer/nir.py
Normal file
305
tinygrad_repo/tinygrad/renderer/nir.py
Normal file
@@ -0,0 +1,305 @@
|
||||
from typing import Callable, cast, Any
|
||||
from tinygrad.dtype import AddrSpace, DType, PtrDType, ImageDType, dtypes, truncate
|
||||
from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.renderer.cstyle import CUDARenderer, OpenCLRenderer
|
||||
from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
|
||||
from tinygrad.runtime.autogen import mesa
|
||||
from tinygrad.runtime.support.c import POINTER
|
||||
import base64, ctypes, ctypes.util, struct, functools, inspect, itertools
|
||||
|
||||
def g(s:str): return getattr(mesa, s)
|
||||
def nsrc(d:mesa.nir_def) -> mesa.nir_src: return mesa.nir_src(ssa=ctypes.pointer(d))
|
||||
|
||||
def glsl_type(t:DType): return mesa.glsl_array_type(glsl_type(t.base), t.size, 0).contents if isinstance(t, PtrDType) else {
|
||||
**{getattr(dtypes,k):g(f"glsl_type_builtin_{v}") for k,v in [('double','double'),('float','float'),('float16','float16_t'),('bool','uint8_t')]},
|
||||
**{d:g(f"glsl_type_builtin_{'u' * (d in dtypes.uints)}int{str(d.bitsize)+'_t' if d.itemsize != 4 else ''}") for d in dtypes.ints}}[t]
|
||||
|
||||
# alu ops, aop[<dtype>][<op>]
|
||||
u_aop = { Ops.ADD: "iadd", Ops.MUL: "imul", Ops.CDIV: "udiv", Ops.CMOD: "umod", Ops.CMPLT: "ult", Ops.CMPNE: "ine", Ops.CMPEQ: "ieq", Ops.OR: "ior",
|
||||
Ops.AND: "iand", Ops.XOR: "ixor", Ops.WHERE: "bcsel", Ops.MAX: "umax", Ops.SHL: "ishl", Ops.SHR: "ushr"}
|
||||
s_aop = {**u_aop, Ops.CMPLT: "ilt", Ops.CDIV: "idiv", Ops.CMOD: "irem", Ops.MAX: "imax", Ops.SHR: "ishr"}
|
||||
f_aop = { Ops.ADD: "fadd", Ops.MUL: "fmul", Ops.CMPLT: "flt", Ops.CMPNE: "fneu", Ops.CMPEQ: "feq", Ops.FDIV: "fdiv", Ops.RECIPROCAL: "frcp",
|
||||
Ops.MAX: "fmax", Ops.TRUNC: "ftrunc", Ops.SIN: "fsin", Ops.EXP2: "fexp2", Ops.LOG2: "flog2"}
|
||||
aop = {**{x:u_aop for x in (dtypes.bool,)+dtypes.uints}, **{x:s_aop for x in dtypes.sints}, **{x:f_aop for x in dtypes.floats}}
|
||||
|
||||
def c(t:DType, u:bool=True) -> str: return "u" if t in dtypes.uints and u else ("i" if t in dtypes.ints else ("f" if t in dtypes.floats else "b"))
|
||||
def ncast(b:mesa.nir_builder, src:mesa.nir_def, it:DType, ot:DType) -> mesa.nir_def:
|
||||
if isinstance(it, PtrDType) and ot == dtypes.long: return src
|
||||
return nalu(b, f"{c(it)}2{c(it) if it in dtypes.ints and ot in dtypes.ints else c(ot, ot == dtypes.bool)}{ot.bitsize}", src)
|
||||
|
||||
def nif(b:mesa.nir_builder, cond:mesa.nir_def, then_fn:Callable, else_fn:Callable):
|
||||
nif = mesa.nir_push_if(b, cond)
|
||||
t = then_fn()
|
||||
mesa.nir_push_else(b, nif)
|
||||
e = else_fn()
|
||||
mesa.nir_pop_if(b, nif)
|
||||
return t, e
|
||||
|
||||
def nalu(b:mesa.nir_builder, op:str, *srcs:mesa.nir_def) -> mesa.nir_def: return g(f"nir_build_alu{len(srcs)}")(b, g(f"nir_op_{op}"), *srcs).contents
|
||||
|
||||
def nir_instr(nc=1, bs=lambda: None, intrins=None, srcs=None, has_def=True, df=None, also=lambda: None, **contents):
|
||||
def dec(f:Callable):
|
||||
@functools.wraps(f)
|
||||
def wrapper(*args, **kwargs) -> mesa.nir_def:
|
||||
(ba:=inspect.signature(f).bind(*args, **kwargs)).apply_defaults()
|
||||
def go(g): return g(**{nm: ba.arguments[nm] for nm in inspect.signature(g).parameters}) if callable(g) else g
|
||||
|
||||
instr = f(*args, **kwargs)
|
||||
if has_def: mesa.nir_def_init(instr.contents.instr, instr.contents._def, go(nc), go(bs))
|
||||
for k, v in go(intrins or {}).items():
|
||||
idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic].index_map[g(f"NIR_INTRINSIC_{k}")]
|
||||
assert idx > 0, "invalid intrinsic. mesa version mismatch?"
|
||||
instr.contents.const_index[idx - 1] = go(v)
|
||||
for i, src in enumerate(go(srcs or [])): ctypes.cast(instr.contents.src, ctypes.POINTER(mesa.nir_src))[i] = go(src)
|
||||
for k,v in {k:vcomp for k,v in contents.items() if (vcomp:=go(v)) is not None}.items(): setattr(instr.contents, k, go(v))
|
||||
mesa.nir_builder_instr_insert(ba.arguments['b'], instr.contents.instr)
|
||||
go(also)
|
||||
return instr.contents._def if has_def else (mesa.nir_def() if df is None else go(df))
|
||||
return wrapper
|
||||
return dec
|
||||
|
||||
@nir_instr(nc=1, bs=lambda src: src.bit_size, exact=lambda b:b.exact, fp_fast_math=lambda b:b.fp_fast_math)
|
||||
def nchannel(b:mesa.nir_builder, src:mesa.nir_def, c:int):
|
||||
alu_src = mesa.nir_alu_src(src=nsrc(src))
|
||||
alu_src.swizzle[0] = c
|
||||
mov = mesa.nir_alu_instr_create(b.shader, mesa.nir_op_mov)
|
||||
ctypes.cast(mov.contents.src, ctypes.POINTER(mesa.nir_alu_src))[0] = alu_src
|
||||
return mov
|
||||
|
||||
def nimm_set(imm:mesa.nir_def, x, dtype:DType):
|
||||
instr = ctypes.cast(imm.parent_instr, ctypes.POINTER(mesa.nir_load_const_instr))
|
||||
struct.pack_into(unwrap(dtype.fmt), (ctypes.c_ubyte * dtype.itemsize).from_address(ctypes.addressof(instr.contents.value)), 0, truncate[dtype](x))
|
||||
|
||||
@nir_instr(nc=1, bs=lambda dtype: dtype.bitsize)
|
||||
def nimm(b:mesa.nir_builder, x, dtype:DType) -> mesa.nir_def:
|
||||
nimm_set((instr:=mesa.nir_load_const_instr_create(b.shader, 1, dtype.bitsize)).contents._def, x, dtype)
|
||||
return instr
|
||||
@nir_instr(nc=1, bs=lambda dtype: dtype.bitsize)
|
||||
def nundef(b, dtype): return mesa.nir_undef_instr_create(b.shader, 1, dtype.bitsize)
|
||||
|
||||
deref_var = nir_instr(nc=1, bs=32, modes=lambda var:var.data.mode, type=lambda var:var.type, var=lambda var:ctypes.pointer(var))( # pylint: disable=W0108
|
||||
lambda b, var: mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_var))
|
||||
|
||||
def iointr(space): return {"ALIGN_MUL":lambda dtype:dtype.itemsize} if space != AddrSpace.REG else {}
|
||||
def scope(space): return 'global' if space == AddrSpace.GLOBAL else ('shared' if space == AddrSpace.LOCAL else 'deref')
|
||||
nstore = nir_instr(has_def=False, df=lambda addr:addr, intrins=lambda space,val: {"WRITE_MASK":(1<<val.num_components)-1, **iointr(space)},
|
||||
num_components=lambda val:val.num_components, srcs=lambda space, addr, val: [nsrc(val), nsrc(addr)][::1 if space != AddrSpace.REG else -1])(
|
||||
lambda b, space, addr, val, dtype: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_store_{scope(space)}")))
|
||||
nload = nir_instr(nc=lambda dtype:dtype.count, bs=lambda dtype:dtype.bitsize//dtype.count, num_components=lambda dtype:dtype.count,
|
||||
intrins=lambda space:{**({"ACCESS":mesa.ACCESS_CAN_REORDER} if space==AddrSpace.GLOBAL else {}), **iointr(space)}, srcs=lambda addr: [nsrc(addr)])(
|
||||
lambda b, space, addr, dtype: mesa.nir_intrinsic_instr_create(b.shader, g(f"nir_intrinsic_load_{scope(space)}")))
|
||||
|
||||
ngid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_id))
|
||||
nlid = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_local_invocation_id))
|
||||
ngsz = nir_instr(nc=3, bs=32)(lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_workgroup_size))
|
||||
def nid(b): return nalu(b, "iadd", nalu(b, "imul", ngid(b), ngsz(b)), nlid(b))
|
||||
|
||||
nbarrier = nir_instr(has_def=False, intrins={"EXECUTION_SCOPE":mesa.SCOPE_WORKGROUP})(
|
||||
lambda b: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_barrier))
|
||||
|
||||
@nir_instr(has_def=False, target=lambda tgt:tgt and ctypes.pointer(tgt), condition=lambda cond:cond and nsrc(cond),
|
||||
else_target=lambda else_tgt: else_tgt and ctypes.pointer(else_tgt))
|
||||
def njump(b:mesa.nir_builder, typ, tgt=None, cond=None, else_tgt=None): return mesa.nir_jump_instr_create(b.shader, typ)
|
||||
|
||||
def if_phi(b:mesa.nir_builder, cond, then_fn, else_fn): return mesa.nir_if_phi(b, *nif(b, cond, then_fn, else_fn)).contents
|
||||
|
||||
def nidx(b:mesa.nir_builder, buf, off, dtype, gate=None) -> mesa.nir_def:
|
||||
@nir_instr(nc=1, bs=32, modes=lambda buf: buf.data.mode, type=lambda buf: mesa.glsl_get_array_element(buf.type))
|
||||
def reg(b, buf):
|
||||
deref = mesa.nir_deref_instr_create(b.shader, mesa.nir_deref_type_array)
|
||||
deref.contents.parent, deref.contents.arr.index = nsrc(deref_var(b, buf)), nsrc(off)
|
||||
return deref
|
||||
f = (functools.partial(reg, b, buf) if dtype.addrspace == AddrSpace.REG else
|
||||
lambda: nalu(b, "iadd", buf, nalu(b, "imul", off, nimm(b, dtype.itemsize, dtypes.long))))
|
||||
return if_phi(b, gate, f, lambda: buf) if gate is not None else f()
|
||||
|
||||
class NIRRenderer(Renderer):
|
||||
suffix = "NIR"
|
||||
nir_options: bytes
|
||||
global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
|
||||
code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
|
||||
|
||||
extra_matcher = PatternMatcher([
|
||||
# handle negative unsigned CONST
|
||||
(UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype, x.dtype.max+x.arg+1) if x.arg < 0 else None),
|
||||
# from ptx
|
||||
(UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
|
||||
# load/store bool -> uint8
|
||||
(UPat(Ops.LOAD, dtypes.bool, name="x"),
|
||||
lambda x: x.replace(dtype=dtypes.uint8, src=x.src[0:1]+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())+x.src[2:]).cast(dtypes.bool)),
|
||||
(UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
|
||||
lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
|
||||
# NIR requires shift amount to be 32 bit: https://docs.mesa3d.org/nir/alu.html#nir-alu-op-ishl
|
||||
(UPat((Ops.SHL, Ops.SHR), name="x"), lambda x: x.replace(src=(x.src[0], x.src[1].cast(dtypes.uint))) if x.src[1].dtype.bitsize != 32 else None),
|
||||
# OpConvertFToU is undefined if Result Type is not wide enough, cast through int32
|
||||
# ref: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpConvertFToU
|
||||
(UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
|
||||
# load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off")), name="x"), lambda x,buf,off: x.replace(
|
||||
src=(buf,off.cast(dtypes.long))) if buf.dtype.addrspace != AddrSpace.REG and off.op not in (Ops.CAST, Ops.STACK) else None),
|
||||
# images need index to be int for nir
|
||||
(UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x")),
|
||||
lambda buf,idx_y,idx_x: buf.index(idx_y.cast(dtypes.int), idx_x.cast(dtypes.int))),
|
||||
])
|
||||
|
||||
def_rewrite = PatternMatcher([
|
||||
(UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)),
|
||||
(UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 8)),
|
||||
(UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: ctx.param(ctx.b, x, 4)),
|
||||
(UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"),UPat.var("off"))).or_casted(), UPat.var("val"))),
|
||||
lambda ctx,buf,off,val: nstore(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), ctx.r[val], val.dtype)),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(), UPat.var("alt"), UPat.var("gate")), name="x"),
|
||||
lambda ctx,x,buf,off,alt,gate: if_phi(ctx.b, ctx.r[gate],
|
||||
lambda: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype, ctx.r[gate]), x.dtype), lambda: ctx.r[alt])),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))).or_casted(),), name="x"),
|
||||
lambda ctx,x,buf,off: nload(ctx.b, buf.addrspace, nidx(ctx.b, ctx.r[buf], ctx.r[off], buf.dtype), x.dtype)),
|
||||
(UPat(Ops.STACK, name="x"), lambda ctx,x: nalu(ctx.b, f"vec{x.dtype.count}", *[ctx.r[src] for src in x.src])),
|
||||
(UPat(GroupOp.ALU, name="x"), lambda ctx,x: nalu(ctx.b, aop[x.src[0].dtype.scalar()][x.op], *[ctx.r[src] for src in x.src])),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: ncast(ctx.b, ctx.r[x.src[0]], x.src[0].dtype, x.dtype)),
|
||||
(UPat(Ops.BITCAST, src=(UPat.var("a"),), allow_any_len=True), lambda ctx,a: ctx.r[a]),
|
||||
(UPat(Ops.GEP, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: nchannel(ctx.b, ctx.r[a], x.arg[0])),
|
||||
(UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:mesa.nir_local_variable_create(ctx.b.impl, glsl_type(x.dtype), f"acc{x.arg}".encode()).contents),
|
||||
(UPat(Ops.BARRIER), lambda ctx: nbarrier(ctx.b)),
|
||||
(UPat(Ops.IF, name="x"), lambda ctx,x: mesa.nir_push_if(ctx.b, ctx.r[x.src[0]])),
|
||||
(UPat(Ops.ENDIF, name="x"), lambda ctx,x: (lambda _: mesa.nir_def())(mesa.nir_pop_if(ctx.b, ctx.r[x.src[0]])))
|
||||
])
|
||||
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
self.compiler = fromimport("tinygrad.runtime.support.compiler_mesa", self.__class__.__name__.replace("Renderer", "Compiler"))(target.arch)
|
||||
if hasattr(self.compiler, "nir_options"): self.nir_options = self.compiler.nir_options
|
||||
mesa.glsl_type_singleton_init_or_ref()
|
||||
self._deinit_types = True
|
||||
|
||||
def __del__(self):
|
||||
if getattr(self, "_deinit_types", False): mesa.glsl_type_singleton_decref()
|
||||
|
||||
def param(self, b:mesa.nir_builder, x, sz:int) -> mesa.nir_def: raise NotImplementedError("needs param")
|
||||
def prerender(self, uops:list[UOp]):
|
||||
self.b = mesa.nir_builder_init_simple_shader(mesa.MESA_SHADER_COMPUTE, mesa.nir_shader_compiler_options.from_buffer_copy(self.nir_options), None)
|
||||
self.b.shader.contents.info.workgroup_size_variable = any([u.op == Ops.SPECIAL and u.arg[0] == 'i' for u in uops])
|
||||
def postrender(self, uops:list[UOp]): pass
|
||||
|
||||
def render(self, uops:list[UOp]):
|
||||
self.prerender(uops)
|
||||
for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].arg
|
||||
self.r: dict[UOp, Any] = {}
|
||||
self.param_idx, ranges = 0, []
|
||||
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP, Ops.INDEX}: pass
|
||||
elif u.op is Ops.CAST and isinstance(u.dtype, PtrDType): pass
|
||||
elif u.op is Ops.AFTER:
|
||||
self.r[u] = self.r[u.src[0]]
|
||||
elif u.op == Ops.SINK:
|
||||
if u.arg is not None:
|
||||
self.b.shader.contents.info.name = ctypes.cast(ctypes.create_string_buffer(u.arg.function_name.encode()), POINTER[ctypes.c_char])
|
||||
elif u.op == Ops.DEFINE_LOCAL:
|
||||
self.r[u] = nimm(self.b, self.b.shader.contents.info.shared_size, dtypes.long)
|
||||
self.b.shader.contents.info.shared_size += u.dtype.nbytes()
|
||||
elif u.op == Ops.RANGE:
|
||||
ranges.append(i:=deref_var(self.b, mesa.nir_local_variable_create(self.b.impl, glsl_type(u.dtype), f"idx{range_str(u)}".encode()).contents))
|
||||
nstore(self.b, AddrSpace.REG, i, nimm(self.b, 0, u.dtype), u.dtype)
|
||||
mesa.nir_push_loop(self.b)
|
||||
self.r[u] = nload(self.b, AddrSpace.REG, i, u.dtype)
|
||||
nif(self.b, nalu(self.b, "ilt", self.r[u], self.r[u.src[0]]), lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
|
||||
elif u.op == Ops.END:
|
||||
r = u.src[1]
|
||||
next_i = nalu(self.b, "iadd", self.r[r], nimm(self.b, 1, r.dtype))
|
||||
# TODO: this nif should be removable ... but TestMultiTensor.test_double_matmul_shard_W_0 segfaults with it gone
|
||||
nif(self.b, nalu(self.b, "ilt", next_i, self.r[r.src[0]]), lambda: None, lambda: njump(self.b, mesa.nir_jump_break))
|
||||
nstore(self.b, AddrSpace.REG, ranges.pop(), next_i, r.dtype),
|
||||
mesa.nir_pop_loop(self.b, None)
|
||||
else:
|
||||
if (d:=self.def_rewrite.rewrite(u, ctx=self)) is None: raise RuntimeError(f"failed to render {u.op} srcs {[x.dtype for x in u.src]}")
|
||||
self.r[u] = cast(mesa.nir_def, d)
|
||||
self.postrender(uops)
|
||||
|
||||
mesa.nir_validate_shader(self.b.shader, b"after render")
|
||||
if DEBUG >= 4: mesa.nir_print_shader(self.b.shader, ctypes.POINTER(mesa.struct__IO_FILE).in_dll(ctypes.CDLL(ctypes.util.find_library('c')),
|
||||
"__stdoutp" if OSX else "stdout"))
|
||||
mesa.nir_serialize(blob:=mesa.struct_blob(), self.b.shader, False)
|
||||
ret = base64.b64encode(ctypes.string_at(blob.data, blob.size)).decode()
|
||||
|
||||
mesa.ralloc_free(self.b.shader)
|
||||
ctypes.CDLL(None).free(blob.data)
|
||||
del self.b, self.r
|
||||
|
||||
return ret
|
||||
|
||||
def supported_dtypes(self): return {d for d in Renderer.supported_dtypes(self) if d not in dtypes.fp8s+(dtypes.bfloat16,)}
|
||||
|
||||
class NAKRenderer(NIRRenderer):
|
||||
param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz),
|
||||
intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])(
|
||||
lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.half or int(self.target.arch[3:]) >= 53)}
|
||||
|
||||
class LVPRenderer(NIRRenderer):
|
||||
has_local = False
|
||||
has_shared = False
|
||||
global_max = (1, 0, 0)
|
||||
nir_options = mesa.lvp_nir_options
|
||||
# gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
|
||||
# https://gitlab.freedesktop.org/mesa/mesa/-/blob/c200b18e876468b51fe80d9660f612dc03a5138e/src/gallium/auxiliary/gallivm/lp_bld_arit.c#L2972
|
||||
code_for_op = {k:v for k,v in NIRRenderer.code_for_op.items() if k != Ops.EXP2}
|
||||
|
||||
param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz},
|
||||
srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
|
||||
setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
|
||||
|
||||
def prerender(self, uops:list[UOp]):
|
||||
super().prerender(uops)
|
||||
self.param_sz = sum([8 if u.op == Ops.PARAM else u.dtype.itemsize for u in uops if u.op in (Ops.PARAM, Ops.DEFINE_VAR)])
|
||||
|
||||
def tovec(b, idx_y, idx_x): return nalu(b, "vec4", idx_x, idx_y, nundef(b, dtypes.int), nundef(b, dtypes.int))
|
||||
def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32
|
||||
nstore_img = nir_instr(has_def=False, df=lambda img:img, num_components=lambda val:val.num_components,
|
||||
intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'SRC_TYPE':nfloat(dtype)},
|
||||
srcs=lambda b,img,idx_y,idx_x,val:[nsrc(x) for x in [img, tovec(b, idx_y, idx_x), nundef(b, dtypes.int), val, nimm(b, 0, dtypes.int)]])(
|
||||
lambda b,img,idx_y,idx_x,val,dtype:mesa.nir_intrinsic_instr_create(b.shader,g("nir_intrinsic_image_store")))
|
||||
|
||||
_nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2D, 'ACCESS':mesa.ACCESS_CAN_REORDER, 'DEST_TYPE':nfloat(dtype)},
|
||||
nc=4, bs=32, num_components=4,
|
||||
srcs=lambda b,img,idx_y,idx_x:[nsrc(x) for x in [img, tovec(b, idx_y, idx_x), nundef(b, dtypes.int), nimm(b, 0, dtypes.int)]])(
|
||||
lambda b,img,idx_y,idx_x,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
|
||||
|
||||
class IR3Renderer(NIRRenderer, OpenCLRenderer):
|
||||
has_aux = True
|
||||
|
||||
def nload_img(ctx,img,idx_y,idx_x):
|
||||
ctx.texs.add(img)
|
||||
return _nload_img(ctx.b, ctx.r[img], ctx.r[idx_y], ctx.r[idx_x], img.dtype)
|
||||
|
||||
def_rewrite = PatternMatcher([
|
||||
(UPat(Ops.STORE, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("val")), allow_any_len=True),
|
||||
lambda ctx,img,idx_y,idx_x,val: nstore_img(ctx.b, ctx.r[img], ctx.r[idx_y], ctx.r[idx_x], ctx.r[val], val.dtype)),
|
||||
(UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("alt"), UPat.var("gate"))),
|
||||
lambda ctx,img,idx_y,idx_x,alt,gate: if_phi(ctx.b, ctx.r[gate], lambda: ctx.nload_img(img, idx_y, idx_x), lambda: ctx.r[alt])),
|
||||
(UPat(Ops.LOAD, src=(UPat.var('img').index(UPat.var('idx_y'), UPat.var('idx_x')),)), nload_img),
|
||||
]) + NIRRenderer.def_rewrite
|
||||
|
||||
_param = LVPRenderer.param
|
||||
def _param_img(self, x):
|
||||
self.img_idx += 1
|
||||
return nimm(self.b, self.img_idx - 1, dtypes.int)
|
||||
|
||||
def param(self, b, x, sz): return self._param_img(x) if isinstance(x.dtype, ImageDType) else self._param(b, x, sz)
|
||||
|
||||
def prerender(self, uops:list[UOp]):
|
||||
super().prerender(uops)
|
||||
self.texs:set[UOp] = set()
|
||||
self.uops, self.ibo_idx, self.img_idx = uops, 0, 0
|
||||
self.param_sz = sum([8 if u.op == Ops.PARAM else u.dtype.itemsize for u in uops if u.op in (Ops.PARAM, Ops.DEFINE_VAR)])
|
||||
|
||||
def postrender(self, uops:list[UOp]):
|
||||
bufs, texs, imgs = [u for u in uops if u.op == Ops.PARAM], itertools.count().__next__, itertools.count().__next__
|
||||
for b in filter(lambda b: isinstance(b.dtype, ImageDType), bufs): nimm_set(self.r[b], texs() if b in self.texs else imgs(), dtypes.int)
|
||||
|
||||
self.b.shader.contents.info.num_ubos = len([u for u in bufs if not isinstance(u.dtype, ImageDType)])
|
||||
self.b.shader.contents.info.num_images = texs() + imgs()
|
||||
|
||||
def supported_dtypes(self): return {d for d in NIRRenderer.supported_dtypes(self) if d != dtypes.double}
|
||||
242
tinygrad_repo/tinygrad/renderer/ptx.py
Normal file
242
tinygrad_repo/tinygrad/renderer/ptx.py
Normal file
@@ -0,0 +1,242 @@
|
||||
from typing import cast, Callable
|
||||
import struct
|
||||
from collections import defaultdict
|
||||
from tinygrad.codegen.opt import tc
|
||||
from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
|
||||
from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.renderer.cstyle import CUDARenderer
|
||||
from tinygrad.helpers import flatten, get_single_element, prod, unwrap, Target
|
||||
|
||||
def render_val(x, dtype):
|
||||
if dtypes.is_float(dtype):
|
||||
if dtype == dtypes.double: return "0d%02X%02X%02X%02X%02X%02X%02X%02X" % tuple(struct.pack("d",x)[::-1])
|
||||
if dtype == dtypes.half: return "0x%02X%02X" % tuple(struct.pack("e",x)[::-1])
|
||||
return "0f%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
|
||||
return str(int(x)) + ("U" if dtypes.is_unsigned(dtype) else "")
|
||||
|
||||
asm_for_op: dict[Ops, Callable] = {
|
||||
Ops.RECIPROCAL: lambda d,a,dt,name: f"rcp{'.approx' if dtypes.is_float(dt) else ''}.{name} {d}, {a};",
|
||||
Ops.EXP2: lambda d,a,dt,name: f"ex2.approx.{name} {d}, {a};", Ops.LOG2: lambda d,a,dt,name: f"lg2.approx.{name} {d}, {a};",
|
||||
Ops.SIN: lambda d,a,dt,name: f"sin.approx.{name} {d}, {a};", Ops.SQRT: lambda d,a,dt,name: f"sqrt.approx.{name} {d}, {a};",
|
||||
Ops.TRUNC: lambda d,a,dt,name: f"cvt.rzi.{name}.{name} {d}, {a};",
|
||||
Ops.SHR: lambda d,a,b,dt,name: f"shr.{name} {d}, {a}, {b};", Ops.SHL: lambda d,a,b,dt,name: f"shl.b{name[1:]} {d}, {a}, {b};",
|
||||
Ops.ADD: lambda d,a,b,dt,name: f"{'or' if dt == dtypes.bool else 'add'}.{name} {d}, {a}, {b};",
|
||||
Ops.MUL: lambda d,a,b,dt,name: f"{'and' if dt == dtypes.bool else 'mul'}{'.lo' if dtypes.is_int(dt) else ''}.{name} {d}, {a}, {b};",
|
||||
Ops.XOR: lambda d,a,b,dt,name: f"xor.pred {d}, {a}, {b};" if dt == dtypes.bool else f"xor.b{name[1:]} {d}, {a}, {b};",
|
||||
Ops.AND: lambda d,a,b,dt, name: f"and.pred {d}, {a}, {b};" if dt == dtypes.bool else f"and.b{name[1:]} {d}, {a}, {b};",
|
||||
Ops.OR: lambda d,a,b,dt, name: f"or.pred {d}, {a}, {b};" if dt == dtypes.bool else f"or.b{name[1:]} {d}, {a}, {b};",
|
||||
Ops.CDIV: lambda d,a,b,dt,name: f"div.{name} {d}, {a}, {b};", Ops.CMOD: lambda d,a,b,dt,name: f"rem.{name} {d}, {a}, {b};",
|
||||
Ops.MAX: lambda d,a,b,dt,name: f"max.{name} {d}, {a}, {b};", Ops.CMPEQ: lambda d,a,b,dt,name: f"setp.eq.{name} {d}, {a}, {b};",
|
||||
Ops.CMPLT: lambda d,a,b,dt,name: f"setp.lt.{name} {d}, {a}, {b};",
|
||||
Ops.CMPNE: lambda d,a,b,dt,name: f"setp.{'neu' if dtypes.is_float(dt) else 'ne'}.{name} {d}, {a}, {b};",
|
||||
Ops.MULACC: lambda d,a,b,c,dt,name: f"{'fma.rn' if dtypes.is_float(dt) else 'mad.lo'}.{name} {d}, {a}, {b}, {c};",
|
||||
Ops.WHERE: lambda d,a,b,c,dt,name: [f"@{a} mov.{name} {d}, {b};", f"@!{a} mov.{name} {d}, {c};"] if dt == dtypes.bool else \
|
||||
f"selp.{'b16' if name == 'f16' else name} {d}, {b}, {c}, {a};"
|
||||
}
|
||||
|
||||
supports_half = (Ops.EXP2, Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPLT, Ops.WHERE, Ops.TRUNC)
|
||||
doesnt_support_half: tuple[Ops, ...] = tuple(op for op in asm_for_op.keys() if op not in supports_half)
|
||||
ptx_matcher = PatternMatcher([
|
||||
# bool CMPNE is XOR, bool CMPLT is XOR+AND (universal makes this slow, this is for renderer only)
|
||||
(UPat.var('x', dtype=dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y),
|
||||
(UPat.var('x', dtype=dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True),
|
||||
(UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
|
||||
# upcast to float32 all the ops that don't support half
|
||||
(UPat(doesnt_support_half, dtype=dtypes.half, name="x"),
|
||||
lambda x: (UOp(x.op, dtypes.float32, tuple(vv.cast(dtypes.float32) for vv in x.src), x.arg).cast(dtypes.half))),
|
||||
# load/store bool -> uint8
|
||||
(UPat(Ops.LOAD, dtypes.bool, src=(UPat(dtype=dtypes.int64),), name="x", allow_any_len=True),
|
||||
lambda x: UOp(x.op, dtypes.uint8, x.src[0:1] + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ()) + x.src[2:]).cast(dtypes.bool)),
|
||||
(UPat(Ops.STORE, src=(UPat(dtype=dtypes.int64), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
|
||||
lambda x: UOp(x.op, dtypes.void, (x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:])),
|
||||
# indexing on PTX is in uint64, we do the math while it's still in the graph
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx")), name="op"), lambda buf,idx,op:
|
||||
UOp(Ops.INDEX, dtype=dtypes.int64, src=(buf, buf.cast(dtypes.int64)+idx.cast(dtypes.int64)*buf.dtype.itemsize)+op.src[2:]) \
|
||||
if op.dtype != dtypes.int64 and buf.dtype.addrspace != AddrSpace.REG else None),
|
||||
# ptx shr and shl instructions require y to be uint
|
||||
(UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
|
||||
(UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
|
||||
])
|
||||
|
||||
def mem_type(x:UOp) -> str:
|
||||
match x.op:
|
||||
case Ops.AFTER: return mem_type(x.src[0])
|
||||
case Ops.DEFINE_LOCAL: return 'shared'
|
||||
case Ops.PARAM: return 'shared' if x.addrspace == AddrSpace.LOCAL else 'global'
|
||||
case _: raise RuntimeError(f"{x.op} needs to be memory")
|
||||
|
||||
def render_wmma(ctx: "PTXRenderer", wmma: UOp):
|
||||
assert ctx.wmma_r, "registry values for wmma must be populated"
|
||||
(N, M, K), dtype_in, dtype_out = wmma.arg[1], wmma.arg[2], wmma.arg[3]
|
||||
|
||||
for src, regs in zip(wmma.src, ctx.wmma_r):
|
||||
for i, reg in enumerate(regs): # pack input and acc registers
|
||||
if (elems_per_reg := 4 // src.dtype.scalar().itemsize) == 1: yield f"mov.b32 {reg}, {ctx.r[src][i]};"
|
||||
else: yield f"mov.b32 {reg}, {{{', '.join(ctx.r[src][i * elems_per_reg : (i+1) * elems_per_reg])}}};"
|
||||
|
||||
dt_map_in, dt_map_out = {dtypes.float: "tf32", dtypes.half: "f16"}, {dtypes.float: "f32", dtypes.half: "f16"}
|
||||
yield f'mma.sync.aligned.m{M}n{N}k{K}.row.col.{dt_map_out[dtype_out]}.{dt_map_in[dtype_in]}.{dt_map_in[dtype_in]}.{dt_map_out[dtype_out]}{" "*12}'+\
|
||||
f'{{{", ".join(ctx.wmma_r[2])}}}, {{{", ".join(ctx.wmma_r[0])}}}, {{{", ".join(ctx.wmma_r[1])}}}, {{{", ".join(ctx.wmma_r[2])}}};'
|
||||
|
||||
for i, reg in enumerate(ctx.wmma_r[2]): # unpack acc registers
|
||||
if (elems_per_reg := 4 // dtype_out.itemsize) == 1: yield f"mov.b32 {ctx.r[wmma][i]}, {reg};"
|
||||
else: yield f"mov.b32 {{{', '.join(ctx.r[wmma][i * elems_per_reg : (i+1) * elems_per_reg])}}}, {reg};"
|
||||
|
||||
def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.is_float(b) else '.rn' if dtypes.is_float(a) and \
|
||||
(a.itemsize < b.itemsize or dtypes.is_int(b) or b == dtypes.bool) else ''
|
||||
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"),
|
||||
(UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"),
|
||||
(UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"),
|
||||
(UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong]} {ctx.r[x]}, [data{x.arg.slot}+0];"),
|
||||
(UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), name="x", allow_any_len=True, src=(UPat.var("src0"),)),
|
||||
lambda ctx, x, src0: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], src0.dtype, ctx.types[src0.dtype])),
|
||||
(UPat(GroupOp.ALU, name="x"), lambda ctx, x: ctx.code_for_op[x.op](ctx.r[x], *[ctx.r[v] for v in x.src], x.dtype, ctx.types[x.dtype])),
|
||||
(UPat(Ops.BITCAST, name="x", src=(UPat.var("a"),), allow_any_len=True), lambda ctx, x, a: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {ctx.r[a]};"),
|
||||
(UPat(Ops.CAST, name="x", src=(UPat(dtype=dtypes.bool, name="a"),)),
|
||||
lambda ctx, x, a: f"selp.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(1, x.dtype)}, {render_val(0, x.dtype)}, {ctx.r[a]};"),
|
||||
(UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
|
||||
lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
|
||||
# store / gated load / load
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(), UPat.var("var"))),
|
||||
lambda ctx, loc, var, buf: f"st.{mem_type(buf)}" + \
|
||||
f"{f'.v{cnt}' if ((cnt:=var.dtype.count)>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
|
||||
f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.dtype.count > 1 else ctx.r[var]};"),
|
||||
(UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(), UPat.var("alt"), UPat.var("gate"))),
|
||||
lambda ctx, x, loc, alt, gate, buf: flatten([
|
||||
[f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
|
||||
[f"@{ctx.r[gate]} ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
|
||||
]) if alt.dtype.count > 1 else [
|
||||
f"@{ctx.r[gate]} ld.{mem_type(buf)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
|
||||
f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
|
||||
(UPat(Ops.LOAD, name="x", src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("loc"))).or_casted(),)),
|
||||
lambda ctx, x, loc, buf: f"ld.{mem_type(buf)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
|
||||
if x.dtype.count > 1 else f"ld.{mem_type(buf)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
|
||||
# simple
|
||||
(UPat(Ops.DEFINE_REG, src=()), lambda ctx: []),
|
||||
(UPat(Ops.RANGE, name="r"), lambda ctx, r: [
|
||||
f"mov.u32 {ctx.r[r]}, -1;",
|
||||
f"bra END_{ctx.r[r][1:]};",
|
||||
"LOOP_" + f"{ctx.r[r][1:]}:"]),
|
||||
(UPat(Ops.END, name="x", src=(UPat(), UPat(Ops.RANGE, name="r"))), lambda ctx, x, r: [
|
||||
"END_" + f"{ctx.r[r][1:]}:",
|
||||
ctx.code_for_op[Ops.ADD](ctx.r[r], ctx.r[r], "1", dtypes.int, ctx.types[dtypes.int]),
|
||||
ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[r], ctx.r[r.src[0]], dtypes.int, ctx.types[dtypes.int]),
|
||||
f"@{ctx.r[x]} bra LOOP_{ctx.r[r][1:]};"]),
|
||||
(UPat(Ops.DEFINE_LOCAL, name="x"),
|
||||
lambda ctx, x: [f".shared .align 16 .b8 local{x.arg}[{x.dtype.size*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, local{x.arg}[0];"]),
|
||||
(UPat(Ops.IF, name="x"), lambda ctx, x: f"@!{ctx.r[x.src[0]]} bra IF_{ctx.r[x.src[0]][1:]}_{ctx.uops.index(x)};"),
|
||||
(UPat(Ops.ENDIF, name="x"), lambda ctx, x: f"IF_{ctx.r[x.src[0].src[0]][1:]}_{ctx.uops.index(x.src[0])}:"),
|
||||
(UPat(Ops.WMMA, name="x"), lambda ctx, x: list(render_wmma(ctx, x))),
|
||||
(UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
|
||||
(UPat(Ops.DEFINE_VAR, name="x"), lambda ctx, x: f"ld.param.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{x.expr}+0];"),
|
||||
])
|
||||
|
||||
class PTXRenderer(Renderer):
|
||||
suffix = "PTX"
|
||||
global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
|
||||
tc_sm80 = [x for x in tc.cuda_sm80 if x.dtype_in in [dtypes.half, dtypes.float]]
|
||||
code_for_op = asm_for_op
|
||||
extra_matcher = ptx_matcher
|
||||
def __init__(self, target:Target):
|
||||
super().__init__(target)
|
||||
from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler, PTXCompiler
|
||||
self.compiler = (PTXCompiler if target.interface.startswith("MOCK") or target.device == "CUDA" else NVPTXCompiler)(target.arch)
|
||||
self.tensor_cores = PTXRenderer.tc_sm80 if (ver:=int(target.arch[3:])) >= 80 else tc.cuda_sm75 if ver >= 75 else []
|
||||
|
||||
# language options
|
||||
kernel_prefix = """.version VERSION
|
||||
.target TARGET
|
||||
.address_size 64
|
||||
.visible .entry"""
|
||||
barrier = "bar.sync\t0;"
|
||||
types: dict[DType, str] = { dtypes.int8: "s16", dtypes.int16: "s16", dtypes.int32: "s32", dtypes.int64: "s64",
|
||||
dtypes.uint8: "u16", dtypes.uint16: "u16", dtypes.uint32: "u32", dtypes.uint64: "u64",
|
||||
dtypes.float16: "f16", dtypes.float32: "f32", dtypes.float64: "f64", dtypes.bool: "pred" }
|
||||
|
||||
mem_types: dict[DType, str] = {**types, dtypes.int8: "s8", dtypes.uint8: "u8", dtypes.bool: "u8", dtypes.float16: "b16"}
|
||||
cast_types: dict[DType, str] = {**types, dtypes.int8: "s8", dtypes.uint8: "u8"}
|
||||
|
||||
def render_kernel(self, kernel, function_name, bufs, regs, uops) -> str:
|
||||
def fmt(line): return line if line[0]=="$" else "\t" + line.replace(" ", "\t" if len(line.split(" ")[0]) > 7 else "\t\t", 1)
|
||||
kernel = '\n'.join(map(fmt, [f".reg .{reg.split('_')[-2]} %{reg}<{cnt}>;" for reg,cnt in regs] + kernel + ["ret;"]))
|
||||
local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
|
||||
launch_bounds = prod([d.vmax for d in local_dims])
|
||||
params = ',\n\t'.join([f".param .{'u64' if dtype.__class__ == PtrDType else self.types[dtype]} {name}" for name,dtype in bufs])
|
||||
return f"{self.kernel_prefix.format(launch_bounds=launch_bounds)} {function_name} (\n\t{params}\n)\n.maxntid {launch_bounds}\n{{\n{kernel}\n}}"
|
||||
|
||||
def render(self, uops:list[UOp]) -> str:
|
||||
kernel:list[str] = []
|
||||
bufs = []
|
||||
|
||||
c: defaultdict[str, int] = defaultdict(int)
|
||||
r: dict[UOp, list[str]|str] = {}
|
||||
self.r = r
|
||||
self.uops = uops
|
||||
|
||||
def ssa(prefix:str, u:UOp|None=None, dtype:str|None=None) -> str:
|
||||
nonlocal c
|
||||
prefix += f"_{dtype if dtype is not None else self.types[unwrap(u).dtype.base]}_"
|
||||
c[prefix] += 1
|
||||
return f"%{prefix}{c[prefix]-1}"
|
||||
|
||||
name = "test"
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP}: continue
|
||||
if u.op is Ops.CAST and isinstance(u.dtype, PtrDType): continue
|
||||
if u.op is Ops.AFTER:
|
||||
self.r[u] = self.r[u.src[0]]
|
||||
continue
|
||||
if u.op is Ops.SINK:
|
||||
if u.arg is not None: name = u.arg.function_name
|
||||
continue
|
||||
if u.op is Ops.STACK:
|
||||
r[u] = [cast(str,r[x]) for x in u.src]
|
||||
continue
|
||||
if u.op is Ops.GEP:
|
||||
r[u] = r[u.src[0]][get_single_element(u.arg)]
|
||||
continue
|
||||
if u.op in {Ops.CAST, Ops.BITCAST} and (u.src[0].dtype == u.dtype or isinstance(u.src[0].dtype, PtrDType)):
|
||||
r[u] = r[u.src[0]]
|
||||
continue
|
||||
if u.op is Ops.DEFINE_REG:
|
||||
r[u] = [ssa("reg", u, self.types[u.dtype.base.scalar()]) for _ in range(u.max_numel())]
|
||||
continue
|
||||
if u.op in {Ops.INDEX, Ops.LOAD, Ops.STORE} and isinstance(u.src[0].dtype, PtrDType) and u.src[0].dtype.addrspace == AddrSpace.REG:
|
||||
if u.op is Ops.INDEX:
|
||||
assert u.src[1].op == Ops.CONST, f"index on REG in ptx only supported on CONST, not {u.src[1].op}"
|
||||
r[u] = r[u.src[0]][u.src[1].arg]
|
||||
else:
|
||||
r[u] = r[u.src[0]]
|
||||
if u.op is Ops.STORE:
|
||||
typ = "pred" if u.src[1].dtype == dtypes.bool else ("b"+self.types[u.src[1].dtype][1:])
|
||||
kernel.append(f"mov.{typ} {self.r[u.src[0]]}, {self.r[u.src[1]]};")
|
||||
continue
|
||||
if u.op is Ops.INDEX: continue # other index we can skip
|
||||
if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
|
||||
elif u.op is Ops.DEFINE_VAR: bufs.append((u.expr, u.dtype))
|
||||
elif u.op is Ops.LOAD:
|
||||
r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)] if u.dtype.count > 1 else ssa('val', u)
|
||||
elif u.op is Ops.PARAM: bufs.append((f"data{u.arg.slot}", u.dtype))
|
||||
elif u.op is Ops.WMMA:
|
||||
# registers for packing/unpacking input and acc
|
||||
self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)],
|
||||
[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.scalar().itemsize)],
|
||||
[ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]]
|
||||
r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)]
|
||||
prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
|
||||
Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL: ("local", self.types[dtypes.ulong]),
|
||||
Ops.PARAM: ("dat", self.types[dtypes.ulong]), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))
|
||||
if prefix: r[u] = ssa(prefix, u, dtype)
|
||||
|
||||
if (l:=cast(str|list[str], string_rewrite.rewrite(u, ctx=self))) is None:
|
||||
raise RuntimeError(f"failed to render {u.op} with {u.dtype} srcs {[x.dtype for x in u.src]}")
|
||||
kernel.extend([l] if isinstance(l, str) else l)
|
||||
|
||||
if u.op is Ops.SPECIAL: kernel = [f".reg .u32 %{u.arg};"] + kernel
|
||||
return self.render_kernel(kernel, name, bufs, c.items(), uops)
|
||||
|
||||
def supported_dtypes(self): return {d for d in super().supported_dtypes()
|
||||
if (d != dtypes.half or int(self.target.arch[3:]) >= 53) and d not in dtypes.fp8s+(dtypes.bfloat16,)}
|
||||
117
tinygrad_repo/tinygrad/renderer/wgsl.py
Normal file
117
tinygrad_repo/tinygrad/renderer/wgsl.py
Normal file
@@ -0,0 +1,117 @@
|
||||
from tinygrad.dtype import DType, PtrDType, dtypes, truncate, AddrSpace
|
||||
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
|
||||
from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite, extra_pm
|
||||
from tinygrad.helpers import strip_parens
|
||||
|
||||
def _mask(dt:DType): return 0xFF if dt.itemsize == 1 else 0xFFFF
|
||||
|
||||
def sign_extend(val:UOp, sext_am:int):
|
||||
return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(dtypes.uint32, 0xffffffff) << sext_am, UOp.const(dtypes.uint32, 0)) \
|
||||
| val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
|
||||
|
||||
# store for char: buf[idx/4] <- (var << (idx%4)*8))
|
||||
def packed_store(bidx:UOp, var:UOp, gate:UOp|None=None):
|
||||
elems, mask = 4//var.dtype.itemsize, _mask(var.dtype)
|
||||
shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*var.dtype.itemsize), bidx.src[1] // elems
|
||||
new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32)
|
||||
idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx))
|
||||
buf = UOp.load(idx, *((UOp.const(dtypes.uint32, 0), gate) if gate is not None else ()), dtype=dtypes.uint32)
|
||||
return UOp.store(idx, (buf & wmask) | new_v, *((gate,) if gate is not None else ()))
|
||||
|
||||
# load for char: sign_extend(buf[idx/4] >> ((idx%4)*8))
|
||||
def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None, gate:UOp|None=None):
|
||||
elems, mask = 4//dtype.itemsize, _mask(dtype)
|
||||
shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*dtype.itemsize), bidx.src[1] // elems
|
||||
idx = UOp(Ops.INDEX, bidx.dtype, (bidx.src[0], div_idx))
|
||||
load = UOp.load(idx, *((var, gate) if var is not None and gate is not None else root.src[1:]), dtype=dtypes.uint32, arg=root.arg)
|
||||
val = (load.cast(dtypes.uint32) >> shift_am) & mask
|
||||
return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)
|
||||
|
||||
def is_packed(dt:DType, odt:DType|None = None) -> bool:
|
||||
if odt is None: odt = dt
|
||||
return dt.itemsize < 4 and dt.base != dtypes.half and (not isinstance(odt, PtrDType) or odt.addrspace != AddrSpace.REG)
|
||||
def _packed_size(dt:PtrDType): return dt.size // (4//dt.itemsize) if is_packed(dt) else dt.size
|
||||
|
||||
def is_nan(a):
|
||||
bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype)
|
||||
return (a.bitcast(getattr(dtypes, f"uint{bs}")) & ((1 << (bs - 1)) - 1)) > (((1 << exp) - 1) << mant)
|
||||
|
||||
wgsl_matcher = PatternMatcher([
|
||||
(UPat((Ops.CMPLT, Ops.XOR), src=(UPat(name="a", dtype=dtypes.bool), UPat.var("b")), name="c"),
|
||||
lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
|
||||
# TODO: load alt value doesnt have to be a const
|
||||
(UPat.load(UPat.var("b"), UPat.cvar("c"), UPat.var("gate"), name="l"),
|
||||
lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l.dtype, b.dtype) else None),
|
||||
(UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None),
|
||||
(UPat.store(UPat.var("bidx"), UPat.var("var"), UPat.var("gate")),
|
||||
lambda bidx,var,gate: packed_store(bidx,var,gate) if is_packed(var.dtype, bidx.dtype) else None),
|
||||
(UPat.store(UPat.var("bidx"), UPat.var("var")),
|
||||
lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None),
|
||||
(UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
|
||||
(UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
|
||||
# fix nan check: 'a != a -> is_nan()'
|
||||
(UPat.var("a") != UPat.var("a"), is_nan),
|
||||
]) + extra_pm
|
||||
|
||||
class WGSLRenderer(CStyleLanguage):
|
||||
global_max = (65535, 65535, 65535)
|
||||
local_max = (256, 256, 64)
|
||||
code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[int(x)]})", "l": lambda x: f"i32(lindex.{'xyz'[int(x)]})"}
|
||||
extra_matcher = wgsl_matcher
|
||||
supports_float4 = False
|
||||
barrier = "workgroupBarrier();"
|
||||
code_for_op = {**CStyleLanguage.code_for_op, Ops.WHERE: lambda a,b,c,dtype: f"select({c},{b},{a})"}
|
||||
nan = "nan()"
|
||||
type_map = { dtypes.float: "f32", dtypes.uchar: "u32", dtypes.ushort: "u32", dtypes.short: "i32",
|
||||
dtypes.char: "i32", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "bool", dtypes.half: "f16" }
|
||||
|
||||
string_rewrite = PatternMatcher([
|
||||
(UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"),
|
||||
(UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.arg else "false"),
|
||||
(UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
|
||||
lambda x: f"bitcast<u32>({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"),
|
||||
(UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"var<workgroup> {ctx[x]}: array<{ctx.buf_map(x.dtype.base)},{_packed_size(x.dtype)}>;"),
|
||||
(UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{_packed_size(x.dtype)}>;"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
|
||||
lambda ctx,x: f"bitcast<vec2<f16>>({ctx[x.src[0]]})[0]"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.uchar, name="x"), lambda ctx,x: f"bitcast<u32>({ctx[x.src[0]]}&0xFF)"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.char, name="x"), lambda ctx,x: f"((i32({ctx[x.src[0]]}&0xFF)<<24)>>24)"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.ushort, name="x"), lambda ctx,x: f"bitcast<u32>(vec2<f16>({ctx[x.src[0]]},0))" \
|
||||
if x.src[0].dtype == dtypes.half else f"bitcast<u32>({ctx[x.src[0]]}&0xFFFF)"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.short, name="x"), lambda ctx,x: f"bitcast<i32>(vec2<f16>({ctx[x.src[0]]},0))" \
|
||||
if x.src[0].dtype == dtypes.half else f"((i32({ctx[x.src[0]]}&0xFFFF)<<16)>>16)"),
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"),
|
||||
# TODO: load alt value doesnt have to be a const
|
||||
(UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate")),
|
||||
lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[gate]})"),
|
||||
(UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)),
|
||||
(UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\
|
||||
# (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
|
||||
f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \
|
||||
else f"{ctx[b]} = {ctx[v]};"),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
|
||||
lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
|
||||
]) + base_rewrite
|
||||
|
||||
def render_cast(self, dt:DType, val: str) -> str: return f"{self.type_map[dt]}({val})"
|
||||
def render_dtype(self, dt:DType, mutable=True) -> str: return "var"
|
||||
def render_load(self, x:str, dt:DType) -> str: return f"atomicLoad(&{x})" if is_packed(dt) else x
|
||||
def buf_map(self, dt:DType) -> str: return "atomic<u32>" if is_packed(dt) else self.type_map[dt.base]
|
||||
def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
|
||||
local_size = [u.src[0].ssimplify() for u in sorted([u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == 'l'], key=lambda u: u.arg)]
|
||||
if not local_size: local_size = [1]
|
||||
bind_it = iter(range(len(bufs)))
|
||||
external_local_bufs = [line.lstrip() for line in kernel if "var<workgroup>" in line]
|
||||
kernel[:] = [line for line in kernel if "var<workgroup>" not in line]
|
||||
prg = "enable f16;\n" if any(uop.dtype.base == dtypes.half for uop in uops) else ""
|
||||
prg += "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\n"
|
||||
prg += "@group(0) @binding(0)\nvar<uniform> INFINITY : f32;\n"
|
||||
prg += "\n".join((external_local_bufs or [])+[f"@group(0) @binding({next(bind_it)+1})" +
|
||||
f"{'var<storage,read_write>' if isinstance(dtype, PtrDType) else 'var<uniform>'}" +
|
||||
f"{name}:{f'array<{self.buf_map(dtype.base)}>' if isinstance(dtype,PtrDType) else self.buf_map(dtype)};" for name,(dtype,_) in bufs])
|
||||
prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>,"
|
||||
return prg + "@builtin(local_invocation_id) lindex: vec3<u32>) {\n" + "\n".join(kernel) + "\n}"
|
||||
|
||||
def supported_dtypes(self): return {dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.int32, dtypes.uint32,
|
||||
dtypes.float, *((dtypes.half,) if "shader-f16" in self.target.arch else ())}
|
||||
Reference in New Issue
Block a user