superrob/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ 0798119

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:42 -05:00
commit b42569dbca
4529 changed files with 1132125 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tinygrad_repo/test/README Normal file
View File

@@ -0,0 +1,5 @@
Three groups of tests run in CI
backend -- tests that run on each backend
null -- tests that don't require any backend
unit -- tests that only run on a single backend in CI

0
tinygrad_repo/test/__init__.py Normal file
View File

0
tinygrad_repo/test/amd/__init__.py Normal file
View File

945
tinygrad_repo/test/amd/disasm.py Normal file
View File

@@ -0,0 +1,945 @@
# RDNA3/RDNA4/CDNA disassembler
from __future__ import annotations
import re
from typing import Callable
from test.amd.helpers import decode_dpp16
from tinygrad.renderer.amd.dsl import Inst, Reg
# Special register mappings for disassembly
SPECIAL_GPRS = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'null', 125: 'm0', 126: 'exec_lo', 127: 'exec_hi',
128: '0', 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: '0x3e22f983', 253: 'scc'}
SPECIAL_GPRS_CDNA = {106: 'vcc_lo', 107: 'vcc_hi', 124: 'm0', 126: 'exec_lo', 127: 'exec_hi',
128: '0', 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: '0x3e22f983', 253: 'scc',
102: 'flat_scratch_lo', 103: 'flat_scratch_hi', 104: 'xnack_mask_lo', 105: 'xnack_mask_hi',
251: 'src_vccz', 252: 'src_execz'}
SPECIAL_PAIRS = {106: 'vcc', 126: 'exec'}
SPECIAL_PAIRS_CDNA = {106: 'vcc', 126: 'exec', 102: 'flat_scratch', 104: 'xnack_mask'}
def decode_src(v, cdna: bool = False) -> str:
"""Decode a source operand encoding to its string representation."""
v = _unwrap(v)
gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
if v in gprs: return gprs[v]
if v < 106: return f's{v}'
if 108 <= v < 124: return f'ttmp{v - 108}'
if 129 <= v <= 192: return str(v - 128) # positive integers 1-64
if 193 <= v <= 208: return str(-(v - 192)) # negative integers -1 to -16
if v >= 256: return f'v{v - 256}'
return f's{v}'
def _unwrap(v) -> int:
"""Unwrap Reg to int offset, or return int as-is."""
return v.offset if isinstance(v, Reg) else v
def _vi(v) -> int:
"""Get VGPR index from Reg or int (for v[N] fields that encode as 256+N)."""
off = _unwrap(v)
return off - 256 if off >= 256 else off
# ═══════════════════════════════════════════════════════════════════════════════
# LITERAL FORMATTING
# ═══════════════════════════════════════════════════════════════════════════════
_FLOAT_DEC = {240: 0.5, 241: -0.5, 242: 1.0, 243: -1.0, 244: 2.0, 245: -2.0, 246: 4.0, 247: -4.0}
def _lit(inst, v, neg=0, cdna=None) -> str:
"""Format literal/inline constant value."""
if cdna is None: cdna = _is_cdna(inst)
v = _unwrap(v)
if v == 255:
lit = inst._literal
if lit is None: return "0"
s = f"0x{lit:x}"
elif v in _FLOAT_DEC: s = str(_FLOAT_DEC[v])
elif 128 <= v <= 192: s = str(v - 128)
elif 193 <= v <= 208: s = str(-(v - 192))
elif v < 128: s = decode_src(v, cdna)
elif v >= 256: s = f"v{v - 256}"
else: s = decode_src(v, cdna)
return f"-{s}" if neg else s
# ═══════════════════════════════════════════════════════════════════════════════
# INSTRUCTION METADATA - fallback functions when inst.num_srcs()/inst.operands unavailable
# ═══════════════════════════════════════════════════════════════════════════════
def _num_srcs(inst) -> int:
"""Fallback: get number of source operands from instruction name."""
name = getattr(inst, 'op_name', '') or ''
n = name.upper()
# FMAC/MAC ops are 2-source (dst is implicit accumulator), but FMA/MAD ops are 3-source
if 'FMAC' in n or 'V_MAC_' in n: return 2
if any(x in n for x in ('FMA', 'MAD', 'CNDMASK', 'BFE', 'BFI', 'LERP', 'MED3', 'SAD', 'DIV_FMAS', 'DIV_FIXUP', 'DIV_SCALE', 'CUBE')): return 3
# PERMLANE_VAR ops are 2-source, but PERMLANE (non-VAR) are 3-source
if 'PERMLANE' in n and '_VAR' not in n: return 3
if any(x in n for x in ('_ADD3', '_LSHL_ADD', '_ADD_LSHL', '_LSHL_OR', '_AND_OR', 'OR3_B32', 'AND_OR_B32', 'ALIGNBIT',
'ALIGNBYTE', 'V_PERM_', 'XOR3', 'XAD', 'MULLIT', 'MINMAX', 'MAXMIN', 'MINIMUMMAXIMUM', 'MAXIMUMMINIMUM',
'MINIMUM3', 'MAXIMUM3', 'MIN3', 'MAX3', 'DOT2', 'CVT_PK_U8_F32', 'DOT4', 'DOT8', 'WMMA', 'SWMMAC')): return 3
return 2
# ═══════════════════════════════════════════════════════════════════════════════
# IMPORTS
# ═══════════════════════════════════════════════════════════════════════════════
from tinygrad.runtime.autogen.amd.rdna3.ins import (VOP1, VOP1_SDST, VOP1_SDST_LIT, VOP1_LIT, VOP2, VOP2_LIT, VOP3, VOP3_SDST, VOP3_SDST_LIT,
VOP3_LIT, VOP3SD, VOP3SD_LIT, VOP3P, VOP3P_LIT, VOPC, VOPC_LIT, VOPD, VOPD_LIT, VINTERP, SOP1, SOP1_LIT, SOP2, SOP2_LIT, SOPC, SOPC_LIT,
SOPK, SOPK_LIT, SOPP, SMEM, DS, FLAT, GLOBAL, SCRATCH, VOP2Op, VOPDOp, SOPPOp, HWREG, MSG)
from tinygrad.runtime.autogen.amd.rdna4.ins import (VOP1 as R4_VOP1, VOP1_SDST as R4_VOP1_SDST,
VOP1_SDST_LIT as R4_VOP1_SDST_LIT, VOP1_LIT as R4_VOP1_LIT,
VOP2 as R4_VOP2, VOP2_LIT as R4_VOP2_LIT, VOP3 as R4_VOP3, VOP3_SDST as R4_VOP3_SDST, VOP3_SDST_LIT as R4_VOP3_SDST_LIT, VOP3_LIT as R4_VOP3_LIT,
VOP3SD as R4_VOP3SD, VOP3SD_LIT as R4_VOP3SD_LIT, VOP3P as R4_VOP3P, VOP3P_LIT as R4_VOP3P_LIT, VOPC as R4_VOPC, VOPC_LIT as R4_VOPC_LIT,
VOPD as R4_VOPD, VOPD_LIT as R4_VOPD_LIT, VINTERP as R4_VINTERP, SOP1 as R4_SOP1, SOP1_LIT as R4_SOP1_LIT, SOP2 as R4_SOP2, SOP2_LIT as R4_SOP2_LIT,
SOPC as R4_SOPC, SOPC_LIT as R4_SOPC_LIT, SOPK as R4_SOPK, SOPK_LIT as R4_SOPK_LIT, SOPP as R4_SOPP, SMEM as R4_SMEM, DS as R4_DS,
VOPDOp as R4_VOPDOp, HWREG as HWREG_RDNA4, VFLAT as R4_FLAT, VGLOBAL as R4_GLOBAL, VSCRATCH as R4_SCRATCH)
from tinygrad.runtime.autogen.amd.cdna.ins import HWREG as HWREG_CDNA
def _is_cdna(inst: Inst) -> bool: return 'cdna' in inst.__class__.__module__
def _is_r4(inst: Inst) -> bool: return 'rdna4' in inst.__class__.__module__
# CDNA opcode name aliases for disasm (new name -> old name expected by tests)
_CDNA_DISASM_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32', 'v_fmamk_f32': 'v_madmk_f32', 'v_fmaak_f32': 'v_madak_f32'}
# ═══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ═══════════════════════════════════════════════════════════════════════════════
def _reg(p: str, b: int, n: int = 1) -> str: return f"{p}{_unwrap(b)}" if n == 1 else f"{p}[{_unwrap(b)}:{_unwrap(b)+n-1}]"
def _sreg(b: int, n: int = 1) -> str: return _reg("s", _unwrap(b), n)
def _vreg(b: int, n: int = 1) -> str:
b = _unwrap(b)
return _reg("v", b - 256 if b >= 256 else b, n)
def _areg(b: int, n: int = 1) -> str:
b = _unwrap(b)
return _reg("a", b - 256 if b >= 256 else b, n) # accumulator registers for GFX90a
def _ttmp(b, n: int = 1) -> str | None:
b = _unwrap(b)
return _reg("ttmp", b - 108, n) if 108 <= b <= 123 else None
def _fmt_sdst(v, n: int = 1, cdna: bool = False) -> str:
v = _unwrap(v)
if t := _ttmp(v, n): return t
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
gprs = SPECIAL_GPRS_CDNA if cdna else SPECIAL_GPRS
if n > 1: return pairs.get(v) or gprs.get(v) or _sreg(v, n) # also check gprs for null/m0
return gprs.get(v, f"s{v}")
def _fmt_src(v, n: int = 1, cdna: bool = False) -> str:
v = _unwrap(v)
if v == 253: return "src_scc" # SCC as source operand
if n == 1: return decode_src(v, cdna)
if v >= 256: return _vreg(v, n)
if v <= 101: return _sreg(v, n) # s0-s101 can be pairs, but 102+ are special on CDNA
pairs = SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS
if n == 2 and v in pairs: return pairs[v]
if v <= 105: return _sreg(v, n) # s102-s105 regular pairs for RDNA
if t := _ttmp(v, n): return t
return decode_src(v, cdna)
def _fmt_v16(v, base: int = 256, hi_thresh: int = 384) -> str:
v = _unwrap(v)
return f"v{(v - base) & 0x7f}.{'h' if v >= hi_thresh else 'l'}"
def _has(op: str, *subs) -> bool: return any(s in op for s in subs)
def _omod(v: int) -> str: return {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(v, "")
def _src16(inst, v: int) -> str:
v = _unwrap(v)
return _fmt_v16(v) if v >= 256 else _lit(inst, v) # format 16-bit src: vgpr.h/l or literal
def _mods(*pairs) -> str: return " ".join(m for c, m in pairs if c)
def _fmt_bits(label: str, val: int, count: int) -> str: return f"{label}:[{','.join(str((val >> i) & 1) for i in range(count))}]"
def _vop3_src(inst, v: int, neg: int, abs_: int, hi: int, n: int, f16: bool) -> str:
"""Format VOP3 source operand with modifiers."""
v = _unwrap(v)
if v == 255: s = _lit(inst, v) # literal constant takes priority
elif n > 1: s = _fmt_src(v, n)
elif f16 and v >= 256: s = f"v{v - 256}.h" if hi else f"v{v - 256}.l"
elif v == 253: s = "src_scc" # VOP3 sources use src_scc not scc
else: s = _lit(inst, v)
if abs_: s = f"|{s}|"
return f"-{s}" if neg else s
def _opsel_str(opsel: int, n: int, need: bool, is16_d: bool) -> str:
"""Format op_sel modifier string."""
if not need: return ""
dst_hi = (opsel >> 3) & 1
if n == 1: return f" op_sel:[{opsel & 1},{dst_hi}]"
# Use 4-element format if bit 2 is set (src2 selection used) or if 3+ sources
if n == 2 and not ((opsel >> 2) & 1): return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{dst_hi}]"
return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{dst_hi}]"
# ═══════════════════════════════════════════════════════════════════════════════
# DISASSEMBLER
# ═══════════════════════════════════════════════════════════════════════════════
def _disasm_vop1(inst: VOP1) -> str:
name, cdna = inst.op_name.lower() or f'vop1_op_{inst.op}', _is_cdna(inst)
name = name.replace('_e32', '') # Strip _e32 suffix
if any(x in name for x in ('v_nop', 'v_pipeflush', 'v_clrexcp')): return name # no operands
if 'readfirstlane' in name:
src = inst.src0.fmt() if inst.src0.offset >= 256 else decode_src(inst.src0.offset, cdna)
vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
return f"{name} {_fmt_sdst(vdst_off, 1, cdna)}, {src}"
bits = inst.canonical_op_bits
is16_dst, is16_src = not cdna and bits['d'] == 16, not cdna and bits['s0'] == 16
# Format dst
if is16_dst: dst = _fmt_v16(inst.vdst)
else: dst = inst.vdst.fmt()
# Format src
if inst.src0.offset == 255: src = _lit(inst, inst.src0)
elif is16_src and inst.src0.offset >= 256: src = _fmt_v16(inst.src0)
elif inst.src0.sz > 1: src = _fmt_src(inst.src0, inst.src0.sz, cdna)
else: src = _lit(inst, inst.src0)
return f"{name} {dst}, {src}"
_VOP2_CARRY_OUT = {'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'} # carry out only
_VOP2_CARRY_INOUT = {'v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'} # carry in and out (CDNA)
_VOP2_CARRY_INOUT_RDNA = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'} # carry in and out (RDNA)
def _disasm_vop2(inst: VOP2) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if cdna: name = _CDNA_DISASM_ALIASES.get(name, name) # apply CDNA aliases
suf = "" if cdna or name.endswith('_e32') or (not cdna and inst.op == VOP2Op.V_DOT2ACC_F32_F16_E32) else "_e32"
lit = inst._literal
is16 = not cdna and inst.canonical_op_bits['d'] == 16
# fmaak/madak: dst = src0 * vsrc1 + K, fmamk/madmk: dst = src0 * K + vsrc1
if 'fmaak' in name or 'madak' in name or (not cdna and inst.op in (VOP2Op.V_FMAAK_F32_E32, VOP2Op.V_FMAAK_F16_E32)):
if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}, 0x{lit:x}"
return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, 0x{lit:x}"
if 'fmamk' in name or 'madmk' in name or (not cdna and inst.op in (VOP2Op.V_FMAMK_F32_E32, VOP2Op.V_FMAMK_F16_E32)):
if lit is None: return f"op_{inst.op.value if hasattr(inst.op, 'value') else inst.op}"
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, 0x{lit:x}, {_fmt_v16(inst.vsrc1)}"
return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, 0x{lit:x}, {inst.vsrc1.fmt()}"
if is16: return f"{name}{suf} {_fmt_v16(inst.vdst)}, {_src16(inst, inst.src0)}, {_fmt_v16(inst.vsrc1)}"
vcc = "vcc" if cdna else "vcc_lo"
basename = name.replace('_e32', '')
if cdna and basename in _VOP2_CARRY_OUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}"
if cdna and basename in _VOP2_CARRY_INOUT: return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
if not cdna and basename in _VOP2_CARRY_INOUT_RDNA:
return f"{name}{suf} {inst.vdst.fmt()}, {vcc}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}, {vcc}"
sn0 = inst.canonical_op_regs.get('s0', 1)
if inst.vdst.sz > 1 or sn0 > 1 or inst.vsrc1.sz > 1:
src0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, sn0, cdna)
return f"{name.replace('_e32', '')} {inst.vdst.fmt()}, {src0}, {inst.vsrc1.fmt()}"
return f"{name}{suf} {inst.vdst.fmt()}, {_lit(inst, inst.src0)}, {inst.vsrc1.fmt()}" + (f", {vcc}" if name == 'v_cndmask_b32' else "")
def _disasm_vopc(inst: VOPC) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
bits = inst.canonical_op_bits
is16 = bits['s0'] == 16
if cdna:
s0 = _lit(inst, inst.src0) if inst.src0.offset == 255 else _fmt_src(inst.src0, inst.src0.sz, cdna)
return f"{name} vcc, {s0}, {inst.vsrc1.fmt()}" # CDNA VOPC always outputs vcc
# RDNA: v_cmpx_* writes to exec (no vcc), v_cmp_* writes to vcc_lo
has_vcc = 'cmpx' not in name
if inst.src0.offset == 255: s0 = _lit(inst, inst.src0)
elif inst.src0.sz > 1: s0 = inst.src0.fmt()
elif is16: s0 = _src16(inst, inst.src0.offset)
else: s0 = _lit(inst, inst.src0)
s1 = inst.vsrc1.fmt() if inst.vsrc1.sz > 1 else _fmt_v16(inst.vsrc1) if is16 else inst.vsrc1.fmt()
suf = "" if name.endswith('_e32') else "_e32"
return f"{name}{suf} vcc_lo, {s0}, {s1}" if has_vcc else f"{name}{suf} {s0}, {s1}"
NO_ARG_SOPP = {SOPPOp.S_BARRIER, SOPPOp.S_WAKEUP, SOPPOp.S_ICACHE_INV,
SOPPOp.S_WAIT_IDLE, SOPPOp.S_ENDPGM_SAVED, SOPPOp.S_CODE_END, SOPPOp.S_ENDPGM_ORDERED_PS_DONE, SOPPOp.S_TTRACEDATA}
def _disasm_sopp(inst: SOPP) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = _is_r4(inst)
# Ops that have no argument when simm16 == 0
no_arg_zero = {'s_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_wait_idle', 's_endpgm_saved',
's_endpgm_ordered_ps_done', 's_code_end'}
if name in no_arg_zero: return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
if name == 's_endpgm': return name if inst.simm16 == 0 else f"{name} {inst.simm16}"
if cdna:
if name == 's_waitcnt':
# GFX9 format: vmcnt[3:0]=bits[3:0], vmcnt[5:4]=bits[15:14], expcnt=bits[6:4], lgkmcnt=bits[11:8] (4 bits, max 15)
vm_lo, exp, lgkm, vm_hi = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 8) & 0xf, (inst.simm16 >> 14) & 0x3
vm = vm_lo | (vm_hi << 4)
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0xf else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
if name == 's_set_gpr_idx_mode':
flags = [n for i, n in enumerate(['SRC0', 'SRC1', 'SRC2', 'DST']) if inst.simm16 & (1 << i)]
return f"{name} gpr_idx({','.join(flags)})"
return f"{name} 0x{inst.simm16:x}" if inst.simm16 else name
# RDNA (use name-based checks instead of enum-based for cross-arch compatibility)
if name == 's_waitcnt':
if is_rdna4:
return f"{name} {inst.simm16}" if inst.simm16 else f"{name} 0"
vm, exp, lgkm = (inst.simm16 >> 10) & 0x3f, inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x3f
p = [f"vmcnt({vm})" if vm != 0x3f else "", f"expcnt({exp})" if exp != 7 else "", f"lgkmcnt({lgkm})" if lgkm != 0x3f else ""]
return f"s_waitcnt {' '.join(x for x in p if x) or '0'}"
if name == 's_delay_alu':
deps = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2',
'TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
skips = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
id0, skip, id1 = inst.simm16 & 0xf, (inst.simm16 >> 4) & 0x7, (inst.simm16 >> 7) & 0xf
def dep(v): return deps[v-1] if 0 < v <= len(deps) else str(v)
p = [f"instid0({dep(id0)})" if id0 else "", f"instskip({skips[skip]})" if skip else "", f"instid1({dep(id1)})" if id1 else ""]
return f"s_delay_alu {' | '.join(x for x in p if x) or '0'}"
if name.startswith(('s_cbranch', 's_branch')): return f"{name} {inst.simm16}"
return f"{name} 0x{inst.simm16:x}"
def _disasm_smem(inst: SMEM) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
if name in ('s_gl1_inv', 's_dcache_inv', 's_dcache_inv_vol', 's_dcache_wb', 's_dcache_wb_vol', 's_icache_inv'): return name
soe, imm = getattr(inst, 'soe', 0) or getattr(inst, 'soffset_en', 0), getattr(inst, 'imm', 1)
is_rdna4 = _is_r4(inst)
offset = inst.ioffset if is_rdna4 else getattr(inst, 'offset', 0) # type: ignore[attr-defined]
if cdna:
if soe and imm: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
elif imm: off_s = f"0x{offset:x}"
elif offset < 256: off_s = decode_src(offset, cdna)
else: off_s = decode_src(inst.soffset, cdna)
elif offset and inst.soffset != 124: off_s = f"{decode_src(inst.soffset, cdna)} offset:0x{offset:x}"
elif offset: off_s = f"0x{offset:x}"
else: off_s = decode_src(inst.soffset, cdna)
is_buffer = 'buffer' in name or 's_atc_probe_buffer' == name
sbase_idx, sbase_count = _unwrap(inst.sbase), 4 if is_buffer else 2
if sbase_count == 2: sbase_str = _fmt_src(sbase_idx, sbase_count, cdna)
elif sbase_idx <= 105: sbase_str = _sreg(sbase_idx, sbase_count)
else: sbase_str = _reg("ttmp", sbase_idx - 108, sbase_count)
if name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{name} {_unwrap(inst.sdata)}, {sbase_str}, {off_s}"
if 'prefetch' in name:
off = getattr(inst, 'ioffset', getattr(inst, 'offset', 0))
if off >= 0x800000: off = off - 0x1000000
off_s = f"0x{off:x}" if off > 255 else str(off)
soff_s = decode_src(inst.soffset, cdna) if inst.soffset != 124 else ("m0" if cdna else "null")
if 'pc_rel' in name: return f"{name} {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
return f"{name} {sbase_str}, {off_s}, {soff_s}, {_unwrap(inst.sdata)}"
# Use get_field_bits for register count
dst_n = inst.canonical_op_regs.get('d', 1)
th, scope = getattr(inst, 'th', 0), getattr(inst, 'scope', 0)
if is_rdna4: # RDNA4 uses th/scope instead of glc/dlc
th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
if th or scope:
th_names = ['TH_LOAD_RT', 'TH_LOAD_NT', 'TH_LOAD_HT', 'TH_LOAD_LU']
scope_names = ['SCOPE_CU', 'SCOPE_SE', 'SCOPE_DEV', 'SCOPE_SYS']
mods = (f" th:{th_names[th]}" if th else "") + (f" scope:{scope_names[scope]}" if scope else "")
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}{mods}"
if 'discard' in name: return f"{name} {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
R4_TH_LOAD = {1: 'TH_LOAD_NT', 2: 'TH_LOAD_HT', 3: 'TH_LOAD_LU', 4: 'TH_LOAD_RT_WB', 5: 'TH_LOAD_NT_WB'}
R4_TH_STORE = {1: 'TH_STORE_NT', 2: 'TH_STORE_HT', 3: 'TH_STORE_ST', 4: 'TH_STORE_RT_WB', 5: 'TH_STORE_NT_WB'}
R4_TH_ATOMIC = {1: 'TH_ATOMIC_RETURN', 2: 'TH_ATOMIC_NT', 3: 'TH_ATOMIC_RETURN_NT',
4: 'TH_ATOMIC_CASCADE_RT', 5: 'TH_ATOMIC_CASCADE_RETURN', 6: 'TH_ATOMIC_CASCADE_NT', 7: 'TH_ATOMIC_CASCADE_RETURN_NT'}
R4_SCOPE = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
def _disasm_flat(inst: FLAT) -> str:
name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
# Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
offset = inst.ioffset if r4 else inst.offset # type: ignore[attr-defined]
if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24) # sign extend 24-bit
elif seg != 'flat':
if cdna:
# CDNA: bit 12 is sign bit but not in offset field
raw = int.from_bytes(inst.to_bytes(), 'little')
off_val = offset | ((raw >> 12) & 1) << 12 # get bit 12
else:
off_val = offset
off_val = off_val if off_val < 4096 else off_val - 8192 # sign extend 13-bit
else:
off_val = offset
# Use get_field_bits: data for stores/atomics, d for loads
regs = inst.canonical_op_regs
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
off_s = f" offset:{off_val}" if off_val else ""
if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}" # type: ignore[attr-defined]
elif r4:
th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
if seg == 'flat': saddr_s = ""
elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
elif seg == 'scratch': saddr_s = f", {decode_src(inst.saddr, cdna)}"
elif _unwrap(inst.saddr) in (SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS):
saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
if 'addtid' in name: return f"{instr} {reg_fn((inst.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
# RDNA4: vaddr instead of addr, vsrc instead of data
addr = inst.vaddr if r4 else inst.addr # type: ignore[attr-defined]
data = inst.vsrc if r4 else inst.data # type: ignore[attr-defined]
# load_lds_* instructions: vaddr, saddr (no vdst, data goes to LDS)
if 'load_lds' in name:
addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
return f"{instr} {addr_s}{saddr_s}{mods}"
if seg == 'flat': addr_w = 2 # flat always uses 64-bit vaddr
elif cdna: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
else: addr_w = 1 if seg == 'scratch' or (_unwrap(inst.saddr) not in (0x7F, 124)) else 2
addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
if 'atomic' in name:
glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc) # type: ignore[attr-defined]
sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
return f"{instr} {reg_fn(inst.vdst, w)}, {addr_s}{saddr_s}{mods}"
def _disasm_ds(inst: DS) -> str:
name = inst.op_name.lower()
acc = getattr(inst, 'acc', 0)
reg_fn = _areg if acc else _vreg
gds = " gds" if getattr(inst, 'gds', 0) else ""
off = f" offset:{inst.offset0 | (inst.offset1 << 8)}" if inst.offset0 or inst.offset1 else ""
off2 = (" offset0:" + str(inst.offset0) if inst.offset0 else "") + (" offset1:" + str(inst.offset1) if inst.offset1 else "")
# Use get_field_bits: data for stores/writes/atomics, d for loads
regs = inst.canonical_op_regs
w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'write' in name or ('load' not in name and 'read' not in name) else regs.get('d', 1)
d0, d1, dst, addr = reg_fn(inst.data0, w), reg_fn(inst.data1, w), reg_fn(inst.vdst, w), _vreg(inst.addr)
if name == 'ds_nop': return name
if name == 'ds_bvh_stack_rtn_b32': return f"{name} {_vreg(inst.vdst)}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, 4)}{off}{gds}"
if 'bvh_stack_push' in name:
d1_regs = 8 if 'push8' in name else 4
vdst_regs = 2 if 'pop2' in name else 1
vdst_s = _vreg(inst.vdst, vdst_regs) if vdst_regs > 1 else _vreg(inst.vdst)
return f"{name} {vdst_s}, {addr}, {_vreg(inst.data0)}, {_vreg(inst.data1, d1_regs)}{off}{gds}"
if 'gws_sema' in name and 'sema_br' not in name: return f"{name}{off}{gds}"
if 'gws_' in name: return f"{name} {addr}{off}{gds}"
if name in ('ds_consume', 'ds_append'): return f"{name} {reg_fn(inst.vdst)}{off}{gds}"
if 'gs_reg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {reg_fn(inst.data0)}{off}{gds}"
if '2addr' in name:
if 'load' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
if 'store' in name and 'xchg' not in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
if 'write2' in name: return f"{name} {addr}, {d0}, {d1}{off2}{gds}"
if 'read2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}{off2}{gds}"
if 'xchg2' in name: return f"{name} {reg_fn(inst.vdst, regs.get('d', 1))}, {addr}, {d0}, {d1}{off2}{gds}"
if 'load' in name or ('read' in name and 'read2' not in name):
return f"{name} {reg_fn(inst.vdst)}{off}{gds}" if 'addtid' in name else f"{name} {dst}, {addr}{off}{gds}"
if ('store' in name or 'write' in name) and not _has(name, 'cmp', 'xchg', 'write2'):
return f"{name} {reg_fn(inst.data0)}{off}{gds}" if 'addtid' in name else f"{name} {addr}, {d0}{off}{gds}"
if 'swizzle' in name or name == 'ds_ordered_count': return f"{name} {reg_fn(inst.vdst)}, {addr}{off}{gds}"
if 'permute' in name: return f"{name} {reg_fn(inst.vdst)}, {addr}, {reg_fn(inst.data0)}{off}{gds}"
if 'condxchg' in name: return f"{name} {reg_fn(inst.vdst, 2)}, {addr}, {reg_fn(inst.data0, 2)}{off}{gds}"
if _has(name, 'cmpst', 'mskor', 'wrap'):
return f"{name} {dst}, {addr}, {d0}, {d1}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}, {d1}{off}{gds}"
return f"{name} {dst}, {addr}, {d0}{off}{gds}" if '_rtn' in name else f"{name} {addr}, {d0}{off}{gds}"
def _disasm_vop3(inst: VOP3) -> str:
name = inst.op_name.lower()
bits = inst.canonical_op_bits
# RDNA4 v_s_* scalar VOP3 instructions - vdst is SGPR (VGPRField adds 256)
if name.startswith('v_s_'):
s0v = _unwrap(inst.src0)
if s0v == 255: src = _lit(inst, inst.src0)
elif s0v == 253: src = "src_scc"
else: src = _fmt_src(inst.src0, max(1, bits['s0'] // 32))
if inst.neg & 1: src = f"-{src}"
if inst.abs & 1: src = f"|{src}|"
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
vdst_raw = _unwrap(inst.vdst)
return f"{name} s{vdst_raw - 256 if vdst_raw >= 256 else vdst_raw}, {src}" + (" clamp" if clamp else "") + _omod(inst.omod)
# Use get_field_bits for register sizes and 16-bit detection
r0, r1, r2 = max(1, bits['s0'] // 32), max(1, bits['s1'] // 32), max(1, bits['s2'] // 32)
is16_d, is16_s, is16_s2 = bits['d'] == 16, bits['s0'] == 16, bits['s2'] == 16
s0 = _vop3_src(inst, inst.src0, inst.neg&1, inst.abs&1, inst.opsel&1, r0, is16_s)
s1 = _vop3_src(inst, inst.src1, inst.neg&2, inst.abs&2, inst.opsel&2, r1, is16_s)
s2 = _vop3_src(inst, inst.src2, inst.neg&4, inst.abs&4, inst.opsel&4, r2, is16_s2)
# Format destination
if 'readlane' in name:
vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
dst = _fmt_sdst(vdst_off, 1)
elif is16_d: dst = f"{inst.vdst.fmt()}.h" if (inst.opsel & 8) else f"{inst.vdst.fmt()}.l"
else: dst = inst.vdst.fmt()
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
cl, om = " clamp" if clamp else "", _omod(inst.omod)
nonvgpr_opsel = ((inst.src0.offset < 256 and (inst.opsel & 1)) or (inst.src1.offset < 256 and (inst.opsel & 2))
or (inst.src2.offset < 256 and (inst.opsel & 4)))
need_opsel = nonvgpr_opsel or (inst.opsel and not is16_s)
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
e64 = "" if name.endswith('_e64') else "_e64"
if op_val < 256: # VOPC
vdst_off = inst.vdst.offset - 256 if inst.vdst.offset >= 256 else inst.vdst.offset
return f"{name}{e64} {s0}, {s1}{cl}" if name.startswith('v_cmpx') else f"{name}{e64} {_fmt_sdst(vdst_off, 1)}, {s0}, {s1}{cl}"
if op_val < 384: # VOP2
n = inst.num_srcs() or 2
os = _opsel_str(inst.opsel, n, need_opsel, is16_d)
return f"{name}{e64} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name}{e64} {dst}, {s0}, {s1}{os}{cl}{om}"
if op_val < 512: # VOP1
if re.match(r'v_cvt_f32_(bf|fp)8', name) and inst.opsel:
os = f" byte_sel:{((inst.opsel & 1) << 1) | ((inst.opsel >> 1) & 1)}"
else:
os = _opsel_str(inst.opsel, 1, need_opsel, is16_d)
if 'v_nop' in name or 'v_pipeflush' in name: return f"{name}{e64}"
return f"{name}{e64} {dst}, {s0}{os}{cl}{om}"
# Native VOP3
n = inst.num_srcs() or 2
os = f" byte_sel:{inst.opsel >> 2}" if 'cvt_sr' in name and inst.opsel else _opsel_str(inst.opsel, n, need_opsel, is16_d)
return f"{name} {dst}, {s0}, {s1}, {s2}{os}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{os}{cl}{om}"
def _disasm_vop3sd(inst: VOP3SD) -> str:
name = inst.op_name.lower()
def src(reg, neg):
s = _lit(inst, reg.offset) if reg.offset == 255 else ("src_scc" if reg.offset == 253 else (reg.fmt() if reg.sz > 1 else _lit(inst, reg.offset)))
return f"neg({s})" if neg and reg.offset == 255 else (f"-{s}" if neg else s)
s0, s1, s2 = src(inst.src0, inst.neg & 1), src(inst.src1, inst.neg & 2), src(inst.src2, inst.neg & 4)
# VOP3SD: _co_ ops (add/sub) without _ci_ have only 2 sources, all others (mad, div_scale, _co_ci_) have 3 sources
has_only_two_srcs = '_co_' in name and '_ci_' not in name and 'mad' not in name
srcs = f"{s0}, {s1}" if has_only_two_srcs else f"{s0}, {s1}, {s2}"
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
return f"{name} {inst.vdst.fmt()}, {_fmt_sdst(inst.sdst, 1)}, {srcs}{' clamp' if clamp else ''}{_omod(inst.omod)}"
def _disasm_vopd(inst: VOPD) -> str:
lit = inst._literal
op_enum = R4_VOPDOp if _is_r4(inst) else VOPDOp
nx, ny = op_enum(inst.opx).name.lower(), op_enum(inst.opy).name.lower()
def half(n, vd, s0, vs1):
vd, vs1 = _vi(vd), _vi(vs1)
if 'mov' in n: return f"{n} v{vd}, {_lit(inst, s0)}"
if 'fmamk' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, 0x{lit:x}, v{vs1}"
if 'fmaak' in n and lit: return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}, 0x{lit:x}"
return f"{n} v{vd}, {_lit(inst, s0)}, v{vs1}"
return f"{half(nx, inst.vdstx, inst.srcx0, inst.vsrcx1)} :: {half(ny, inst.vdsty, inst.srcy0, inst.vsrcy1)}"
def _disasm_vop3p(inst: VOP3P) -> str:
name = inst.op_name.lower()
is_swmmac, n, is_fma_mix = 'swmmac' in name, inst.num_srcs() or 2, 'fma_mix' in name
def get_src(reg):
return _lit(inst, reg.offset) if reg.offset == 255 else reg.fmt()
src0, src1, src2, dst = get_src(inst.src0), get_src(inst.src1), get_src(inst.src2), inst.vdst.fmt()
opsel_hi = inst.opsel_hi | (inst.opsel_hi2 << 2)
clamp = getattr(inst, 'cm', None) or getattr(inst, 'clmp', 0)
if is_fma_mix:
def m(s, neg, abs_): return f"-{f'|{s}|' if abs_ else s}" if neg else (f"|{s}|" if abs_ else s)
src0, src1, src2 = m(src0, inst.neg & 1, inst.neg_hi & 1), m(src1, inst.neg & 2, inst.neg_hi & 2), m(src2, inst.neg & 4, inst.neg_hi & 4)
mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
+ ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi else []) + (["clamp"] if clamp else []))
elif is_swmmac:
mods = ([f"index_key:{inst.opsel}"] if inst.opsel else []) + ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else []) + \
([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else [])
else:
opsel_hi_default = 7 if n == 3 else 3
mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
+ ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else [])
+ ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else [])
+ ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if clamp else []))
mod_s = ' ' + ' '.join(mods) if mods else ''
return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
def _disasm_sop1(inst: SOP1) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
# Use get_field_bits for register sizes
regs = inst.canonical_op_regs
dst_regs, src_regs = regs.get('d', 1), regs.get('s0', 1)
src = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, src_regs, cdna)
if not cdna:
if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}"
if 'setpc_b64' in name or 'rfe_b64' in name: return f"{name} {src}"
if 'swappc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2)}, {src}"
if 'sendmsg_rtn' in name:
v = _unwrap(inst.ssrc0)
try: msg_str = MSG(v).name if v != 255 else None # MSG_RTN_ILLEGAL_MSG (255) not supported by LLVM
except ValueError: msg_str = None
return f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, sendmsg({msg_str})" if msg_str else f"{name} {_fmt_sdst(inst.sdst, dst_regs)}, 0x{v:x}"
sop1_src_only = ('S_ALLOC_VGPR', 'S_SLEEP_VAR', 'S_BARRIER_SIGNAL', 'S_BARRIER_SIGNAL_ISFIRST',
'S_BARRIER_INIT', 'S_BARRIER_JOIN', 'S_SET_GPR_IDX_IDX', 'S_CBRANCH_JOIN')
if inst.op_name in sop1_src_only: return f"{name} {src}"
if cdna:
if 'getpc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}"
if 'setpc_b64' in name or 'rfe_b64' in name: return f"{name} {src}"
if 'swappc_b64' in name: return f"{name} {_fmt_sdst(inst.sdst, 2, cdna)}, {src}"
return f"{name} {_fmt_sdst(inst.sdst, dst_regs, cdna)}, {src}"
def _disasm_sop2(inst: SOP2) -> str:
cdna, name = _is_cdna(inst), inst.op_name.lower()
lit = inst._literal
# Use get_field_bits for register sizes
regs = inst.canonical_op_regs
dn, s0n, s1n = regs['d'], regs['s0'], regs['s1']
s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, s0n, cdna)
s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, s1n, cdna)
dst = _fmt_sdst(inst.sdst, dn, cdna)
if 'fmamk' in name and lit is not None: return f"{name} {dst}, {s0}, 0x{lit:x}, {s1}"
if 'fmaak' in name and lit is not None: return f"{name} {dst}, {s0}, {s1}, 0x{lit:x}"
if name in ('s_cbranch_g_fork', 's_rfe_restore_b64'): return f"{name} {s0}, {s1}" # no destination
return f"{name} {dst}, {s0}, {s1}"
def _disasm_sopc(inst: SOPC) -> str:
cdna, regs, name = _is_cdna(inst), inst.canonical_op_regs, inst.op_name.lower()
s0 = _lit(inst, inst.ssrc0) if _unwrap(inst.ssrc0) == 255 else _fmt_src(inst.ssrc0, regs['s0'], cdna)
if name == 's_set_gpr_idx_on':
imm = _unwrap(inst.ssrc1) & 0xf
flags = [n for i, n in enumerate(['SRC0', 'SRC1', 'SRC2', 'DST']) if imm & (1 << i)]
return f"{name} {s0}, gpr_idx({','.join(flags)})"
s1 = _lit(inst, inst.ssrc1) if _unwrap(inst.ssrc1) == 255 else _fmt_src(inst.ssrc1, regs['s1'], cdna)
return f"{name} {s0}, {s1}"
_HWREG_BLACKLIST = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_REG_FLUSH_IB', 'HW_REG_SHADER_TBA_LO', 'HW_REG_SHADER_TBA_HI',
'HW_REG_SHADER_FLAT_SCRATCH_LO', 'HW_REG_SHADER_FLAT_SCRATCH_HI', 'HW_REG_SHADER_CYCLES'}
_HWREG_BLACKLIST_CDNA = {'HW_REG_PC_LO', 'HW_REG_PC_HI', 'HW_REG_IB_DBG1', 'HW_REG_FLUSH_IB', 'HW_REG_SQ_SHADER_TBA_LO', 'HW_REG_SQ_SHADER_TBA_HI',
'HW_REG_SQ_SHADER_TMA_LO', 'HW_REG_SQ_SHADER_TMA_HI', 'HW_REG_SQ_PERF_SNAPSHOT_DATA', 'HW_REG_SQ_PERF_SNAPSHOT_DATA1',
'HW_REG_SQ_PERF_SNAPSHOT_PC_LO', 'HW_REG_SQ_PERF_SNAPSHOT_PC_HI', 'HW_REG_XCC_ID'}
def _disasm_sopk(inst: SOPK) -> str:
name, cdna = inst.op_name.lower(), _is_cdna(inst)
is_rdna4 = _is_r4(inst)
hw = HWREG_CDNA if cdna else (HWREG_RDNA4 if is_rdna4 else HWREG)
blacklist = _HWREG_BLACKLIST_CDNA if cdna else _HWREG_BLACKLIST
def fmt_hwreg(hid, hoff, hsz):
try: hr_name = hw(hid).name.replace("HW_REG_WAVE_", "HW_REG_")
except ValueError: return f"0x{inst.simm16:x}"
if hr_name in blacklist: return f"0x{inst.simm16:x}"
return f"hwreg({hr_name})" if hoff == 0 and hsz == 32 else f"hwreg({hr_name}, {hoff}, {hsz})"
if name == 's_setreg_imm32_b32':
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
return f"{name} {fmt_hwreg(hid, hoff, hsz)}, 0x{inst._literal:x}"
if name == 's_version': return f"{name} 0x{inst.simm16:x}"
if name in ('s_setreg_b32', 's_getreg_b32'):
hid, hoff, hsz = inst.simm16 & 0x3f, (inst.simm16 >> 6) & 0x1f, ((inst.simm16 >> 11) & 0x1f) + 1
hs = fmt_hwreg(hid, hoff, hsz)
return f"{name} {hs}, {_fmt_sdst(inst.sdst, 1, cdna)}" if 'setreg' in name else f"{name} {_fmt_sdst(inst.sdst, 1, cdna)}, {hs}"
if name in ('s_subvector_loop_begin', 's_subvector_loop_end'):
return f"{name} {_fmt_sdst(inst.sdst, 1)}, 0x{inst.simm16:x}"
return f"{name} {_fmt_sdst(inst.sdst, inst.canonical_op_regs['d'], cdna)}, 0x{inst.simm16:x}"
def _disasm_vinterp(inst: VINTERP) -> str:
mods = _mods((inst.waitexp, f"wait_exp:{inst.waitexp}"), (inst.clmp, "clamp"))
s0, s1, s2 = _lit(inst, inst.src0, inst.neg & 1), _lit(inst, inst.src1, inst.neg & 2), _lit(inst, inst.src2, inst.neg & 4)
return f"{inst.op_name.lower()} {inst.vdst.fmt()}, {s0}, {s1}, {s2}" + (" " + mods if mods else "")
DISASM_HANDLERS: dict[type, Callable[..., str]] = {
VOP1: _disasm_vop1, VOP1_SDST: _disasm_vop1, VOP1_SDST_LIT: _disasm_vop1, VOP1_LIT: _disasm_vop1,
VOP2: _disasm_vop2, VOP2_LIT: _disasm_vop2, VOPC: _disasm_vopc, VOPC_LIT: _disasm_vopc,
VOP3: _disasm_vop3, VOP3_SDST: _disasm_vop3, VOP3_SDST_LIT: _disasm_vop3, VOP3_LIT: _disasm_vop3,
VOP3SD: _disasm_vop3sd, VOP3SD_LIT: _disasm_vop3sd,
VOPD: _disasm_vopd, VOPD_LIT: _disasm_vopd, VOP3P: _disasm_vop3p, VOP3P_LIT: _disasm_vop3p,
VINTERP: _disasm_vinterp, SOPP: _disasm_sopp, SMEM: _disasm_smem, DS: _disasm_ds, FLAT: _disasm_flat, GLOBAL: _disasm_flat, SCRATCH: _disasm_flat,
SOP1: _disasm_sop1, SOP1_LIT: _disasm_sop1, SOP2: _disasm_sop2, SOP2_LIT: _disasm_sop2,
SOPC: _disasm_sopc, SOPC_LIT: _disasm_sopc, SOPK: _disasm_sopk, SOPK_LIT: _disasm_sopk,
# RDNA4
R4_VOP1: _disasm_vop1, R4_VOP1_SDST: _disasm_vop1, R4_VOP1_SDST_LIT: _disasm_vop1, R4_VOP1_LIT: _disasm_vop1,
R4_VOP2: _disasm_vop2, R4_VOP2_LIT: _disasm_vop2, R4_VOPC: _disasm_vopc, R4_VOPC_LIT: _disasm_vopc,
R4_VOP3: _disasm_vop3, R4_VOP3_SDST: _disasm_vop3, R4_VOP3_SDST_LIT: _disasm_vop3, R4_VOP3_LIT: _disasm_vop3,
R4_VOP3SD: _disasm_vop3sd, R4_VOP3SD_LIT: _disasm_vop3sd, R4_VOP3P: _disasm_vop3p, R4_VOP3P_LIT: _disasm_vop3p,
R4_FLAT: _disasm_flat, R4_GLOBAL: _disasm_flat, R4_SCRATCH: _disasm_flat,
R4_VOPD: _disasm_vopd, R4_VOPD_LIT: _disasm_vopd, R4_VINTERP: _disasm_vinterp, R4_SOPP: _disasm_sopp, R4_SMEM: _disasm_smem, R4_DS: _disasm_ds,
R4_SOP1: _disasm_sop1, R4_SOP1_LIT: _disasm_sop1, R4_SOP2: _disasm_sop2, R4_SOP2_LIT: _disasm_sop2,
R4_SOPC: _disasm_sopc, R4_SOPC_LIT: _disasm_sopc, R4_SOPK: _disasm_sopk, R4_SOPK_LIT: _disasm_sopk}
def disasm(inst: Inst) -> str: return DISASM_HANDLERS[type(inst)](inst)
# ═══════════════════════════════════════════════════════════════════════════════
# CDNA DISASSEMBLER SUPPORT
# ═══════════════════════════════════════════════════════════════════════════════
from tinygrad.runtime.autogen.amd.cdna.ins import (VOP1 as CDNA_VOP1, VOP1_LIT as CDNA_VOP1_LIT,
VOP1_SDWA as CDNA_VOP1_SDWA, VOP1_DPP16 as CDNA_VOP1_DPP16,
VOP2 as CDNA_VOP2, VOP2_LIT as CDNA_VOP2_LIT, VOP2_SDWA as CDNA_VOP2_SDWA, VOP2_DPP16 as CDNA_VOP2_DPP16,
VOPC as CDNA_VOPC, VOPC_LIT as CDNA_VOPC_LIT, VOPC_SDWA_SDST as CDNA_VOPC_SDWA_SDST,
VOP3 as CDNA_VOP3, VOP3_SDST as CDNA_VOP3_SDST, VOP3SD as CDNA_VOP3SD, VOP3P as CDNA_VOP3P, VOP3P_MFMA as CDNA_VOP3P_MFMA, VOP3PX2 as CDNA_VOP3PX2,
SOP1 as CDNA_SOP1, SOP1_LIT as CDNA_SOP1_LIT, SOP2 as CDNA_SOP2, SOP2_LIT as CDNA_SOP2_LIT,
SOPC as CDNA_SOPC, SOPC_LIT as CDNA_SOPC_LIT, SOPK as CDNA_SOPK, SOPK_LIT as CDNA_SOPK_LIT,
SOPP as CDNA_SOPP, SMEM as CDNA_SMEM, DS as CDNA_DS,
FLAT as CDNA_FLAT, GLOBAL as CDNA_GLOBAL, SCRATCH as CDNA_SCRATCH, MUBUF as CDNA_MUBUF)
def _cdna_src(inst, v, neg, abs_=0, n=1):
s = _lit(inst, v) if v == 255 else _fmt_src(v, n, cdna=True)
if abs_: s = f"|{s}|"
return f"neg({s})" if neg and v == 255 else (f"-{s}" if neg else s)
_CDNA_VOP3_ALIASES = {'v_fmac_f64': 'v_mul_legacy_f32', 'v_dot2c_f32_bf16': 'v_mac_f32'}
def _disasm_vop3a(inst) -> str:
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
name = inst.op_name.lower() or f'vop3a_op_{op_val}'
n = inst.num_srcs() or _num_srcs(inst)
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
# _sr_ instructions use 4-element op_sel (src2 for byte selection)
opsel_n = 3 if '_sr_' in name and n == 2 else n
opsel = _opsel_str(inst.opsel, opsel_n, inst.opsel != 0, False)
orig_name = name
name = _CDNA_VOP3_ALIASES.get(name, name)
if name != orig_name:
s0, s1 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, 1), _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, 1)
s2 = ""
dst = _vreg(inst.vdst)
else:
regs = inst.canonical_op_regs
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
s0 = _cdna_src(inst, inst.src0, inst.neg&1, inst.abs&1, r0)
s1 = _cdna_src(inst, inst.src1, inst.neg&2, inst.abs&2, r1)
s2 = _cdna_src(inst, inst.src2, inst.neg&4, inst.abs&4, r2)
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
if op_val >= 512:
return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
if op_val < 256:
# VOPC: vdst is actually sdst (SGPR pair), but VGPRField adds 256 to the offset
sdst_val = _unwrap(inst.vdst)
if sdst_val >= 256: sdst_val -= 256
sdst = _fmt_sdst(sdst_val, 2, cdna=True)
return f"{name} {sdst}, {s0}, {s1}{cl}"
if 320 <= op_val < 512:
if name in ('v_nop', 'v_clrexcp', 'v_nop_e64', 'v_clrexcp_e64'): return name.replace('_e64', '')
return f"{name} {dst}, {s0}{cl}{om}"
if name == 'v_cndmask_b32':
s2 = _fmt_src(inst.src2, 2, cdna=True)
return f"{name} {dst}, {s0}, {s1}, {s2}{cl}{om}"
return f"{name} {dst}, {s0}, {s1}, {s2}{opsel}{cl}{om}" if n == 3 else f"{name} {dst}, {s0}, {s1}{opsel}{cl}{om}"
def _disasm_vop3b(inst) -> str:
op_val = inst.op.value if hasattr(inst.op, 'value') else inst.op
name, cdna = inst.op_name.lower() or f'vop3b_op_{op_val}', _is_cdna(inst)
n = inst.num_srcs() or _num_srcs(inst)
regs = inst.canonical_op_regs
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
s0 = _cdna_src(inst, inst.src0, inst.neg&1, n=r0)
s1 = _cdna_src(inst, inst.src1, inst.neg&2, n=r1)
s2 = _cdna_src(inst, inst.src2, inst.neg&4, n=r2)
# CDNA VOP3_SDST uses vdst field for sdst (but vdst adds 256), RDNA uses separate sdst field
sdst_val = getattr(inst, 'sdst', None)
if sdst_val is None and hasattr(inst, 'vdst'):
sdst_val = _unwrap(inst.vdst)
if sdst_val >= 256: sdst_val -= 256 # VGPRField adds 256, remove it for SGPR
# For CDNA VOP3_SDST (VOPC->VOP3), vdst is the scalar dest (sdst), there's no vdst output
if cdna and 'v_cmp' in name:
sdst = _fmt_sdst(sdst_val, 2, cdna=True)
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
return f"{name} {sdst}, {s0}, {s1}{cl}{om}"
dst = _vreg(inst.vdst, dregs) if dregs > 1 else _vreg(inst.vdst)
sdst = _fmt_sdst(sdst_val, 2, cdna=cdna)
cl, om = " clamp" if inst.clmp else "", _omod(inst.omod)
if name in ('v_addc_co_u32', 'v_subb_co_u32', 'v_subbrev_co_u32'):
s2 = _fmt_src(inst.src2, 2, cdna=cdna)
return f"{name} {dst}, {sdst}, {s0}, {s1}, {s2}{cl}{om}" if n == 3 else f"{name} {dst}, {sdst}, {s0}, {s1}{cl}{om}"
def _disasm_cdna_vop3p(inst) -> str:
name, n = inst.op_name.lower(), inst.num_srcs() or 2
is_mfma = 'mfma' in name or 'smfmac' in name
is_accvgpr = 'accvgpr' in name
def get_src(v, sc): return _lit(inst, v) if v == 255 else _fmt_src(v, sc, cdna=True)
# Handle accvgpr read/write (accumulator register operations)
if is_accvgpr:
src0_off = _unwrap(inst.src0)
vdst_off = _vi(inst.vdst)
if 'read' in name:
# v_accvgpr_read_b32 vN, aM - reads from accumulator to VGPR
return f"{name}_b32 v{vdst_off}, a{src0_off - 256 if src0_off >= 256 else src0_off}"
if 'write' in name:
# v_accvgpr_write_b32 aM, src - writes to accumulator from source
src = _lit(inst, inst.src0) if src0_off == 255 else (f"v{src0_off - 256}" if src0_off >= 256 else decode_src(src0_off, cdna=True))
return f"{name}_b32 a{vdst_off}, {src}"
# Handle v_mfma_ld_scale_b32 - special 2-operand format: v_mfma_ld_scale_b32 src0, src1
if 'ld_scale' in name:
src0, src1 = get_src(inst.src0, 1), get_src(inst.src1, 1)
mods = ([_fmt_bits("op_sel", inst.opsel, 2)] if inst.opsel else []) + \
([_fmt_bits("op_sel_hi", inst.opsel_hi, 2)] if inst.opsel_hi != 3 else [])
return f"{name} {src0}, {src1}{' ' + ' '.join(mods) if mods else ''}"
# Handle MFMA instructions with accumulator destinations
if is_mfma:
regs = inst.canonical_op_regs
dregs, r0, r1, r2 = regs['d'], regs['s0'], regs['s1'], regs['s2']
# Infer register counts from instruction name if not in operands table (e.g., v_mfma_f32_32x32x4_xf32)
if dregs == 1:
if '32x32' in name: dregs, r0, r1, r2 = 16, 2, 2, 16
elif '16x16' in name: dregs, r0, r1, r2 = 4, 2, 2, 4
# MFMA reuses VOP3P fields differently: clmp -> acc_cd (dest is acc), opsel_hi -> acc (src1/src2 are acc)
# acc field (bits 60-59): bit 0 = src2 is acc (always for MFMA), bit 1 = src1 is acc
acc = inst.opsel_hi # opsel_hi field maps to acc for MFMA
acc_cd = inst.clmp # clmp field maps to acc_cd for MFMA (dest is accumulator)
is_smfmac = 'smfmac' in name # SMFMAC has different operand semantics
# Format sources: src0 is always VGPR, src1/src2 depend on acc bits
def mfma_src(v, sc, is_acc):
v = _unwrap(v)
if v == 255: return _lit(inst, v)
if 128 <= v <= 208 or 240 <= v <= 248: return _lit(inst, v)
base = v - 256 if v >= 256 else v
if is_acc: return _areg(base, sc)
return _vreg(base, sc)
src0 = get_src(inst.src0, r0) # src0 is always VGPR
src1 = mfma_src(inst.src1, r1, acc & 2) # bit 1 = src1 is acc
# For SMFMAC, src2 is always a VGPR index (1 register), not accumulator
src2 = _vreg(inst.src2) if is_smfmac else mfma_src(inst.src2, r2, acc_cd)
dst = _areg(inst.vdst, dregs) if acc_cd else _vreg(inst.vdst, dregs)
# MFMA uses neg:[...] not neg_lo:[...], and doesn't support op_sel_hi or clamp
# Only f64 MFMA instructions support neg modifier
# f8f6f4 MFMA instructions support cbsz/blgp modifiers
mods = []
if 'f8f6f4' in name:
if inst.neg_hi: mods.append(f"cbsz:{inst.neg_hi}")
if inst.neg: mods.append(f"blgp:{inst.neg}")
elif inst.neg and 'f64' in name:
mods.append(_fmt_bits("neg", inst.neg, n))
return f"{name} {dst}, {src0}, {src1}, {src2}{' ' + ' '.join(mods) if mods else ''}"
# Standard VOP3P instructions
src0, src1, src2, dst = get_src(inst.src0, 1), get_src(inst.src1, 1), get_src(inst.src2, 1), _vreg(inst.vdst)
opsel_hi = inst.opsel_hi # CDNA VOP3P only has 2 bits for opsel_hi (no opsel_hi2)
opsel_hi_default = 3 # CDNA default is 0b11 (2 bits), not 0b111 like RDNA
mods = (([_fmt_bits("op_sel", inst.opsel, n)] if inst.opsel else [])
+ ([_fmt_bits("op_sel_hi", opsel_hi, n)] if opsel_hi != opsel_hi_default else [])
+ ([_fmt_bits("neg_lo", inst.neg, n)] if inst.neg else [])
+ ([_fmt_bits("neg_hi", inst.neg_hi, n)] if inst.neg_hi else []) + (["clamp"] if inst.clmp else []))
mod_s = ' ' + ' '.join(mods) if mods else ''
return f"{name} {dst}, {src0}, {src1}, {src2}{mod_s}" if n == 3 else f"{name} {dst}, {src0}, {src1}{mod_s}"
def _disasm_mubuf(inst) -> str:
name = inst.op_name.lower()
# Determine vdata register count from instruction name
nregs = 4 if 'xyzw' in name else 3 if 'xyz' in name else 2 if 'xy' in name or 'x2' in name or 'f64' in name or 'dwordx2' in name else 1
vdata = _vreg(inst.vdata, nregs)
vaddr = _vreg(inst.vaddr) if inst.offen or inst.idxen else None
srsrc = str(inst.srsrc)
soffset_val = _unwrap(inst.soffset)
soffset = f"s{soffset_val}" if soffset_val < 128 else "off"
offset = f" offset:{inst.offset}" if inst.offset else ""
offen = " offen" if inst.offen else ""
idxen = " idxen" if inst.idxen else ""
lds = " lds" if inst.lds else ""
sc0 = " sc0" if inst.sc0 else ""
sc1 = " sc1" if inst.sc1 else ""
nt = " nt" if inst.nt else ""
# Handle special cases
if name in ('buffer_wbl2', 'buffer_inv'):
return f"{name}{sc0}{sc1}"
if vaddr:
return f"{name} {vdata}, {vaddr}, {srsrc}, {soffset}{offen}{idxen}{offset}{sc0}{nt}{sc1}{lds}"
return f"{name} {vdata}, off, {srsrc}, {soffset}{offset}{sc0}{nt}{sc1}{lds}"
_SDWA_SEL = {0: 'BYTE_0', 1: 'BYTE_1', 2: 'BYTE_2', 3: 'BYTE_3', 4: 'WORD_0', 5: 'WORD_1', 6: 'DWORD'}
def _disasm_vop1_sdwa(inst) -> str:
name = inst.op_name.lower().replace('_e32', '')
regs = inst.canonical_op_regs
dst = _vreg(inst.vdst, regs['d'])
# When s0=1, vsrc0 is SGPR/constant (VGPRField adds 256, so subtract it back)
if inst.s0 == 0: src0 = _vreg(inst.vsrc0, regs['s0'])
else:
raw = _unwrap(inst.vsrc0) - 256 # VGPRField adds 256
src0 = decode_src(raw, cdna=True) # handles SGPRs, constants, specials
src0_sel = _SDWA_SEL.get(inst.src0_sel, f'SEL{inst.src0_sel}')
mods = []
if inst.clmp: mods.append("clamp")
if inst.omod == 1: mods.append("mul:2")
elif inst.omod == 2: mods.append("mul:4")
elif inst.omod == 3: mods.append("div:2")
mods.append(f"src0_sel:{src0_sel}")
return f"{name}_sdwa {dst}, {src0} {' '.join(mods)}"
def _decode_dpp(dpp: int) -> str:
"""Decode DPP control value to string."""
op, arg = decode_dpp16(dpp)
if op == "quad_perm": return f"quad_perm:[{','.join(str(x) for x in arg)}]"
if op in ("row_mirror", "row_half_mirror"): return op
if op == "dpp": return f"dpp:{arg:#x}"
return f"{op}:{arg}"
def _disasm_vop1_dpp(inst) -> str:
name = inst.op_name.lower().replace('_e32', '')
regs = inst.canonical_op_regs
dst, src0 = _vreg(inst.vdst, regs['d']), _vreg(inst.vsrc0, regs['s0'])
dpp_str = _decode_dpp(inst.dpp)
mods = [dpp_str]
if inst.row_mask != 0xf: mods.append(f"row_mask:{inst.row_mask:#x}")
if inst.bank_mask != 0xf: mods.append(f"bank_mask:{inst.bank_mask:#x}")
if inst.bc: mods.append("bound_ctrl:1")
return f"{name}_dpp {dst}, {src0} {' '.join(mods)}"
def _disasm_vop2_sdwa(inst) -> str:
name, cdna = inst.op_name.lower().replace('_e32', ''), _is_cdna(inst)
regs = inst.canonical_op_regs
dst = _vreg(inst.vdst, regs['d'])
# When s0/s1=1, vsrc is SGPR/constant (VGPRField adds 256, so subtract it back)
src0 = _vreg(inst.vsrc0, regs['s0']) if inst.s0 == 0 else decode_src(_unwrap(inst.vsrc0) - 256, cdna)
src1 = _vreg(inst.vsrc1, regs['s1']) if inst.s1 == 0 else decode_src(_unwrap(inst.vsrc1) - 256, cdna)
src0_sel = _SDWA_SEL.get(inst.src0_sel, f'SEL{inst.src0_sel}')
src1_sel = _SDWA_SEL.get(inst.src1_sel, f'SEL{inst.src1_sel}')
mods = []
if inst.clmp: mods.append("clamp")
if inst.omod == 1: mods.append("mul:2")
elif inst.omod == 2: mods.append("mul:4")
elif inst.omod == 3: mods.append("div:2")
if inst.src0_sel != 6: mods.append(f"src0_sel:{src0_sel}")
if inst.src1_sel != 6: mods.append(f"src1_sel:{src1_sel}")
mods_str = ' '.join(mods) if mods else ""
# CDNA carry instructions and cndmask need vcc operands
if cdna and name in _VOP2_CARRY_OUT: return f"{name}_sdwa {dst}, vcc, {src0}, {src1} {mods_str}".strip()
if cdna and name in _VOP2_CARRY_INOUT: return f"{name}_sdwa {dst}, vcc, {src0}, {src1}, vcc {mods_str}".strip()
if cdna and name == 'v_cndmask_b32': return f"{name}_sdwa {dst}, {src0}, {src1}, vcc {mods_str}".strip()
return f"{name}_sdwa {dst}, {src0}, {src1} {mods_str}".strip()
def _disasm_vop2_dpp(inst) -> str:
name, cdna = inst.op_name.lower().replace('_e32', ''), _is_cdna(inst)
regs = inst.canonical_op_regs
dst, src0, src1 = _vreg(inst.vdst, regs['d']), _vreg(inst.vsrc0, regs['s0']), _vreg(inst.vsrc1, regs['s1'])
dpp_str = _decode_dpp(inst.dpp)
mods = [dpp_str]
if inst.row_mask != 0xf: mods.append(f"row_mask:{inst.row_mask:#x}")
if inst.bank_mask != 0xf: mods.append(f"bank_mask:{inst.bank_mask:#x}")
if inst.bc: mods.append("bound_ctrl:1")
# CDNA carry instructions and cndmask need vcc operands
if cdna and name in _VOP2_CARRY_OUT: return f"{name}_dpp {dst}, vcc, {src0}, {src1} {' '.join(mods)}"
if cdna and name in _VOP2_CARRY_INOUT: return f"{name}_dpp {dst}, vcc, {src0}, {src1}, vcc {' '.join(mods)}"
if cdna and name == 'v_cndmask_b32': return f"{name}_dpp {dst}, {src0}, {src1}, vcc {' '.join(mods)}"
return f"{name}_dpp {dst}, {src0}, {src1} {' '.join(mods)}"
def _disasm_vopc_sdwa(inst) -> str:
name = inst.op_name.lower().replace('_e32', '')
regs = inst.canonical_op_regs
sdst = _fmt_sdst(inst.sdst, 2, cdna=True)
src0 = _vreg(inst.vsrc0, regs['s0']) if getattr(inst, 's0', 0) == 0 else decode_src(_unwrap(inst.vsrc0) - 256, cdna=True)
src1 = _vreg(inst.vsrc1, regs['s1']) if getattr(inst, 's1', 0) == 0 else decode_src(_unwrap(inst.vsrc1) - 256, cdna=True)
src0_sel = _SDWA_SEL.get(inst.src0_sel, f'SEL{inst.src0_sel}')
src1_sel = _SDWA_SEL.get(inst.src1_sel, f'SEL{inst.src1_sel}')
mods = []
if inst.src0_sel != 6: mods.append(f"src0_sel:{src0_sel}")
if inst.src1_sel != 6: mods.append(f"src1_sel:{src1_sel}")
return f"{name}_sdwa {sdst}, {src0}, {src1} {' '.join(mods)}".strip()
def _disasm_vop3px2(inst) -> str:
"""VOP3PX2 disassembler for scaled MFMA instructions."""
name = inst.op_name.lower()
regs = inst.canonical_op_regs
dregs, r2 = regs['d'], regs['s2']
# 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
vgprs = {0: 8, 1: 8, 2: 6, 3: 6, 4: 4}
r0, r1 = vgprs.get(inst.cbsz, 8), vgprs.get(inst.blgp, 8)
def mfma_src(v, sc, is_acc):
v = _unwrap(v)
if v == 255: return _lit(inst, v)
base = v - 256 if v >= 256 else v
return _areg(base, sc) if is_acc else _vreg(base, sc)
src0, src1, src2 = mfma_src(inst.src0, r0, False), mfma_src(inst.src1, r1, inst.acc & 2), mfma_src(inst.src2, r2, inst.acc_cd)
dst = _areg(inst.vdst, dregs) if inst.acc_cd else _vreg(inst.vdst, dregs)
scale_src0, scale_src1 = _vreg(inst.scale_src0), _vreg(inst.scale_src1)
mods = []
if inst.opsel: mods.append(_fmt_bits("op_sel", inst.opsel, 3))
if inst.opsel_hi != 0: mods.append(_fmt_bits("op_sel_hi", inst.opsel_hi, 3))
if inst.neg: mods.append(_fmt_bits("neg", inst.neg, 3))
if inst.cbsz: mods.append(f"cbsz:{inst.cbsz}")
if inst.blgp: mods.append(f"blgp:{inst.blgp}")
return f"{name} {dst}, {src0}, {src1}, {src2}, {scale_src0}, {scale_src1}{' ' + ' '.join(mods) if mods else ''}"
DISASM_HANDLERS.update({CDNA_VOP1: _disasm_vop1, CDNA_VOP1_LIT: _disasm_vop1,
CDNA_VOP1_SDWA: _disasm_vop1_sdwa, CDNA_VOP1_DPP16: _disasm_vop1_dpp,
CDNA_VOP2: _disasm_vop2, CDNA_VOP2_LIT: _disasm_vop2,
CDNA_VOP2_SDWA: _disasm_vop2_sdwa, CDNA_VOP2_DPP16: _disasm_vop2_dpp,
CDNA_VOPC: _disasm_vopc, CDNA_VOPC_LIT: _disasm_vopc, CDNA_VOPC_SDWA_SDST: _disasm_vopc_sdwa,
CDNA_SOP1: _disasm_sop1, CDNA_SOP1_LIT: _disasm_sop1, CDNA_SOP2: _disasm_sop2, CDNA_SOP2_LIT: _disasm_sop2,
CDNA_SOPC: _disasm_sopc, CDNA_SOPC_LIT: _disasm_sopc, CDNA_SOPK: _disasm_sopk, CDNA_SOPK_LIT: _disasm_sopk, CDNA_SOPP: _disasm_sopp,
CDNA_SMEM: _disasm_smem, CDNA_DS: _disasm_ds, CDNA_FLAT: _disasm_flat, CDNA_GLOBAL: _disasm_flat, CDNA_SCRATCH: _disasm_flat,
CDNA_VOP3: _disasm_vop3a, CDNA_VOP3_SDST: _disasm_vop3b, CDNA_VOP3SD: _disasm_vop3b,
CDNA_VOP3P: _disasm_cdna_vop3p, CDNA_VOP3P_MFMA: _disasm_cdna_vop3p,
CDNA_MUBUF: _disasm_mubuf, CDNA_VOP3PX2: _disasm_vop3px2})

132
tinygrad_repo/test/amd/helpers.py Normal file
View File

@@ -0,0 +1,132 @@
"""Shared test helpers for AMD tests."""
import ctypes
from tinygrad.helpers import unwrap
from tinygrad.runtime.autogen import llvm
from tinygrad.runtime.support.elf import elf_loader
ARCH_TO_TARGET:dict[str, list[str]] = {
"rdna3":["gfx1100", "gfx1151"],
"rdna4":["gfx1200", "gfx1201"],
"cdna":["gfx950", "gfx942"],
}
TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items() for t in targets}
_DPP16_RANGE_OPS = {0x100: "row_shl", 0x110: "row_shr", 0x120: "row_ror", 0x150: "row_newbcast", 0x160: "row_share", 0x170: "row_xmask"}
_DPP16_EXACT_OPS = {0x130: ("wave_shl", 1), 0x134: ("wave_rol", 1), 0x138: ("wave_shr", 1), 0x13c: ("wave_ror", 1),
0x140: ("row_mirror", 0), 0x141: ("row_half_mirror", 0), 0x142: ("row_bcast", 15), 0x143: ("row_bcast", 31)}
def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
def decode_dpp16(dpp: int) -> tuple[str, int | tuple[int, int, int, int]]:
"""Decode a DPP16 control word into a symbolic operation and argument."""
if dpp < 0x100: return "quad_perm", ((dpp >> 0) & 0x3, (dpp >> 2) & 0x3, (dpp >> 4) & 0x3, (dpp >> 6) & 0x3)
if dpp in _DPP16_EXACT_OPS: return _DPP16_EXACT_OPS[dpp]
if (base := dpp & 0x1f0) in _DPP16_RANGE_OPS: return _DPP16_RANGE_OPS[base], dpp & 0xf
return "dpp", dpp
def get_mattr(arch:str) -> str:
return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
# LLVM in-process assembler/disassembler (replaces llvm-mc and llvm-objdump subprocesses)
_SENTINEL = b'\xde\xad\xbe\xef'
_SENTINEL_ASM = '.byte 0xde, 0xad, 0xbe, 0xef'
def _cerr(): return ctypes.pointer(ctypes.pointer(ctypes.c_char()))
def _expect(x, err, ret=None):
if x: raise RuntimeError(unwrap(ctypes.cast(err.contents, ctypes.c_char_p).value).decode() if not isinstance(err, str) else err)
return ret
def _init_llvm():
for component in ['Target', 'TargetInfo', 'TargetMC', 'AsmParser', 'AsmPrinter', 'Disassembler']:
getattr(llvm, f'LLVMInitializeAMDGPU{component}')()
def _create_target_machine(mcpu:str, mattr:str) -> llvm.LLVMTargetMachineRef:
target = _expect(llvm.LLVMGetTargetFromTriple(b'amdgcn-amd-amdhsa', ctypes.pointer(tgt:=llvm.LLVMTargetRef()), err:=_cerr()), err, tgt)
return llvm.LLVMCreateTargetMachine(target, b'amdgcn-amd-amdhsa', mcpu.encode(), mattr.encode(),
llvm.LLVMCodeGenLevelDefault, llvm.LLVMRelocDefault, llvm.LLVMCodeModelDefault)
def _emit_obj(asm_text:str, mcpu:str, mattr:str, diag_errors:list[str]|None=None) -> bytes:
"""Assemble raw asm text into an ELF object using LLVM in-process."""
_init_llvm()
tm = _create_target_machine(mcpu, mattr)
ctx = llvm.LLVMContextCreate()
try:
errors = diag_errors if diag_errors is not None else []
@llvm.LLVMDiagnosticHandler
def handle_diag(diag_ref, _arg):
if llvm.LLVMGetDiagInfoSeverity(diag_ref) == llvm.LLVMDSError:
errors.append(ctypes.string_at(llvm.LLVMGetDiagInfoDescription(diag_ref)).decode())
llvm.LLVMContextSetDiagnosticHandler(ctx, handle_diag, None)
mod = llvm.LLVMModuleCreateWithNameInContext(b'asm', ctx)
llvm.LLVMSetTarget(mod, b'amdgcn-amd-amdhsa')
asm_bytes = asm_text.encode()
llvm.LLVMSetModuleInlineAsm2(mod, asm_bytes, len(asm_bytes))
buf = llvm.LLVMMemoryBufferRef()
_expect(llvm.LLVMTargetMachineEmitToMemoryBuffer(tm, mod, llvm.LLVMObjectFile, err:=_cerr(), ctypes.pointer(buf)), err)
obj = ctypes.string_at(llvm.LLVMGetBufferStart(buf), llvm.LLVMGetBufferSize(buf))
llvm.LLVMDisposeMemoryBuffer(buf)
llvm.LLVMDisposeModule(mod)
return obj
finally:
llvm.LLVMContextDispose(ctx)
llvm.LLVMDisposeTargetMachine(tm)
def _extract_text(obj:bytes) -> bytes:
"""Extract .text section from ELF object bytes."""
return next(s.content for s in elf_loader(obj)[1] if s.name == ".text")
def llvm_assemble(instrs:list[str], mcpu:str, mattr:str) -> list[bytes]:
"""Assemble instructions in one LLVM emission, return per-instruction bytes."""
if not instrs: return []
parts = []
for instr in instrs:
parts.append(instr)
parts.append(_SENTINEL_ASM)
text = _extract_text(_emit_obj('.text\n' + '\n'.join(parts) + '\n', mcpu, mattr))
results, start = [], 0
for _ in instrs:
idx = text.find(_SENTINEL, start)
assert idx != -1, "sentinel not found in .text section"
results.append(bytes(text[start:idx]))
start = idx + len(_SENTINEL)
return results
def llvm_disasm(code:bytes, mcpu:str, mattr:str) -> list[str]:
"""Disassemble raw bytes into instruction strings using LLVM."""
_init_llvm()
dc = llvm.LLVMCreateDisasmCPUFeatures(b'amdgcn-amd-amdhsa', mcpu.encode(), mattr.encode(), None, 0,
llvm.LLVMOpInfoCallback(0), llvm.LLVMSymbolLookupCallback(0))
if not dc: raise RuntimeError(f"failed to create disasm context for {mcpu}")
llvm.LLVMSetDisasmOptions(dc, 2 | 4) # PrintImmHex | AsmPrinterVariant
try:
buf = ctypes.create_string_buffer(256)
arr = (ctypes.c_uint8 * len(code)).from_buffer_copy(code)
results, offset = [], 0
while offset < len(code):
size = llvm.LLVMDisasmInstruction(dc, ctypes.cast(ctypes.addressof(arr) + offset, ctypes.POINTER(ctypes.c_uint8)),
len(code) - offset, 0, buf, 256)
if size == 0: break
results.append(buf.value.decode().strip())
offset += size
return results
finally:
llvm.LLVMDisasmDispose(dc)
def llvm_filter_valid_asm(tests:list[tuple[str, bytes]], mcpu:str, mattr:str) -> list[tuple[str, bytes]]:
"""Filter out tests where original ASM isn't valid on target, and where LLVM roundtrip doesn't match."""
if not tests: return []
# Assemble all instructions at once with sentinels and diagnostic handler to detect failures
parts, diag_errors = [], [] # type: ignore[var-annotated]
for asm, _ in tests:
parts.append(asm)
parts.append(_SENTINEL_ASM)
text = _extract_text(_emit_obj('.text\n' + '\n'.join(parts) + '\n', mcpu, mattr, diag_errors))
results, start = [], 0
for _ in tests:
idx = text.find(_SENTINEL, start)
assert idx != -1, "sentinel not found in .text section"
results.append(bytes(text[start:idx]))
start = idx + len(_SENTINEL)
# Invalid instructions produce 0 bytes; also filter where LLVM roundtrip doesn't match original
return [(asm, data) for (asm, data), chunk in zip(tests, results) if len(chunk) > 0 and chunk == data]

1
tinygrad_repo/test/amd/hw/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Hardware-validated emulator tests for RDNA3 instructions."""

287
tinygrad_repo/test/amd/hw/helpers.py Normal file
View File

@@ -0,0 +1,287 @@
"""Test infrastructure for hardware-validated RDNA3 emulator tests.
Uses run_asm() with memory output, so tests can run on both emulator and real hardware.
Set USE_HW=1 to run on both emulator and hardware, comparing results.
"""
import ctypes, math, os, struct
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from test.mockgpu.amd.emu import run_asm
from tinygrad.renderer.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
def _f32(i: int) -> float: return struct.unpack('<f', struct.pack('<I', i & 0xFFFFFFFF))[0]
# f16 conversion helpers
def f16(i: int) -> float: return struct.unpack('<e', struct.pack('<H', i & 0xFFFF))[0]
def f32_to_f16(f: float) -> int:
f = float(f)
if math.isnan(f): return 0x7e00
if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
try: return struct.unpack('<H', struct.pack('<e', f))[0]
except OverflowError: return 0x7c00 if f > 0 else 0xfc00
# For backwards compatibility with tests using SrcEnum.NULL etc.
class SrcEnum:
NULL = NULL
VCC_LO = VCC_LO
VCC_HI = VCC_HI
EXEC_LO = EXEC_LO
EXEC_HI = EXEC_HI
SCC = SCC
M0 = M0
POS_HALF = 0.5
NEG_HALF = -0.5
POS_ONE = 1.0
NEG_ONE = -1.0
POS_TWO = 2.0
NEG_TWO = -2.0
POS_FOUR = 4.0
NEG_FOUR = -4.0
VCC = VCC_LO # For VOP3SD sdst field (VCC_LO is exported from dsl)
USE_HW = os.environ.get("USE_HW", "0") == "1"
FLOAT_TOLERANCE = 1e-5
def get_gpu_target() -> tuple[int, int, int]:
"""Get the GPU target as (major, minor, stepping) tuple."""
if not USE_HW: return (0, 0, 0)
from tinygrad.device import Device
return Device["AMD"].target # type: ignore[attr-defined]
def skip_unless_gfx(min_major: int, min_minor: int = 0, reason: str = ""):
"""Skip test if GPU target is below the minimum required version."""
import unittest
def decorator(test_func):
if not USE_HW: return test_func
target = get_gpu_target()
if target[0] < min_major or (target[0] == min_major and target[1] < min_minor):
return unittest.skip(reason or f"requires gfx{min_major}{min_minor}0+")(test_func)
return test_func
return decorator
# Output buffer layout: vgpr[N_VGPRS][n_lanes], sgpr[N_SGPRS], vcc, scc, exec
N_VGPRS, N_SGPRS, WAVE_SIZE = 16, 16, 32
SGPR_BYTES = N_SGPRS * 4 # 16 regs * 4 bytes = 64
_VGPR_REGION = N_VGPRS * WAVE_SIZE * 4 # minimum vgpr region size (tests may use as scratch)
def _out_bytes(n_lanes: int) -> int: return max(N_VGPRS * n_lanes * 4, _VGPR_REGION) + SGPR_BYTES + 12
OUT_BYTES = _out_bytes(WAVE_SIZE) # default for single-wave (backward compat)
# Float conversion helpers
def f2i(f: float) -> int: return _i32(f)
def i2f(i: int) -> float: return _f32(i)
def f2i64(f: float) -> int: return struct.unpack('<Q', struct.pack('<d', f))[0]
def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
def assemble(instructions: list) -> bytes:
return b''.join(inst.to_bytes() for inst in instructions)
# Simple WaveState class for test output parsing (mirrors test/mockgpu/amd/emu.py interface for tests)
class WaveState:
def __init__(self, n_lanes: int = 32):
self.vgpr = [[0] * 256 for _ in range(n_lanes)] # vgpr[lane][reg]
self.sgpr = [0] * 128
self.vcc = 0
self.scc = 0
def get_prologue_epilogue(n_lanes: int) -> tuple[list, list]:
"""Generate prologue and epilogue instructions for state capture."""
prologue = [
s_mov_b32(s[80], s[0]),
s_mov_b32(s[81], s[1]),
v_mov_b32_e32(v[255], v[0]),
]
for i in range(N_VGPRS):
prologue.append(v_mov_b32_e32(v[i], 0))
for i in range(N_SGPRS):
prologue.append(s_mov_b32(s[i], 0))
prologue.append(s_mov_b32(VCC_LO, 0))
epilogue = [
s_mov_b32(s[90], VCC_LO),
s_cselect_b32(s[91], 1, 0),
# Save EXEC early (before we modify it for VGPR stores)
s_mov_b32(s[95], EXEC_LO),
# Restore EXEC to all active lanes for VGPR stores (test may have modified EXEC)
s_mov_b32(EXEC_LO, (1 << min(n_lanes, WAVE_SIZE)) - 1),
s_load_b64(s[92:93], s[80:81], 0, soffset=NULL),
s_waitcnt(0), # simm16=0 waits for all
v_lshlrev_b32_e32(v[240], 2, v[255]),
]
vgpr_bytes = N_VGPRS * n_lanes * 4
for i in range(N_VGPRS):
epilogue.append(global_store_b32(addr=v[240], data=v[i], saddr=s[92:93], offset=i * n_lanes * 4))
epilogue.append(v_mov_b32_e32(v[241], 0))
epilogue.append(v_cmp_eq_u32_e32(v[255], v[241]))
epilogue.append(s_and_saveexec_b32(s[94], VCC_LO))
# Scalar stores: only thread 0. Use v[240]=vgpr_bytes as base offset so immediate offsets stay small.
epilogue.append(v_mov_b32_e32(v[240], vgpr_bytes))
for i in range(N_SGPRS):
epilogue.append(v_mov_b32_e32(v[243], s[i]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=i * 4))
epilogue.append(v_mov_b32_e32(v[243], s[90]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES))
epilogue.append(v_mov_b32_e32(v[243], s[91]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES + 4))
# Store EXEC (saved earlier in s[95])
epilogue.append(v_mov_b32_e32(v[243], s[95]))
epilogue.append(global_store_b32(addr=v[240], data=v[243], saddr=s[92:93], offset=SGPR_BYTES + 8))
epilogue.append(s_mov_b32(EXEC_LO, s[94]))
epilogue.append(s_endpgm())
return prologue, epilogue
def parse_output(out_buf: bytes, n_lanes: int) -> WaveState:
"""Parse output buffer into WaveState."""
vgpr_bytes = N_VGPRS * n_lanes * 4
st = WaveState(n_lanes)
for i in range(N_VGPRS):
for lane in range(n_lanes):
off = i * n_lanes * 4 + lane * 4
st.vgpr[lane][i] = struct.unpack_from('<I', out_buf, off)[0]
for i in range(N_SGPRS):
st.sgpr[i] = struct.unpack_from('<I', out_buf, vgpr_bytes + i * 4)[0]
st.vcc = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES)[0]
st.scc = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES + 4)[0]
# Store EXEC in its proper location (index 126)
st.sgpr[EXEC_LO.offset] = struct.unpack_from('<I', out_buf, vgpr_bytes + SGPR_BYTES + 8)[0]
return st
def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions via emulator run_asm, dump state to memory, return WaveState."""
buf_sz = _out_bytes(n_lanes)
out_buf = (ctypes.c_uint8 * buf_sz)(*([0] * buf_sz))
out_addr = ctypes.addressof(out_buf)
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
args = (ctypes.c_uint64 * 1)(out_addr)
args_ptr = ctypes.addressof(args)
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
lib_ptr = ctypes.addressof(kernel_buf)
# rsrc2: USER_SGPR_COUNT=2, ENABLE_SGPR_WORKGROUP_ID_X/Y/Z=1, LDS_SIZE=128 (64KB)
rsrc2 = 0x19c | (128 << 15)
scratch_size = 0x10000 # 64KB per lane, matches .amdhsa_private_segment_fixed_size in run_program_hw
result = run_asm(lib_ptr, len(code), 1, 1, 1, n_lanes, 1, 1, args_ptr, rsrc2, scratch_size)
assert result == 0, f"run_asm failed with {result}"
return parse_output(bytes(out_buf), n_lanes)
def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
"""Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram
from tinygrad.runtime.support.compiler_amd import HIPCompiler
from tinygrad.helpers import flat_mv
dev = Device["AMD"]
compiler = HIPCompiler(dev.arch) # type: ignore[attr-defined]
prologue, epilogue = get_prologue_epilogue(n_lanes)
code = assemble(prologue + instructions + epilogue)
byte_str = ', '.join(f'0x{b:02x}' for b in code)
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 96
.amdhsa_wavefront_size32 1
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_kernarg_size 8
.amdhsa_group_segment_fixed_size 65536
.amdhsa_private_segment_fixed_size 65536
.amdhsa_enable_private_segment 1
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 8
.group_segment_fixed_size: 65536
.private_segment_fixed_size: 65536
.kernarg_segment_align: 8
.wavefront_size: 32
.sgpr_count: 96
.vgpr_count: 256
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
lib = compiler.compile(asm_src)
prg = AMDProgram(dev, "test", lib) # type: ignore[arg-type]
buf_sz = _out_bytes(n_lanes)
out_gpu = dev.allocator.alloc(buf_sz)
assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
out_buf = bytearray(buf_sz)
dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
return parse_output(bytes(out_buf), n_lanes)
def compare_wave_states(emu_st: WaveState, hw_st: WaveState, n_lanes: int, n_vgprs: int = N_VGPRS, ulp_tolerance: int = 0) -> list[str]:
"""Compare two WaveStates and return list of differences.
Args:
ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
"""
import math
diffs = []
for i in range(n_vgprs):
for lane in range(n_lanes):
emu_val = emu_st.vgpr[lane][i]
hw_val = hw_st.vgpr[lane][i]
if emu_val != hw_val:
emu_f, hw_f = _f32(emu_val), _f32(hw_val)
if math.isnan(emu_f) and math.isnan(hw_f):
continue
# Check ULP difference for floats (only for same-sign values)
if ulp_tolerance > 0 and (emu_val < 0x80000000) == (hw_val < 0x80000000):
ulp_diff = abs(int(emu_val) - int(hw_val))
if ulp_diff <= ulp_tolerance:
continue
diffs.append(f"v[{i}] lane {lane}: emu=0x{emu_val:08x} ({emu_f:.6g}) hw=0x{hw_val:08x} ({hw_f:.6g})")
for i in range(N_SGPRS):
emu_val = emu_st.sgpr[i]
hw_val = hw_st.sgpr[i]
if emu_val != hw_val:
diffs.append(f"s[{i}]: emu=0x{emu_val:08x} hw=0x{hw_val:08x}")
if emu_st.vcc != hw_st.vcc:
diffs.append(f"vcc: emu=0x{emu_st.vcc:08x} hw=0x{hw_st.vcc:08x}")
if emu_st.scc != hw_st.scc:
diffs.append(f"scc: emu={emu_st.scc} hw={hw_st.scc}")
return diffs
def run_program(instructions: list, n_lanes: int = 1, ulp_tolerance: int = 0) -> WaveState:
"""Run instructions and return WaveState.
If USE_HW=1, runs on both emulator and hardware, compares results, and raises if they differ.
Otherwise, runs only on emulator.
Args:
ulp_tolerance: Allow up to this many ULPs difference for float comparisons (0 = exact match required)
"""
emu_st = run_program_emu(instructions, n_lanes)
if USE_HW:
hw_st = run_program_hw(instructions, n_lanes)
diffs = compare_wave_states(emu_st, hw_st, n_lanes, ulp_tolerance=ulp_tolerance)
if diffs:
raise AssertionError("Emulator vs Hardware mismatch:\n" + "\n".join(diffs))
return hw_st
return emu_st

20
tinygrad_repo/test/amd/hw/test_cdna_sdwa.py Normal file
View File

@@ -0,0 +1,20 @@
#!/usr/bin/env python3
import unittest
import tinygrad.runtime.autogen.amd.cdna.ins as cdna
from test.amd.hw.test_cdna_vop3 import run_cdna
class TestCDNASDWA(unittest.TestCase):
def test_v_add_co_u32_e32_writes_vcc(self):
out = run_cdna([
cdna.s_mov_b32(cdna.s[0], 0xffffffff),
cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
cdna.v_mov_b32_e32(cdna.v[13], 1),
cdna.v_add_co_u32_e32(cdna.v[0], cdna.SDWA, cdna.v[13], vsrc0=cdna.v[0], dst_sel=6, src0_sel=6),
cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO),
cdna.v_lshlrev_b32_e32(cdna.v[2], 31, cdna.v[2]),
cdna.v_or_b32_e32(cdna.v[2], cdna.v[2], cdna.v[0]),
])
self.assertEqual(out, 0x80000000)
if __name__ == "__main__":
unittest.main()

130
tinygrad_repo/test/amd/hw/test_cdna_vop3.py Normal file
View File

@@ -0,0 +1,130 @@
"""CDNA VOP3 instruction coverage.
Exercises generated CDNA pcode end-to-end in the emulator and compares against
gfx950 hardware when USE_HW=1.
"""
import ctypes, struct, unittest
import tinygrad.runtime.autogen.amd.cdna.ins as cdna
from tinygrad.helpers import flat_mv
from tinygrad.renderer.amd.dsl import NULL
from test.amd.hw.helpers import USE_HW, assemble
from test.mockgpu.amd.emu import run_asm
LANES = 1
def _code(instructions: list, out_reg: int = 2, out_addr: int | None = None) -> bytes:
load_out_addr = [
cdna.s_mov_b32(cdna.s[92], out_addr & 0xffffffff),
cdna.s_mov_b32(cdna.s[93], out_addr >> 32),
] if out_addr is not None else [
cdna.s_load_dwordx2(cdna.s[92:93], cdna.s[80:81], 0, soffset=NULL),
cdna.s_waitcnt(0),
]
return assemble([
cdna.s_mov_b32(cdna.s[80], cdna.s[0]),
cdna.s_mov_b32(cdna.s[81], cdna.s[1]),
cdna.v_mov_b32_e32(cdna.v[255], cdna.v[0]),
*instructions,
*load_out_addr,
cdna.v_lshlrev_b32_e32(cdna.v[240], 2, cdna.v[255]),
cdna.global_store_dword(addr=cdna.v[240], data=cdna.v[out_reg], saddr=cdna.s[92:93], offset=0),
cdna.s_endpgm(),
])
def _run_emu(instructions: list, out_reg: int = 2) -> int:
out_buf = (ctypes.c_uint32 * LANES)(*([0] * LANES))
args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
code = _code(instructions, out_reg)
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, LANES, 1, 1, ctypes.addressof(args),
0x19c | (128 << 15), 0x10000, arch="cdna")
assert result == 0, f"run_asm failed with {result}"
return out_buf[0]
def _run_hw(instructions: list, out_reg: int = 2) -> int:
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram
from tinygrad.runtime.support.compiler_amd import HIPCompiler
dev = Device["AMD"]
if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware")
out_gpu = dev.allocator.alloc(LANES * 4)
code = _code(instructions, out_reg, out_gpu.va_addr)
byte_str = ", ".join(f"0x{b:02x}" for b in code)
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 96
.amdhsa_accum_offset 256
.amdhsa_kernarg_size 0
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 0
.group_segment_fixed_size: 0
.private_segment_fixed_size: 0
.kernarg_segment_align: 8
.wavefront_size: 64
.sgpr_count: 96
.vgpr_count: 256
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src))
prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
out = bytearray(LANES * 4)
dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
return struct.unpack("<I", out)[0]
def run_cdna(instructions: list, out_reg: int = 2) -> int:
emu = _run_emu(instructions, out_reg)
if not USE_HW: return emu
hw = _run_hw(instructions, out_reg)
if emu != hw: raise AssertionError(f"Emulator vs Hardware mismatch: emu=0x{emu:08x} hw=0x{hw:08x}")
return hw
class TestCDNAVOP3(unittest.TestCase):
def test_cvt_pk_fp8_f32_preserves_upper_half(self):
"""V_CVT_PK_FP8_F32 with OPSEL[3]=0 writes only D[15:0]."""
out = run_cdna([
cdna.s_mov_b32(cdna.s[0], 0xdeadbeef),
cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]),
cdna.v_mov_b32_e32(cdna.v[0], 1.0),
cdna.v_mov_b32_e32(cdna.v[1], 2.0),
cdna.v_cvt_pk_fp8_f32(cdna.v[2], cdna.v[0], cdna.v[1]),
])
self.assertEqual(out, 0xdead4038)
def test_cvt_pk_bf8_f32_overflow_and_inf(self):
"""V_CVT_PK_BF8_F32 converts finite overflow and infinities to E5M2 infinities."""
for name, bits, expected in [
("finite_overflow", 0x47700000, 0x7c),
("pos_inf", 0x7f800000, 0x7c),
("neg_inf", 0xff800000, 0xfc),
]:
with self.subTest(name=name):
out = run_cdna([
cdna.s_mov_b32(cdna.s[0], 0xdeadbeef),
cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]),
cdna.s_mov_b32(cdna.s[0], bits),
cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
cdna.v_mov_b32_e32(cdna.v[1], 1.0),
cdna.v_cvt_pk_bf8_f32(cdna.v[2], cdna.v[0], cdna.v[1]),
])
self.assertEqual(out, 0xdead3c00 | expected)

187
tinygrad_repo/test/amd/hw/test_dpp.py Normal file
View File

@@ -0,0 +1,187 @@
"""Tests for DPP16 source swizzles.
These instructions trap in the default wave32 hw helper, so this file uses a
minimal wave64 lane-store harness and compares emulator vs hardware directly
when USE_HW=1.
"""
import ctypes, unittest
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from tinygrad.helpers import flat_mv
from test.amd.hw.helpers import USE_HW, assemble
from test.mockgpu.amd.emu import run_asm
WAVE64 = 64
def _wave64_code(instructions: list, out_reg: int = 1) -> bytes:
return assemble([
s_mov_b32(s[80], s[0]),
s_mov_b32(s[81], s[1]),
v_mov_b32_e32(v[255], v[0]),
*instructions,
s_load_b64(s[92:93], s[80:81], 0, soffset=NULL),
s_waitcnt(0),
v_lshlrev_b32_e32(v[240], 2, v[255]),
global_store_b32(addr=v[240], data=v[out_reg], saddr=s[92:93], offset=0),
s_endpgm(),
])
def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
out_buf = (ctypes.c_uint32 * WAVE64)(*([0] * WAVE64))
args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
code = _wave64_code(instructions, out_reg)
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
rsrc2 = 0x19c | (128 << 15)
scratch_size = 0x10000
result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, WAVE64, 1, 1, ctypes.addressof(args), rsrc2, scratch_size)
assert result == 0, f"run_asm failed with {result}"
return list(out_buf)
def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram
from tinygrad.runtime.support.compiler_amd import HIPCompiler
dev = Device["AMD"]
compiler = HIPCompiler(dev.arch) # type: ignore[attr-defined]
code = _wave64_code(instructions, out_reg)
byte_str = ', '.join(f'0x{b:02x}' for b in code)
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 96
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_kernarg_size 8
.amdhsa_group_segment_fixed_size 65536
.amdhsa_private_segment_fixed_size 65536
.amdhsa_enable_private_segment 1
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 8
.group_segment_fixed_size: 65536
.private_segment_fixed_size: 65536
.kernarg_segment_align: 8
.wavefront_size: 64
.sgpr_count: 96
.vgpr_count: 256
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
lib = compiler.compile(asm_src)
prg = AMDProgram(dev, "test", lib) # type: ignore[arg-type]
out_gpu = dev.allocator.alloc(WAVE64 * 4)
prg(out_gpu, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
out = bytearray(WAVE64 * 4)
dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(WAVE64)]
def run_wave64(instructions: list, out_reg: int = 1) -> list[int]:
emu = _run_wave64_emu(instructions, out_reg)
if not USE_HW: return emu
hw = _run_wave64_hw(instructions, out_reg)
if emu != hw:
diffs = [f"lane {i}: emu=0x{e:08x} hw=0x{h:08x}" for i, (e, h) in enumerate(zip(emu, hw)) if e != h]
raise AssertionError("Emulator vs Hardware mismatch:\n" + '\n'.join(diffs[:16]))
return hw
class TestDPP16(unittest.TestCase):
def _run_copy(self, dpp: int, *, row_mask: int = 0xf, bank_mask: int = 0xf, bc: int = 1, dst_seed: int | None = None) -> list[int]:
instructions = [
v_mul_u32_u24_e32(v[0], 10, v[255]),
v_add_nc_u32_e32(v[0], 3, v[0]),
]
if dst_seed is not None: instructions.append(v_mov_b32_e32(v[1], dst_seed))
instructions += [v_mov_b32_e32(v[2], 0), v_or_b32_e32(v[1], DPP, v[2], vsrc0=v[0], dpp=dpp, row_mask=row_mask, bank_mask=bank_mask, bc=bc)]
return run_wave64(instructions)
def test_quad_perm_reverse(self):
out = self._run_copy(0x1b)
self.assertEqual(out[0], 33)
self.assertEqual(out[1], 23)
self.assertEqual(out[2], 13)
self.assertEqual(out[3], 3)
self.assertEqual(out[4], 73)
def test_row_shl(self):
out = self._run_copy(0x101)
self.assertEqual(out[0], 13)
self.assertEqual(out[7], 83)
self.assertEqual(out[14], 153)
self.assertEqual(out[15], 0)
self.assertEqual(out[16], 173)
def test_row_shr(self):
out = self._run_copy(0x111)
self.assertEqual(out[0], 0)
self.assertEqual(out[1], 3)
self.assertEqual(out[8], 73)
self.assertEqual(out[15], 143)
self.assertEqual(out[16], 0)
self.assertEqual(out[17], 163)
def test_row_ror(self):
out = self._run_copy(0x121)
self.assertEqual(out[0], 153)
self.assertEqual(out[1], 3)
self.assertEqual(out[15], 143)
self.assertEqual(out[16], 313)
def test_row_mirror(self):
out = self._run_copy(0x140)
self.assertEqual(out[0], 153)
self.assertEqual(out[5], 103)
self.assertEqual(out[8], 73)
self.assertEqual(out[16], 313)
def test_row_half_mirror(self):
out = self._run_copy(0x141)
self.assertEqual(out[0], 73)
self.assertEqual(out[7], 3)
self.assertEqual(out[8], 153)
self.assertEqual(out[15], 83)
self.assertEqual(out[16], 233)
def test_row_mask(self):
out = self._run_copy(0x101, row_mask=0x5, dst_seed=0xDEADBEEF)
self.assertEqual(out[0], 13)
self.assertEqual(out[15], 0)
self.assertEqual(out[16], 0xDEADBEEF)
self.assertEqual(out[32], 333)
self.assertEqual(out[47], 0)
self.assertEqual(out[48], 0xDEADBEEF)
def test_bank_mask(self):
out = self._run_copy(0x101, bank_mask=0x5, dst_seed=0xDEADBEEF)
self.assertEqual(out[0], 13)
self.assertEqual(out[3], 43)
self.assertEqual(out[4], 0xDEADBEEF)
self.assertEqual(out[8], 93)
self.assertEqual(out[12], 0xDEADBEEF)
class TestVOPCDPP16(unittest.TestCase):
def test_row_bcast15_materializes_vcc(self):
out = run_wave64([
v_mov_b32_e32(v[0], v[255]),
v_cmp_eq_u32_e32(DPP, v[0], vsrc0=v[0], dpp=0x142, row_mask=0xf, bank_mask=0xf, bc=1),
v_mov_b32_e32(v[2], 0),
v_mov_b32_e32(v[3], 1),
v_cndmask_b32_e32(v[1], v[2], v[3]),
])
for lane in (0, 16, 32, 48): self.assertEqual(out[lane], 1)
for lane in (1, 15, 31, 47, 63): self.assertEqual(out[lane], 0)

960
tinygrad_repo/test/amd/hw/test_ds.py Normal file
View File

@@ -0,0 +1,960 @@
"""Tests for DS instructions - data share (LDS) operations.
Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
ds_add_*, ds_max_*, ds_min_*, ds_and_*, ds_or_*, ds_xor_*,
ds_inc_*, ds_dec_*, ds_cmpstore_*, ds_storexchg_*
"""
import unittest
from test.amd.hw.helpers import *
class TestDS2Addr(unittest.TestCase):
"""Tests for DS_*_2ADDR instructions."""
def test_ds_store_load_2addr_b32(self):
"""DS_STORE_2ADDR_B32 and DS_LOAD_2ADDR_B32 with offset * 4."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
def test_ds_store_load_2addr_b64(self):
"""DS_STORE_2ADDR_B64 and DS_LOAD_2ADDR_B64."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_B64, addr=v[10], data0=v[0:1], data1=v[2:3], vdst=v[0], offset0=0, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], offset0=0, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
class TestDS2AddrMore(unittest.TestCase):
"""Additional DS_*_2ADDR tests."""
def test_ds_store_load_2addr_b32_nonzero_offsets(self):
"""DS_STORE_2ADDR_B32 with non-zero offsets (offset*4 scaling)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=2, offset1=5),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=2, offset1=5),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have value from offset 8 (2*4)")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have value from offset 20 (5*4)")
def test_ds_2addr_b64_no_overlap(self):
"""DS_LOAD_2ADDR_B64 with adjacent offsets should not overlap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x11111111),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0x22222222),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_mov_b32(s[2], 0x33333333),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
s_mov_b32(s[2], 0x44444444),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[10], vdst=v[4:7], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should be 0x11111111")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should be 0x22222222")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should be 0x33333333")
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should be 0x44444444")
def test_ds_load_2addr_b32_no_overwrite(self):
"""DS_LOAD_2ADDR_B32 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[2]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[4], s[2]), # Sentinel
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[2:3], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
def test_ds_load_2addr_b64_addr_overlaps_vdst(self):
"""DS_LOAD_2ADDR_B64 where addr register overlaps vdst range.
Hardware reads the address before writing any results, so addr=v[4]
with vdst=v[4:7] must load all 4 dwords using the original v[4] value.
"""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_mov_b32(s[2], 0xCCCCCCCC),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=8),
s_mov_b32(s[2], 0xDDDDDDDD),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# addr=v[4] overlaps vdst=v[4:7]
v_mov_b32_e32(v[4], 0),
DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[4], vdst=v[4:7], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 = LDS[0:4]")
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB, "v5 = LDS[4:8]")
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC, "v6 = LDS[8:12]")
self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD, "v7 = LDS[12:16]")
def test_ds_load_2addr_b32_addr_overlaps_vdst(self):
"""DS_LOAD_2ADDR_B32 where addr register overlaps vdst range."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_mov_b32(s[2], 0xBBBBBBBB),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=4),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# addr=v[2] overlaps vdst=v[2:3]
v_mov_b32_e32(v[2], 0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[2], vdst=v[2:3], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 = LDS[0:4]")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 = LDS[4:8]")
def test_ds_load_b64_no_overwrite(self):
"""DS_LOAD_B64 should only write 2 VGPRs."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[2]),
ds_store_b64(addr=v[10], data0=v[0:1], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0x12345678),
v_mov_b32_e32(v[4], s[2]), # Sentinel
ds_load_b64(addr=v[10], vdst=v[2:3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][3], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][4], 0x12345678, "v4 should be untouched")
class TestDSB96(unittest.TestCase):
"""Tests for DS_STORE_B96 and DS_LOAD_B96 (96-bit / 3 dwords)."""
def test_ds_store_load_b96(self):
"""DS_STORE_B96 stores 3 VGPRs, DS_LOAD_B96 loads them back."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
ds_store_b96(addr=v[10], data0=v[0:2]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b96(addr=v[10], vdst=v[4:6]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
def test_ds_store_b96_with_offset(self):
"""DS_STORE_B96 with non-zero offset."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
DS(DSOp.DS_STORE_B96, addr=v[10], data0=v[0:2], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B96, addr=v[10], vdst=v[4:6], offset0=12),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
class TestDSB128(unittest.TestCase):
"""Tests for DS_STORE_B128 and DS_LOAD_B128 (128-bit / 4 dwords)."""
def test_ds_store_load_b128(self):
"""DS_STORE_B128 stores 4 VGPRs, DS_LOAD_B128 loads them back."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x44444444),
v_mov_b32_e32(v[3], s[0]),
ds_store_b128(addr=v[10], data0=v[0:3]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b128(addr=v[10], vdst=v[4:7]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have first dword")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have second dword")
self.assertEqual(st.vgpr[0][6], 0x33333333, "v6 should have third dword")
self.assertEqual(st.vgpr[0][7], 0x44444444, "v7 should have fourth dword")
def test_ds_store_b128_with_offset(self):
"""DS_STORE_B128 with non-zero offset."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xDDDDDDDD),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_B128, addr=v[10], data0=v[0:3], offset0=16),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B128, addr=v[10], vdst=v[4:7], offset0=16),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC)
self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD)
class TestDSAtomic(unittest.TestCase):
"""Tests for DS atomic operations."""
def test_ds_max_rtn_u32(self):
"""DS_MAX_RTN_U32: atomically store max and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_min_rtn_u32(self):
"""DS_MIN_RTN_U32: atomically store min and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[1], s[2]),
ds_min_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 200)
self.assertEqual(st.vgpr[0][3], 100)
def test_ds_and_rtn_b32(self):
"""DS_AND_RTN_B32: atomically AND and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xFF00FF00),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0xFFFF0000),
v_mov_b32_e32(v[1], s[2]),
ds_and_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFF00FF00)
self.assertEqual(st.vgpr[0][3], 0xFF000000)
def test_ds_or_rtn_b32(self):
"""DS_OR_RTN_B32: atomically OR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0x00FF0000),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0x000000FF),
v_mov_b32_e32(v[1], s[2]),
ds_or_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x00FF0000)
self.assertEqual(st.vgpr[0][3], 0x00FF00FF)
def test_ds_xor_rtn_b32(self):
"""DS_XOR_RTN_B32: atomically XOR and return old value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0xFFFFFFFF),
v_mov_b32_e32(v[1], s[2]),
ds_xor_rtn_b32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0x55555555)
def test_ds_inc_rtn_u32(self):
"""DS_INC_RTN_U32: increment with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_inc_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 6)
def test_ds_dec_rtn_u32(self):
"""DS_DEC_RTN_U32: decrement with wrap."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 5),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vgpr[0][3], 4)
def test_ds_cmpstore_b32_match(self):
"""DS_CMPSTORE_B32: conditional store when compare matches."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[2], s[2]), # compare = 100 (matches)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 200)
def test_ds_cmpstore_b32_no_match(self):
"""DS_CMPSTORE_B32: no store when compare doesn't match."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]), # new value
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[2], s[2]), # compare = 50 (doesn't match)
ds_cmpstore_b32(addr=v[10], data0=v[1], data1=v[2], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[4], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 100)
def test_ds_max_u32_no_rtn(self):
"""DS_MAX_U32 (no RTN): atomically store max, no return value."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 200),
v_mov_b32_e32(v[1], s[2]),
ds_max_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][3], 200, "v3 should have max(100, 200) = 200")
def test_ds_add_u32_no_rtn_preserves_vdst(self):
"""DS_ADD_U32 (no RTN) should NOT write to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF, "v2 should preserve sentinel")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_add_rtn_u32_writes_vdst(self):
"""DS_ADD_RTN_U32 should write old value to vdst."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[2]), # sentinel
s_mov_b32(s[2], 100),
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 50),
v_mov_b32_e32(v[1], s[2]),
ds_add_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "v2 should have old value (100)")
self.assertEqual(st.vgpr[0][3], 150, "v3 should have 100 + 50 = 150")
def test_ds_dec_rtn_u32_wrap(self):
"""DS_DEC_RTN_U32: decrement wraps when value is 0 or > limit."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[2], 0), # Start at 0
v_mov_b32_e32(v[0], s[2]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 10), # limit
v_mov_b32_e32(v[1], s[2]),
ds_dec_rtn_u32(addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0, "v2 should have old value (0)")
# When mem == 0 or mem > limit, result = limit
self.assertEqual(st.vgpr[0][3], 10, "v3 should wrap to limit (10)")
class TestDSStorexchg(unittest.TestCase):
"""Tests for DS_STOREXCHG instructions."""
def test_ds_storexchg_rtn_b32(self):
"""DS_STOREXCHG_RTN_B32: exchange value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STOREXCHG_RTN_B32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
class TestDSRegisterWidth(unittest.TestCase):
"""Regression tests: DS loads should only write correct number of VGPRs."""
def test_ds_load_b32_no_overwrite(self):
"""DS_LOAD_B32 should only write 1 VGPR."""
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]), # sentinel
ds_store_b32(addr=v[0], data0=v[1], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[0], vdst=v[1], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
class TestDS2AddrStride64(unittest.TestCase):
"""Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
def test_ds_store_load_2addr_stride64_b32(self):
"""DS_STORE_2ADDR_STRIDE64_B32: stores at ADDR + offset*256."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[2:3], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 from addr 256")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 from addr 512")
def test_ds_store_load_2addr_stride64_b64(self):
"""DS_STORE_2ADDR_STRIDE64_B64: stores at ADDR + offset*512."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[2:3], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[4:7], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][6], 0x12345678)
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0)
def test_ds_storexchg_2addr_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_RTN_B32: exchange at two addresses."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4:5], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[10], vdst=v[6:7], offset0=0, offset1=1),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "old val 0")
self.assertEqual(st.vgpr[0][5], 0x22222222, "old val 1")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "new val 0")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "new val 1")
def test_ds_storexchg_rtn_b64(self):
"""DS_STOREXCHG_RTN_B64: exchange 64-bit value and return old."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]), # initial low
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[1], s[0]), # initial high
DS(DSOp.DS_STORE_B64, addr=v[10], data0=v[0:1], vdst=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[2], s[0]), # new low
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[3], s[0]), # new high
DS(DSOp.DS_STOREXCHG_RTN_B64, addr=v[10], data0=v[2:3], vdst=v[4:5], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B64, addr=v[10], vdst=v[6:7], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should have old low dword")
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE, "v5 should have old high dword")
self.assertEqual(st.vgpr[0][6], 0x12345678, "v6 should have new low dword")
self.assertEqual(st.vgpr[0][7], 0x9ABCDEF0, "v7 should have new high dword")
def test_ds_store_load_2addr_stride64_b64_roundtrip(self):
"""DS_STORE_2ADDR_STRIDE64_B64 followed by DS_LOAD_2ADDR_STRIDE64_B64 works correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B64, addr=v[10], vdst=v[2:5], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should have val1 low")
self.assertEqual(st.vgpr[0][3], 0x22222222, "v3 should have val1 high")
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have val2 low")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have val2 high")
def test_ds_storexchg_2addr_stride64_rtn_b32(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32: exchange at two addresses (offset*256)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B32, addr=v[10], data0=v[0], data1=v[1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B32, addr=v[10], data0=v[2], data1=v[3], vdst=v[4:5], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_2ADDR_STRIDE64_B32, addr=v[10], vdst=v[6:7], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0x11111111, "v4 should have old value")
self.assertEqual(st.vgpr[0][5], 0x22222222, "v5 should have old value")
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have new value")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have new value")
def test_ds_storexchg_2addr_stride64_rtn_b64_returns_old(self):
"""DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64: returns old values correctly."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
DS(DSOp.DS_STORE_2ADDR_STRIDE64_B64, addr=v[10], data0=v[0:1], data1=v[0:1], vdst=v[0], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[6], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[7], s[0]),
DS(DSOp.DS_STOREXCHG_2ADDR_STRIDE64_RTN_B64, addr=v[10], data0=v[6:7], data1=v[6:7], vdst=v[8:11], offset0=1, offset1=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][8], 0x11111111, "v8 should have old val1 low")
self.assertEqual(st.vgpr[0][9], 0x22222222, "v9 should have old val1 high")
self.assertEqual(st.vgpr[0][10], 0x11111111, "v10 should have old val2 low")
self.assertEqual(st.vgpr[0][11], 0x22222222, "v11 should have old val2 high")
class TestAtomicOrdering(unittest.TestCase):
"""Tests for atomic operation return values and ordering."""
def test_ds_add_rtn_sequence(self):
"""DS_ADD_RTN returns correct old values in sequence."""
instructions = [
v_mov_b32_e32(v[10], 0),
v_mov_b32_e32(v[0], 100),
DS(DSOp.DS_STORE_B32, addr=v[10], data0=v[0], vdst=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[1], 25),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[2], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_ADD_RTN_U32, addr=v[10], data0=v[1], vdst=v[3], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
DS(DSOp.DS_LOAD_B32, addr=v[10], vdst=v[4], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100, "First add should return 100")
self.assertEqual(st.vgpr[0][3], 125, "Second add should return 125")
self.assertEqual(st.vgpr[0][4], 150, "Final value should be 150")
class TestDsPermute(unittest.TestCase):
"""Tests for DS_PERMUTE_B32 and DS_BPERMUTE_B32 instructions."""
def test_ds_permute_b32_identity(self):
"""DS_PERMUTE_B32 with identity permutation (lane 0 sends to lane 0)."""
# For simplicity, test with single lane
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (lane 0)
v_mov_b32_e32(v[1], 0xDEADBEEF), # data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 sends to lane 0, so lane 0 gets 0xDEADBEEF
self.assertEqual(st.vgpr[0][2], 0xDEADBEEF)
def test_ds_bpermute_b32_identity(self):
"""DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
instructions = [
v_mov_b32_e32(v[0], 0), # addr = 0 (read from lane 0)
v_mov_b32_e32(v[1], 0xCAFEBABE), # data in lane 0
ds_bpermute_b32(v[2], v[0], v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
# Lane 0 reads from lane 0's v[1]
self.assertEqual(st.vgpr[0][2], 0xCAFEBABE)
def test_ds_permute_b32_broadcast(self):
"""DS_PERMUTE_B32 broadcast - all lanes send to lane 0."""
# With 4 lanes, all sending to lane 0, highest lane wins
instructions = [
v_mov_b32_e32(v[0], 0), # All lanes send to addr 0 (lane 0)
v_mov_b32_e32(v[1], 0x11111111), # All lanes send same data
ds_permute_b32(v[2], v[0], v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=4)
# Lane 0 receives data (highest numbered active lane wins)
self.assertEqual(st.vgpr[0][2], 0x11111111)
def test_ds_bpermute_b32_xor_swap(self):
"""DS_BPERMUTE_B32 with XOR-1 pattern — each lane reads from lane^1.
This is the pattern used by warp_shfl_xor in flash attention for reduce_max/reduce_sum.
Each lane has a unique value (lane_id + 100), and reads from the adjacent lane.
"""
instructions = [
# v[0] = (lane_id ^ 1) * 4 (byte offset for bpermute)
v_xor_b32_e32(v[0], 1, v[255]),
v_lshlrev_b32_e32(v[0], 2, v[0]),
# v[1] = lane_id + 100 (unique per-lane value)
s_mov_b32(s[0], 100),
v_add_nc_u32_e32(v[1], s[0], v[255]),
# ds_bpermute: v[2] = v[1] from lane (lane_id ^ 1)
ds_bpermute_b32(vdst=v[2], addr=v[0], data0=v[1]),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=32)
for lane in range(32):
src_lane = lane ^ 1
expected = src_lane + 100
self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
class TestDSSubDword(unittest.TestCase):
"""Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
def test_ds_store_b16_and_d16_hi(self):
"""DS_STORE_B16 stores low 16 bits, DS_STORE_B16_D16_HI stores high 16 bits to adjacent LDS half-words."""
instructions = [
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0xBEEF1234),
DS(DSOp.DS_STORE_B16, addr=v[0], data0=v[1], offset0=0),
DS(DSOp.DS_STORE_B16_D16_HI, addr=v[0], data0=v[1], offset0=2),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(vdst=v[2], addr=v[0], offset0=0),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xBEEF1234, "lo=0x1234 at byte 0, hi=0xBEEF at byte 2")
class TestDSLargeOffset(unittest.TestCase):
"""Tests for DS instructions with offsets > 255 (offset1 > 0).
The DS offset is a 16-bit value encoded as (offset1 << 8) | offset0.
These tests verify that offset1 is used correctly, not just offset0.
"""
def test_ds_store_load_b32_offset_256(self):
"""DS_STORE_B32/DS_LOAD_B32 with offset=256 (offset0=0, offset1=1)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[1], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
def test_ds_store_load_b32_offset_300(self):
"""DS_STORE_B32/DS_LOAD_B32 with offset=300 (offset0=44, offset1=1)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=44, offset1=1), # offset = 300
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[1], offset0=44, offset1=1), # offset = 300
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
def test_ds_store_load_b64_offset_512(self):
"""DS_STORE_B64/DS_LOAD_B64 with offset=512 (offset0=0, offset1=2)."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[1], s[0]),
ds_store_b64(addr=v[10], data0=v[0:1], offset0=0, offset1=2), # offset = 512
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b64(addr=v[10], vdst=v[2:3], offset0=0, offset1=2), # offset = 512
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x11111111)
self.assertEqual(st.vgpr[0][3], 0x22222222)
def test_ds_large_offset_distinct_from_small(self):
"""Verify offset=256 and offset=0 address different LDS locations."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[1], s[0]),
# Store 0xAAAAAAAA at offset=0, 0xBBBBBBBB at offset=256
ds_store_b32(addr=v[10], data0=v[0], offset0=0, offset1=0), # offset = 0
ds_store_b32(addr=v[10], data0=v[1], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Read back both
ds_load_b32(addr=v[10], vdst=v[2], offset0=0, offset1=0), # offset = 0
ds_load_b32(addr=v[10], vdst=v[3], offset0=0, offset1=1), # offset = 256
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "offset=0 should read 0xAAAAAAAA")
self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "offset=256 should read 0xBBBBBBBB")
def test_ds_store_load_b32_offset_448(self):
"""DS_STORE_B32/DS_LOAD_B32 with offset=448 (offset0=192, offset1=1) - matches matmul B tile."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[0], s[0]),
ds_store_b32(addr=v[10], data0=v[0], offset0=192, offset1=1), # offset = 448
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b32(addr=v[10], vdst=v[1], offset0=192, offset1=1), # offset = 448
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x12345678)
def test_ds_load_b64_offset_392(self):
"""DS_LOAD_B64 with offset=392 (offset0=136, offset1=1) - matches matmul B tile load."""
instructions = [
v_mov_b32_e32(v[10], 0),
s_mov_b32(s[0], 0xAABBCCDD),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[0], 0x11223344),
v_mov_b32_e32(v[1], s[0]),
ds_store_b64(addr=v[10], data0=v[0:1], offset0=136, offset1=1), # offset = 392
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
ds_load_b64(addr=v[10], vdst=v[2:3], offset0=136, offset1=1), # offset = 392
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xAABBCCDD)
self.assertEqual(st.vgpr[0][3], 0x11223344)
if __name__ == '__main__':
unittest.main()

363
tinygrad_repo/test/amd/hw/test_flat.py Normal file
View File

@@ -0,0 +1,363 @@
"""Tests for FLAT instructions - flat memory operations.
Includes: flat_load_*, flat_store_*, flat_atomic_*
"""
import unittest
from test.amd.hw.helpers import *
class TestFlatAtomic(unittest.TestCase):
"""Tests for FLAT atomic instructions."""
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
"""Helper to create atomic test instructions."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
] + setup_instrs + [atomic_instr, s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
check_fn(st)
def test_flat_atomic_add_u32(self):
"""FLAT_ATOMIC_ADD_U32 adds to memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_swap_b32(self):
"""FLAT_ATOMIC_SWAP_B32 swaps memory value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_and_b32(self):
"""FLAT_ATOMIC_AND_B32 ANDs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xFF00FF00),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xFFFF0000),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_AND_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xFF00FF00)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_or_b32(self):
"""FLAT_ATOMIC_OR_B32 ORs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0x00FF0000),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0x0000FF00),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_OR_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0x00FF0000)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_inc_u32(self):
"""FLAT_ATOMIC_INC_U32 increments and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 10),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 100), # threshold
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 10)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_dec_u32(self):
"""FLAT_ATOMIC_DEC_U32 decrements and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 10),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[3], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_DEC_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 10)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_sub_u32(self):
"""FLAT_ATOMIC_SUB_U32 subtracts from memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 30),
v_mov_b32_e32(v[3], s[0]), # sub 30
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SUB_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_xor_b32(self):
"""FLAT_ATOMIC_XOR_B32 XORs with memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[3], s[0]), # XOR mask
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_XOR_B32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 should have old value")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_min_u32(self):
"""FLAT_ATOMIC_MIN_U32 stores min and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]), # compare value (smaller)
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_MIN_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100, "v4 should have old value (100)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_max_u32(self):
"""FLAT_ATOMIC_MAX_U32 stores max and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[3], s[0]), # compare value (larger)
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_MAX_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 50, "v4 should have old value (50)")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_inc_u64_returns_old_value(self):
"""FLAT_ATOMIC_INC_U64 should return full 64-bit old value."""
TEST_OFFSET = 2000
setup = [
# Store initial 64-bit value: 0xCAFEBABE_DEADBEEF
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Threshold: 0xFFFFFFFF_FFFFFFFF
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[4], s[0]),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_INC_U64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF, "v6 should have old value low dword")
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE, "v7 should have old value high dword")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_add_u64(self):
"""FLAT_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0x00000001), # add 1
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_ADD_U64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0x11111111, "v6 should have old value low")
self.assertEqual(st.vgpr[0][7], 0x22222222, "v7 should have old value high")
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_flat_atomic_swap_b64(self):
"""FLAT_ATOMIC_SWAP_B64 swaps 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0xDDDDDDDD),
v_mov_b32_e32(v[5], s[0]),
]
atomic = FLAT(FLATOp.FLAT_ATOMIC_SWAP_B64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xAAAAAAAA, "v6 should have old value low")
self.assertEqual(st.vgpr[0][7], 0xBBBBBBBB, "v7 should have old value high")
self._make_test(setup, atomic, check, TEST_OFFSET)
class TestFlatLoad(unittest.TestCase):
"""Tests for FLAT load instructions."""
def test_flat_load_b32(self):
"""FLAT_LOAD_B32 loads 32-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
FLAT(FLATOp.FLAT_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
def test_flat_load_b64(self):
"""FLAT_LOAD_B64 loads 64-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
FLAT(FLATOp.FLAT_LOAD_B64, addr=v[0:1], vdst=v[4:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][4], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][5], 0xCAFEBABE)
def test_flat_load_b96(self):
"""FLAT_LOAD_B96 loads 96-bit (3 dword) value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[4], s[0]),
global_store_b96(addr=v[0:1], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
FLAT(FLATOp.FLAT_LOAD_B96, addr=v[0:1], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][5], 0x11111111)
self.assertEqual(st.vgpr[0][6], 0x22222222)
self.assertEqual(st.vgpr[0][7], 0x33333333)
def test_flat_load_b128(self):
"""FLAT_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x33333333),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x44444444),
v_mov_b32_e32(v[5], s[0]),
global_store_b128(addr=v[0:1], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
FLAT(FLATOp.FLAT_LOAD_B128, addr=v[0:1], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][6], 0x11111111)
self.assertEqual(st.vgpr[0][7], 0x22222222)
self.assertEqual(st.vgpr[0][8], 0x33333333)
self.assertEqual(st.vgpr[0][9], 0x44444444)
if __name__ == '__main__':
unittest.main()

679
tinygrad_repo/test/amd/hw/test_global.py Normal file
View File

@@ -0,0 +1,679 @@
"""Tests for GLOBAL instructions - global memory operations.
Includes: global_load_*, global_store_*, global_atomic_*, global_load_d16_*
"""
import unittest
from test.amd.hw.helpers import *
class TestGlobalAtomic(unittest.TestCase):
"""Tests for GLOBAL atomic instructions."""
def _make_test(self, setup_instrs, atomic_instr, check_fn, test_offset=2000):
"""Helper to create atomic test instructions."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
] + setup_instrs + [atomic_instr, s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
check_fn(st)
def test_global_atomic_add_u32(self):
"""GLOBAL_ATOMIC_ADD_U32 adds to memory and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 100),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 50),
v_mov_b32_e32(v[3], s[0]),
]
atomic = GLOBAL(GLOBALOp.GLOBAL_ATOMIC_ADD_U32, addr=v[0:1], data=v[3], vdst=v[4], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][4], 100)
self._make_test(setup, atomic, check, TEST_OFFSET)
def test_global_atomic_add_u64(self):
"""GLOBAL_ATOMIC_ADD_U64 adds 64-bit value and returns old value."""
TEST_OFFSET = 2000
setup = [
s_mov_b32(s[0], 0xFFFFFFFF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[3], s[0]),
global_store_b64(addr=v[0:1], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[0], 0x00000001),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x00000000),
v_mov_b32_e32(v[5], s[0]),
]
atomic = GLOBAL(GLOBALOp.GLOBAL_ATOMIC_ADD_U64, addr=v[0:1], data=v[4:5], vdst=v[6:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET, glc=1)
def check(st):
self.assertEqual(st.vgpr[0][6], 0xFFFFFFFF)
self.assertEqual(st.vgpr[0][7], 0x00000000)
self._make_test(setup, atomic, check, TEST_OFFSET)
class TestGlobalLoad(unittest.TestCase):
"""Tests for GLOBAL load instructions."""
def test_global_load_b96(self):
"""GLOBAL_LOAD_B96 loads 96-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[0]),
global_store_b96(addr=v[0:1], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B96, addr=v[0:1], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][5], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][6], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][7], 0xCCCCCCCC)
def test_global_load_b128(self):
"""GLOBAL_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 2000
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[0], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[0]),
s_mov_b32(s[0], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[0]),
s_mov_b32(s[0], 0x12345678),
v_mov_b32_e32(v[4], s[0]),
s_mov_b32(s[0], 0x9ABCDEF0),
v_mov_b32_e32(v[5], s[0]),
global_store_b128(addr=v[0:1], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B128, addr=v[0:1], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][6], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][7], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][8], 0x12345678)
self.assertEqual(st.vgpr[0][9], 0x9ABCDEF0)
class TestGlobalStore(unittest.TestCase):
"""Tests for GLOBAL store instructions."""
def test_global_store_b8_basic(self):
"""GLOBAL_STORE_B8 stores a single byte from VDATA[7:0]."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# First store 0xDEADBEEF to memory
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Now store single byte 0x42 to same address (should only change byte 0)
v_mov_b32_e32(v[2], 0x42),
global_store_b8(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Read back and check
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Only byte 0 should change from 0xEF to 0x42
self.assertEqual(st.vgpr[0][0], 0xDEADBE42, "Only byte 0 should be modified")
def test_global_store_b8_byte1(self):
"""GLOBAL_STORE_B8 at offset+1 stores to byte 1."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[2], 0x42),
global_store_b8(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+1),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEAD42EF, "Only byte 1 should be modified")
def test_global_store_b16_basic(self):
"""GLOBAL_STORE_B16 stores a 16-bit value from VDATA[15:0]."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADCAFE, "Only lower 16 bits should be modified")
def test_global_store_b16_high_half(self):
"""GLOBAL_STORE_B16 at offset+2 stores to high 16 bits."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+2),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xCAFEBEEF, "Only upper 16 bits should be modified")
def test_global_store_b16_byte_offset_1(self):
"""GLOBAL_STORE_B16 at byte offset 1 stores bytes 1-2 within the same word."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDDCCBBAA),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Store 0xBEEF at byte offset 1 (bytes 1-2)
s_mov_b32(s[4], 0xBEEF),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+1),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Bytes 1-2 should be 0xBEEF (0xEF at byte 1, 0xBE at byte 2)
# Original: 0xDDCCBBAA -> bytes [AA, BB, CC, DD]
# After: 0xDDBEEFAA -> bytes [AA, EF, BE, DD]
self.assertEqual(st.vgpr[0][0], 0xDDBEEFAA, "Bytes 1-2 should be 0xBEEF")
def test_global_store_b16_cross_word_boundary(self):
"""GLOBAL_STORE_B16 at byte offset 3 crosses word boundary (byte 3 of word N, byte 0 of word N+1)."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Initialize two consecutive words
s_mov_b32(s[4], 0xDDCCBBAA),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[4], 0x44332211),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Store 0xBEEF at byte offset 3 (crosses word boundary)
# Low byte (0xEF) goes to byte 3 of first word
# High byte (0xBE) goes to byte 0 of second word
s_mov_b32(s[4], 0xBEEF),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+3),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back both words
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[4], data=v[4], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], v[4]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# First word: 0xDDCCBBAA -> 0xEFCCBBAA (byte 3 becomes 0xEF)
# Second word: 0x44332211 -> 0x443322BE (byte 0 becomes 0xBE)
self.assertEqual(st.vgpr[0][0], 0xEFCCBBAA, "Byte 3 of first word should be 0xEF")
self.assertEqual(st.vgpr[0][1], 0x443322BE, "Byte 0 of second word should be 0xBE")
def test_global_store_b64_basic(self):
"""GLOBAL_STORE_B64 stores 8 bytes from v[n:n+1] to memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
s_mov_b32(s[5], 0xCAFEBABE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b64(addr=v[0], data=v[2:3], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B64, addr=v[0], vdst=v[4:5], data=v[4:5], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
class TestD16HiLoads(unittest.TestCase):
"""Tests for D16_HI load instructions that load into high 16 bits."""
def test_global_load_d16_hi_b16_preserves_low_bits(self):
"""GLOBAL_LOAD_D16_HI_B16 must preserve low 16 bits of destination."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
global_store_b16(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[3], s[4]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[0:1], vdst=v[3], data=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFEBEEF, f"Expected 0xCAFEBEEF, got 0x{result:08x}")
def test_global_load_d16_hi_b16_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_B16 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b16(addr=v[3], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[0], s[4]), # data field - should NOT affect result
v_mov_b32_e32(v[1], 0), # vdst - low bits should be preserved
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
def test_global_load_d16_hi_u8_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_U8 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xAB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b8(addr=v[3], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[4], s[4]), # data field
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[5], s[4]), # vdst
v_mov_b32_e32(v[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_U8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0x00ABBEEF, f"Expected 0x00ABBEEF, got 0x{result:08x}")
def test_global_load_d16_hi_b16_same_addr_and_dst_zero_addr(self):
"""GLOBAL_LOAD_D16_HI_B16 with same register for addr and vdst, addr value=0."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b16(addr=v[3], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[1], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
self.assertEqual(result, 0xCAFE0000, f"Expected 0xCAFE0000, got 0x{result:08x}")
def test_global_load_d16_hi_b16_tril_exact_pattern(self):
"""Exact pattern from tril() failure: data=v0 differs from vdst=v1."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x01010101),
v_mov_b32_e32(v[10], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b32(addr=v[3], data=v[10], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[3], data=v[10], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Set v[0] to 0x0101 (simulating prior u16 load result)
s_mov_b32(s[4], 0x0101),
v_mov_b32_e32(v[0], s[4]),
# Set v[1] to 0
v_mov_b32_e32(v[1], 0),
# Load using v[1] as addr AND vdst, but v[0] as data
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[0], saddr=s[2:3], offset=TEST_OFFSET+6),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[1]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
# Expected: hi=0x0101 (loaded), lo=0x0000 (from v1) -> 0x01010000
self.assertEqual(result, 0x01010000, f"Expected 0x01010000, got 0x{result:08x}")
def test_global_load_d16_hi_i8_data_differs_from_vdst(self):
"""GLOBAL_LOAD_D16_HI_I8 where data field differs from vdst."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x80), # negative signed byte = -128
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], 0),
global_store_b8(addr=v[3], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x0000DEAD),
v_mov_b32_e32(v[4], s[4]), # data field
s_mov_b32(s[4], 0x0000BEEF),
v_mov_b32_e32(v[5], s[4]), # vdst
v_mov_b32_e32(v[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_I8, addr=v[3], vdst=v[5], data=v[4], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][0]
# 0x80 sign-extended = 0xFF80, lo=0xBEEF -> 0xFF80BEEF
self.assertEqual(result, 0xFF80BEEF, f"Expected 0xFF80BEEF, got 0x{result:08x}")
def test_global_store_b64_tril_pattern(self):
"""Test the exact pattern from tril() kernel that was failing."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0x01010101),
v_mov_b32_e32(v[10], s[4]),
v_mov_b32_e32(v[11], s[4]),
s_mov_b32(s[4], 0x01),
v_mov_b32_e32(v[12], s[4]),
v_mov_b32_e32(v[0], 0),
global_store_b64(addr=v[0], data=v[10:11], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b8(addr=v[0], data=v[12], saddr=s[2:3], offset=TEST_OFFSET+8),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[2], 0),
v_mov_b32_e32(v[1], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_U16, addr=v[2], vdst=v[0], data=v[0], saddr=s[2:3], offset=TEST_OFFSET+3),
GLOBAL(GLOBALOp.GLOBAL_LOAD_D16_HI_B16, addr=v[1], vdst=v[1], data=v[1], saddr=s[2:3], offset=TEST_OFFSET+6),
GLOBAL(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[3], data=v[3], saddr=s[2:3], offset=TEST_OFFSET),
GLOBAL(GLOBALOp.GLOBAL_LOAD_U8, addr=v[2], vdst=v[4], data=v[4], saddr=s[2:3], offset=TEST_OFFSET+8),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_and_b32_e32(v[5], 0xffff, v[0]),
v_lshlrev_b32_e32(v[0], 24, v[0]),
v_lshrrev_b32_e32(v[5], 8, v[5]),
v_or_b32_e32(v[0], v[3], v[0]),
v_or_b32_e32(v[1], v[5], v[1]),
global_store_b64(addr=v[2], data=v[0:1], saddr=s[2:3], offset=TEST_OFFSET+16),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B64, addr=v[2], vdst=v[6:7], data=v[6:7], saddr=s[2:3], offset=TEST_OFFSET+16),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
v0 = st.vgpr[0][0]
v1 = st.vgpr[0][1]
self.assertEqual(v0, 0x01000001, f"v0: expected 0x01000001, got 0x{v0:08x}")
self.assertEqual(v1, 0x01010001, f"v1: expected 0x01010001, got 0x{v1:08x}")
byte5 = (v1 >> 8) & 0xff
self.assertEqual(byte5, 0x00, f"byte5: expected 0x00, got 0x{byte5:02x}")
class TestGlobalOffset(unittest.TestCase):
"""Tests for GLOBAL instructions with different offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load/store wrong data.
"""
def test_global_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load from offset 100 -> should get 0xAAAAAAAA
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[3], saddr=SrcEnum.NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=200),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
v_mov_b32_e32(v[1], v[4]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_store_different_offsets(self):
"""Store to two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0x11111111 at offset 300
s_mov_b32(s[0], 0x11111111),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=300),
# Store 0x22222222 at offset 400
s_mov_b32(s[0], 0x22222222),
v_mov_b32_e32(v[3], s[0]),
global_store_b32(addr=v[0:1], data=v[3], saddr=SrcEnum.NULL, offset=400),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back to verify
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[4], saddr=SrcEnum.NULL, offset=300),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0:1], vdst=v[5], saddr=SrcEnum.NULL, offset=400),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0x11111111, f"offset 300: expected 0x11111111, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0x22222222, f"offset 400: expected 0x22222222, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_no_saddr(self):
"""Test negative offset without saddr (VGPR pair for address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from vaddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from vaddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], s[2]),
v_mov_b32_e32(v[1], s[3]),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0:1], data=v[2], saddr=SrcEnum.NULL, offset=200),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# vaddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[6], saddr=SrcEnum.NULL, offset=-100),
# vaddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
v_mov_b32_e32(v[4], s[4]),
v_mov_b32_e32(v[5], s[5]),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[4:5], vdst=v[7], saddr=SrcEnum.NULL, offset=-100),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[4], 0),
v_mov_b32_e32(v[5], 0),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
def test_global_negative_offset_with_saddr(self):
"""Test negative offset with saddr (SGPR pair for base address).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from saddr pointing to base+200 -> should get 0xAAAA (at 100).
Load with offset -100 from saddr pointing to base+300 -> should get 0xBBBB (at 200)."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[0], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[0], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[0]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# saddr = base+200, load with offset -100 -> should get value at 100
s_add_u32(s[4], s[2], 200),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[6], saddr=s[4:5], offset=-100),
# saddr = base+300, load with offset -100 -> should get value at 200
s_add_u32(s[4], s[2], 300),
s_addc_u32(s[5], s[3], 0),
GLOBAL(GLOBALOp.GLOBAL_LOAD_B32, addr=v[0], vdst=v[7], saddr=s[4:5], offset=-100),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[6], 0),
v_mov_b32_e32(v[7], 0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[4], 0),
s_mov_b32(s[5], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.vgpr[0][0]:08x}")
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.vgpr[0][1]:08x}")
if __name__ == '__main__':
unittest.main()

129
tinygrad_repo/test/amd/hw/test_rdna4_permlane_var.py Normal file
View File

@@ -0,0 +1,129 @@
"""RDNA4 V_PERMLANE16_VAR_B32 / V_PERMLANEX16_VAR_B32 coverage.
Exercises the generated pcode path end-to-end in the emulator and compares against
real RDNA4 hardware when USE_HW=1.
"""
import ctypes, unittest
import tinygrad.runtime.autogen.amd.rdna4.ins as r4
from tinygrad.helpers import flat_mv
from tinygrad.renderer.amd.dsl import NULL
from test.amd.hw.helpers import USE_HW, assemble
from test.mockgpu.amd.emu import run_asm
LANES = 32
def _code(instructions: list, out_reg: int = 2) -> bytes:
return assemble([
r4.s_mov_b32(r4.s[80], r4.s[0]),
r4.s_mov_b32(r4.s[81], r4.s[1]),
r4.v_mov_b32_e32(r4.v[255], r4.v[0]),
*instructions,
r4.s_load_b64(r4.s[92:93], r4.s[80:81], soffset=NULL),
r4.s_wait_kmcnt(simm16=0),
r4.v_lshlrev_b32_e32(r4.v[240], 2, r4.v[255]),
r4.v_mov_b32_e32(r4.v[241], 0),
r4.global_store_b32(vaddr=r4.v[240:241], saddr=r4.s[92:93], vsrc=r4.v[out_reg]),
r4.s_endpgm(),
])
def _run_emu(instructions: list, out_reg: int = 2) -> list[int]:
out_buf = (ctypes.c_uint32 * LANES)(*([0] * LANES))
args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
code = _code(instructions, out_reg)
kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, LANES, 1, 1, ctypes.addressof(args), arch='rdna4')
assert result == 0, f"run_asm failed with {result}"
return list(out_buf)
def _run_hw(instructions: list, out_reg: int = 2) -> list[int]:
from tinygrad.device import Device
from tinygrad.runtime.ops_amd import AMDProgram
from tinygrad.runtime.support.compiler_amd import HIPCompiler
dev = Device['AMD']
if not dev.arch.startswith('gfx12'): raise unittest.SkipTest('requires RDNA4 hardware')
compiler = HIPCompiler(dev.arch)
code = _code(instructions, out_reg)
byte_str = ', '.join(f'0x{b:02x}' for b in code)
asm_src = f""".text
.globl test
.p2align 8
.type test,@function
test:
.byte {byte_str}
.rodata
.p2align 6
.amdhsa_kernel test
.amdhsa_next_free_vgpr 256
.amdhsa_next_free_sgpr 96
.amdhsa_wavefront_size32 1
.amdhsa_user_sgpr_kernarg_segment_ptr 1
.amdhsa_kernarg_size 8
.amdhsa_group_segment_fixed_size 65536
.amdhsa_private_segment_fixed_size 65536
.amdhsa_enable_private_segment 1
.end_amdhsa_kernel
.amdgpu_metadata
---
amdhsa.version:
- 1
- 0
amdhsa.kernels:
- .name: test
.symbol: test.kd
.kernarg_segment_size: 8
.group_segment_fixed_size: 65536
.private_segment_fixed_size: 65536
.kernarg_segment_align: 8
.wavefront_size: 32
.sgpr_count: 96
.vgpr_count: 256
.max_flat_workgroup_size: 1024
...
.end_amdgpu_metadata
"""
lib = compiler.compile(asm_src)
prg = AMDProgram(dev, 'test', lib)
out_gpu = dev.allocator.alloc(LANES * 4)
prg(out_gpu, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
out = bytearray(LANES * 4)
dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(LANES)]
def run_rdna4(instructions: list, out_reg: int = 2) -> list[int]:
emu = _run_emu(instructions, out_reg)
if not USE_HW: return emu
hw = _run_hw(instructions, out_reg)
if emu != hw:
diffs = [f"lane {i}: emu=0x{e:08x} hw=0x{h:08x}" for i, (e, h) in enumerate(zip(emu, hw)) if e != h]
raise AssertionError("Emulator vs Hardware mismatch:\n" + '\n'.join(diffs[:16]))
return hw
class TestPermlaneVarRDNA4(unittest.TestCase):
def test_v_permlane16_var_b32_reverse(self):
out = run_rdna4([
r4.v_mov_b32_e32(r4.v[0], r4.v[255]),
r4.v_xor_b32_e32(r4.v[1], 15, r4.v[255]),
r4.v_permlane16_var_b32(r4.v[2], r4.v[0], r4.v[1]),
])
self.assertEqual(out[0], 15)
self.assertEqual(out[5], 10)
self.assertEqual(out[15], 0)
self.assertEqual(out[16], 31)
self.assertEqual(out[21], 26)
self.assertEqual(out[31], 16)
def test_v_permlanex16_var_b32_cross_row(self):
out = run_rdna4([
r4.v_mov_b32_e32(r4.v[0], r4.v[255]),
r4.v_mov_b32_e32(r4.v[1], r4.v[255]),
r4.v_permlanex16_var_b32(r4.v[2], r4.v[0], r4.v[1]),
])
self.assertEqual(out[0], 16)
self.assertEqual(out[5], 21)
self.assertEqual(out[15], 31)
self.assertEqual(out[16], 0)
self.assertEqual(out[21], 5)
self.assertEqual(out[31], 15)

355
tinygrad_repo/test/amd/hw/test_scratch.py Normal file
View File

@@ -0,0 +1,355 @@
"""Tests for SCRATCH instructions - scratch (private) memory operations.
Includes: scratch_load_*, scratch_store_*
"""
import unittest
from test.amd.hw.helpers import *
class TestScratchStore(unittest.TestCase):
"""Tests for SCRATCH store instructions."""
def test_scratch_store_b32_basic(self):
"""SCRATCH_STORE_B32 stores 32-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
# Store via scratch
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back via scratch
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
def test_scratch_store_b64_basic(self):
"""SCRATCH_STORE_B64 stores 64-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
s_mov_b32(s[5], 0xCAFEBABE),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
scratch_store_b64(addr=v[0], data=v[2:3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_b64(addr=v[0], vdst=v[4:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[4]),
v_mov_b32_e32(v[1], v[5]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
def test_scratch_store_b8_basic(self):
"""SCRATCH_STORE_B8 stores single byte to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# First store full word
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Store single byte
v_mov_b32_e32(v[2], 0x42),
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Only byte 0 should change from 0xEF to 0x42
self.assertEqual(st.vgpr[0][0], 0xDEADBE42)
def test_scratch_store_b16_basic(self):
"""SCRATCH_STORE_B16 stores 16-bit value to scratch memory."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[0], 0),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[4], 0xCAFE),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADCAFE)
class TestScratchLoad(unittest.TestCase):
"""Tests for SCRATCH load instructions."""
def test_scratch_load_b96(self):
"""SCRATCH_LOAD_B96 loads 96-bit value correctly."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[3], s[4]),
s_mov_b32(s[4], 0xCCCCCCCC),
v_mov_b32_e32(v[4], s[4]),
scratch_store_b96(addr=v[0], data=v[2:4], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_b96(addr=v[0], vdst=v[5:7], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[5]),
v_mov_b32_e32(v[1], v[6]),
v_mov_b32_e32(v[2], v[7]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA)
self.assertEqual(st.vgpr[0][1], 0xBBBBBBBB)
self.assertEqual(st.vgpr[0][2], 0xCCCCCCCC)
def test_scratch_load_b128(self):
"""SCRATCH_LOAD_B128 loads 128-bit value correctly."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[4]),
s_mov_b32(s[4], 0xCAFEBABE),
v_mov_b32_e32(v[3], s[4]),
s_mov_b32(s[4], 0x12345678),
v_mov_b32_e32(v[4], s[4]),
s_mov_b32(s[4], 0x9ABCDEF0),
v_mov_b32_e32(v[5], s[4]),
scratch_store_b128(addr=v[0], data=v[2:5], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_b128(addr=v[0], vdst=v[6:9], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[6]),
v_mov_b32_e32(v[1], v[7]),
v_mov_b32_e32(v[2], v[8]),
v_mov_b32_e32(v[3], v[9]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xDEADBEEF)
self.assertEqual(st.vgpr[0][1], 0xCAFEBABE)
self.assertEqual(st.vgpr[0][2], 0x12345678)
self.assertEqual(st.vgpr[0][3], 0x9ABCDEF0)
def test_scratch_load_u8(self):
"""SCRATCH_LOAD_U8 loads unsigned byte with zero extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADBEAB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_u8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xAB)
def test_scratch_load_i8(self):
"""SCRATCH_LOAD_I8 loads signed byte with sign extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0x80), # -128 as signed byte
v_mov_b32_e32(v[2], s[4]),
scratch_store_b8(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_i8(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFF80)
def test_scratch_load_u16(self):
"""SCRATCH_LOAD_U16 loads unsigned 16-bit with zero extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xDEADCAFE),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_u16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xCAFE)
def test_scratch_load_i16(self):
"""SCRATCH_LOAD_I16 loads signed 16-bit with sign extension."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0x8000), # -32768 as signed 16-bit
v_mov_b32_e32(v[2], s[4]),
scratch_store_b16(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
scratch_load_i16(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFF8000)
class TestScratchSVE(unittest.TestCase):
"""Tests for SCRATCH SVE (Scratch VGPR Enable) bit behavior."""
def test_scratch_sve_zero_ignores_vaddr(self):
"""With SVE=0, VADDR should be ignored in address calculation."""
TEST_OFFSET = 256
# Store a marker value at offset 256 (where SVE=0 should go)
# Then set v[0] to a non-zero value (100) and store via scratch with SVE=0
# If SVE=0 is handled correctly, the VADDR (100) should be IGNORED,
# and the store should go to offset 256, not 256+100=356
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Now set v[0] to 100 (non-zero) and store 0xBBBBBBBB with SVE=0
# With SVE=0, v[0] should be IGNORED, so this should overwrite offset 256
v_mov_b32_e32(v[0], 100),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back from offset 256 (with v[0]=0) - should get 0xBBBBBBBB
v_mov_b32_e32(v[0], 0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=0),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# If SVE=0 works correctly, v[0] should be 0xBBBBBBBB (the second store overwrote the first)
# If SVE=0 is wrong (VADDR used), v[0] would be 0xAAAAAAAA (stores went to different locations)
self.assertEqual(st.vgpr[0][0], 0xBBBBBBBB, "SVE=0 should ignore VADDR, both stores should go to same location")
def test_scratch_sve_one_uses_vaddr(self):
"""With SVE=1, VADDR should be used as offset in address calculation."""
TEST_OFFSET = 256
# Store at offset 256 with v[0]=0, then store at offset 256 with v[0]=100 and SVE=1
# With SVE=1, the second store should go to 256+100=356, not 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# First, store 0xAAAAAAAA at offset 256 with v[0]=0
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Now set v[0] to 100 and store 0xBBBBBBBB with SVE=1
# With SVE=1, v[0] IS used, so this should go to offset 256+100=356
v_mov_b32_e32(v[0], 100),
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back from offset 256 (with v[0]=0) - should still be 0xAAAAAAAA
v_mov_b32_e32(v[0], 0),
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET, sve=1),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# If SVE=1 works correctly, v[0] should be 0xAAAAAAAA (stores went to different locations)
self.assertEqual(st.vgpr[0][0], 0xAAAAAAAA, "SVE=1 should use VADDR, stores should go to different locations")
class TestScratchMultiLane(unittest.TestCase):
"""Tests for SCRATCH operations with multiple lanes."""
def test_scratch_store_load_multi_lane(self):
"""SCRATCH store/load works correctly with multiple lanes (private per-lane memory)."""
TEST_OFFSET = 256
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=SrcEnum.NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Each lane stores its lane ID
v_mov_b32_e32(v[0], 0),
v_mov_b32_e32(v[2], v[255]), # v[255] has packed workitem IDs, low 10 bits = x
v_and_b32_e32(v[2], 0x3FF, v[2]), # extract lane ID
scratch_store_b32(addr=v[0], data=v[2], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
# Load back
scratch_load_b32(addr=v[0], vdst=v[3], saddr=SrcEnum.NULL, offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], v[3]),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=4)
# Each lane should have loaded its own lane ID
for lane in range(4):
self.assertEqual(st.vgpr[lane][0], lane, f"Lane {lane} should have value {lane}")
if __name__ == '__main__':
unittest.main()

448
tinygrad_repo/test/amd/hw/test_smem.py Normal file
View File

@@ -0,0 +1,448 @@
"""Tests for SMEM instructions - scalar memory operations.
Includes: s_load_b32, s_load_b64, s_load_b128, s_load_b256, s_load_b512
Tests both immediate and register offset addressing modes.
"""
import unittest
from test.amd.hw.helpers import *
# Use offset into output buffer for test data (output buffer is 2124 bytes)
TEST_OFFSET = 2000
# Cache invalidation sequence for scalar loads after vector stores
# s_wait_idle waits for all outstanding memory operations including cache flushes
CACHE_INV = [s_gl1_inv(), s_dcache_inv(), s_wait_idle()]
class TestSLoadRegisterOffset(unittest.TestCase):
"""Tests for s_load with register offset (soffset field).
Bug: s_load_b32(s[dst], s[base:base+1], s[off]) ignores the register offset
and only uses the immediate offset field. This causes incorrect memory loads
when the offset comes from a register.
"""
def test_s_load_b32_register_offset_basic(self):
"""s_load_b32 with register offset should load from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test values to output buffer: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[4], 0xAAAAAAAA),
s_mov_b32(s[5], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Now test s_load with register offset
# Put offset value in s[4]: offset = 4 bytes (1 dword)
s_mov_b32(s[4], 4),
# Load from out_ptr + TEST_OFFSET + s[4] (should load 0xBBBBBBBB)
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB,
f"s_load with reg offset 4 should load 0xBBBBBBBB: s[5]=0x{st.sgpr[5]:08x}")
def test_s_load_b32_register_offset_different_from_immediate(self):
"""s_load_b32 with register offset loads different data than immediate offset 0."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[4], 0xAAAAAAAA),
s_mov_b32(s[5], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
v_mov_b32_e32(v[3], s[5]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load with immediate offset 0
s_load_b32(s[5], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Load with register offset 4
s_mov_b32(s[4], 4),
s_load_b32(s[6], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# s[5] has dword at offset 0 (0xAAAAAAAA), s[6] has dword at offset 4 (0xBBBBBBBB)
self.assertEqual(st.sgpr[5], 0xAAAAAAAA)
self.assertEqual(st.sgpr[6], 0xBBBBBBBB)
self.assertNotEqual(st.sgpr[5], st.sgpr[6],
f"s_load with reg offset 4 should load different value than offset 0: "
f"s[5]=0x{st.sgpr[5]:08x}, s[6]=0x{st.sgpr[6]:08x}")
def test_s_load_b32_register_offset_same_as_dst(self):
"""s_load_b32 where soffset register is same as destination.
This is the exact pattern that exposes the bug:
s_load_b32(s[8], s[2:3], s[8])
The offset should be read BEFORE the destination is overwritten.
"""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[6], 0xAAAAAAAA),
s_mov_b32(s[7], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[6]),
v_mov_b32_e32(v[3], s[7]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Set up s[4] = 4 (offset in bytes)
s_mov_b32(s[4], 4),
# Load using s[4] as both offset and destination
# Should load from base + 4, then store result in s[4]
s_load_b32(s[4], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Also load with immediate offset 4 for comparison
s_load_b32(s[5], s[2:3], NULL, offset=TEST_OFFSET+4),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# s[4] and s[5] should have the same value (both loaded from offset 4 = 0xBBBBBBBB)
self.assertEqual(st.sgpr[4], 0xBBBBBBBB)
self.assertEqual(st.sgpr[4], st.sgpr[5],
f"s_load with reg offset s[4]=4 should match immediate offset=4: "
f"s[4]=0x{st.sgpr[4]:08x}, s[5]=0x{st.sgpr[5]:08x}")
def test_s_load_b32_register_offset_zero(self):
"""s_load_b32 with register offset = 0 should be same as immediate offset 0."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test value: 0xDEADBEEF at offset
s_mov_b32(s[7], 0xDEADBEEF),
v_mov_b32_e32(v[2], s[7]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load with register offset 0
s_mov_b32(s[4], 0),
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Load with immediate offset 0
s_load_b32(s[6], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[5], 0xDEADBEEF)
self.assertEqual(st.sgpr[5], st.sgpr[6],
f"s_load with reg offset 0 should match immediate offset 0: "
f"s[5]=0x{st.sgpr[5]:08x}, s[6]=0x{st.sgpr[6]:08x}")
def test_s_load_b32_register_plus_immediate_offset(self):
"""s_load_b32 with both register and immediate offset should add them."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test values: 0xAAAAAAAA at offset, 0xBBBBBBBB at offset+4
s_mov_b32(s[8], 0xAAAAAAAA),
s_mov_b32(s[9], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[8]),
v_mov_b32_e32(v[3], s[9]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# reg offset = 4, imm offset = 0 -> total offset = 4
s_mov_b32(s[4], 4),
s_load_b32(s[5], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# reg offset = 0, imm offset = 4 -> total offset = 4
s_mov_b32(s[6], 0),
s_load_b32(s[7], s[2:3], s[6], offset=TEST_OFFSET+4),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
# Both should load from offset 4 (0xBBBBBBBB)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB)
self.assertEqual(st.sgpr[7], 0xBBBBBBBB)
self.assertEqual(st.sgpr[5], st.sgpr[7],
f"reg_off=4 + imm_off=0 should equal reg_off=0 + imm_off=4: "
f"s[5]=0x{st.sgpr[5]:08x}, s[7]=0x{st.sgpr[7]:08x}")
class TestSLoadMultiDword(unittest.TestCase):
"""Tests for multi-dword s_load with register offset."""
def test_s_load_b64_register_offset(self):
"""s_load_b64 with register offset should load 2 dwords from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test values: 0xAAAAAAAA, 0xBBBBBBBB at offset
s_mov_b32(s[10], 0xAAAAAAAA),
s_mov_b32(s[11], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[10]),
v_mov_b32_e32(v[3], s[11]),
v_mov_b32_e32(v[0], 0),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
global_store_b32(addr=v[0], data=v[3], saddr=s[2:3], offset=TEST_OFFSET+4),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load with register offset 0
s_mov_b32(s[4], 0),
s_load_b64(s[6:7], s[2:3], s[4], offset=TEST_OFFSET),
s_waitcnt(0),
# Compare with immediate offset
s_load_b64(s[8:9], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], 0xAAAAAAAA)
self.assertEqual(st.sgpr[7], 0xBBBBBBBB)
self.assertEqual(st.sgpr[6], st.sgpr[8])
self.assertEqual(st.sgpr[7], st.sgpr[9])
def test_s_load_b128_register_offset(self):
"""s_load_b128 with register offset should load 4 dwords from base + reg_offset."""
instructions = [
# Load output buffer pointer from args
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Store test values: 0xAAAAAAAA, 0xBBBBBBBB, 0xCCCCCCCC, 0xDDDDDDDD at offset
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[14], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[14], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_mov_b32(s[14], 0xCCCCCCCC),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+8),
s_mov_b32(s[14], 0xDDDDDDDD),
v_mov_b32_e32(v[2], s[14]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+12),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load with register offset 0 (s_load_b128 requires 4-aligned dest: s[4], s[8], s[12], ...)
s_mov_b32(s[15], 0),
s_load_b128(s[4:7], s[2:3], s[15], offset=TEST_OFFSET),
s_waitcnt(0),
# Compare with immediate offset
s_load_b128(s[8:11], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt(0),
# Zero out pointer regs (different addresses in emu vs hw)
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA)
self.assertEqual(st.sgpr[5], 0xBBBBBBBB)
self.assertEqual(st.sgpr[6], 0xCCCCCCCC)
self.assertEqual(st.sgpr[7], 0xDDDDDDDD)
self.assertEqual(st.sgpr[4], st.sgpr[8])
self.assertEqual(st.sgpr[5], st.sgpr[9])
class TestSLoadLarge(unittest.TestCase):
"""Tests for large s_load operations (s_load_b256, s_load_b512)."""
def test_s_load_b256_basic(self):
"""s_load_b256 loads 8 consecutive dwords."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
# Store 8 test values
s_mov_b32(s[20], 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET),
s_mov_b32(s[20], 0x22222222),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+4),
s_mov_b32(s[20], 0x33333333),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+8),
s_mov_b32(s[20], 0x44444444),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+12),
s_mov_b32(s[20], 0x55555555),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+16),
s_mov_b32(s[20], 0x66666666),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+20),
s_mov_b32(s[20], 0x77777777),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+24),
s_mov_b32(s[20], 0x88888888),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET+28),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load all 8 dwords with s_load_b256
s_load_b256(s[4:11], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111)
self.assertEqual(st.sgpr[5], 0x22222222)
self.assertEqual(st.sgpr[6], 0x33333333)
self.assertEqual(st.sgpr[7], 0x44444444)
self.assertEqual(st.sgpr[8], 0x55555555)
self.assertEqual(st.sgpr[9], 0x66666666)
self.assertEqual(st.sgpr[10], 0x77777777)
self.assertEqual(st.sgpr[11], 0x88888888)
def test_s_load_b512_basic(self):
"""s_load_b512 loads 16 consecutive dwords."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
# Store 16 test values (use a pattern: 0x10, 0x20, ..., 0x100)
*[instr for i in range(16) for instr in [
s_mov_b32(s[20], (i + 1) * 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET + i * 4),
]],
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load all 16 dwords with s_load_b512
s_load_b512(s[64:79], s[2:3], NULL, offset=TEST_OFFSET),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
# Copy results to lower regs for verification (since st.sgpr only has 16 regs in test)
s_mov_b32(s[4], s[64]),
s_mov_b32(s[5], s[65]),
s_mov_b32(s[6], s[78]),
s_mov_b32(s[7], s[79]),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111, "first dword")
self.assertEqual(st.sgpr[5], 0x22222222, "second dword")
self.assertEqual(st.sgpr[6], 0xFFFFFFFF & (15 * 0x11111111), "15th dword")
self.assertEqual(st.sgpr[7], 0xFFFFFFFF & (16 * 0x11111111), "16th dword")
def test_s_load_b256_with_register_offset(self):
"""s_load_b256 with register offset should add reg offset to address."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
# Store pattern at TEST_OFFSET+8: skip first 2 dwords
*[instr for i in range(8) for instr in [
s_mov_b32(s[20], (i + 1) * 0x11111111),
v_mov_b32_e32(v[2], s[20]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=TEST_OFFSET + 8 + i * 4),
]],
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load with register offset 8
s_mov_b32(s[20], 8),
s_load_b256(s[4:11], s[2:3], s[20], offset=TEST_OFFSET),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0x11111111, "first dword at offset+8")
self.assertEqual(st.sgpr[5], 0x22222222, "second dword at offset+8")
self.assertEqual(st.sgpr[11], 0x88888888, "last dword at offset+8")
class TestSLoadOffset(unittest.TestCase):
"""Tests for s_load with different immediate offsets.
These tests verify that instruction deduplication correctly handles different offset values.
If offset is made dynamic incorrectly, instructions with different offsets may load wrong data.
"""
def test_s_load_different_offsets(self):
"""Load from two different offsets and verify correct values."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100
s_mov_b32(s[4], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
# Store 0xBBBBBBBB at offset 200
s_mov_b32(s[4], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[4]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# Load from offset 100 -> should get 0xAAAAAAAA
s_load_b32(s[4], s[2:3], NULL, offset=100),
# Load from offset 200 -> should get 0xBBBBBBBB
s_load_b32(s[5], s[2:3], NULL, offset=200),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA, f"offset 100: expected 0xAAAAAAAA, got 0x{st.sgpr[4]:08x}")
self.assertEqual(st.sgpr[5], 0xBBBBBBBB, f"offset 200: expected 0xBBBBBBBB, got 0x{st.sgpr[5]:08x}")
def test_s_load_negative_offset(self):
"""Test negative offset (21-bit signed).
Store 0xAAAA at offset 100, 0xBBBB at offset 200.
Load with offset -100 from base+200 -> should get 0xAAAA.
Load with offset -100 from base+300 -> should get 0xBBBB."""
instructions = [
s_load_b64(s[2:3], s[80:81], 0, soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[0], 0),
# Store 0xAAAAAAAA at offset 100, 0xBBBBBBBB at offset 200
s_mov_b32(s[8], 0xAAAAAAAA),
v_mov_b32_e32(v[2], s[8]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=100),
s_mov_b32(s[8], 0xBBBBBBBB),
v_mov_b32_e32(v[2], s[8]),
global_store_b32(addr=v[0], data=v[2], saddr=s[2:3], offset=200),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
*CACHE_INV,
# base+200, load with offset -100 -> should get value at 100
s_add_u32(s[6], s[2], 200),
s_addc_u32(s[7], s[3], 0),
s_load_b32(s[4], s[6:7], NULL, offset=-100),
# base+300, load with offset -100 -> should get value at 200
s_add_u32(s[6], s[2], 300),
s_addc_u32(s[7], s[3], 0),
s_load_b32(s[5], s[6:7], NULL, offset=-100),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
s_mov_b32(s[2], 0),
s_mov_b32(s[3], 0),
s_mov_b32(s[6], 0),
s_mov_b32(s[7], 0),
s_mov_b32(s[8], 0),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[4], 0xAAAAAAAA, f"offset 200-100=100: expected 0xAAAAAAAA, got 0x{st.sgpr[4]:08x}")
self.assertEqual(st.sgpr[5], 0xBBBBBBBB, f"offset 300-100=200: expected 0xBBBBBBBB, got 0x{st.sgpr[5]:08x}")
if __name__ == '__main__':
unittest.main()

1007
tinygrad_repo/test/amd/hw/test_sop.py Normal file
View File

File diff suppressed because it is too large Load Diff

35
tinygrad_repo/test/amd/hw/test_vinterp.py Normal file
View File

@@ -0,0 +1,35 @@
"""Tests for VINTERP instructions."""
import unittest
from test.amd.hw.helpers import *
class TestVInterp(unittest.TestCase):
def test_v_interp_p10_f32(self):
instructions = [
v_mov_b32_e32(v[10], v[255]),
v_cvt_f32_u32_e32(v[1], v[10]),
s_mov_b32(s[0], f2i(100.0)),
v_add_f32_e32(v[1], s[0], v[1]),
v_cvt_f32_u32_e32(v[3], v[10]),
s_mov_b32(s[1], f2i(10.0)),
v_add_f32_e32(v[3], s[1], v[3]),
s_mov_b32(s[2], f2i(2.0)),
v_interp_p10_f32(v[4], v[1], s[2], v[3]),
]
st = run_program(instructions, n_lanes=8)
for lane in range(4): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 212.0, places=5)
for lane in range(4, 8): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 224.0, places=5)
def test_v_interp_p10_f16_f32(self):
instructions = [
v_mov_b32_e32(v[10], v[255]),
v_cvt_f32_u32_e32(v[11], v[10]),
v_cvt_f16_f32_e32(v[1], v[11]),
s_mov_b32(s[0], f2i(10.0)),
v_add_f32_e32(v[12], s[0], v[11]),
v_cvt_f16_f32_e32(v[3], v[12]),
s_mov_b32(s[1], f2i(2.0)),
v_interp_p10_f16_f32(v[4], v[1], s[1], v[3]),
]
st = run_program(instructions, n_lanes=8)
for lane in range(4): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 12.0, places=5)
for lane in range(4, 8): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 24.0, places=5)

1633
tinygrad_repo/test/amd/hw/test_vop1.py Normal file
View File

File diff suppressed because it is too large Load Diff

994
tinygrad_repo/test/amd/hw/test_vop2.py Normal file
View File

@@ -0,0 +1,994 @@
"""Tests for VOP2 instructions - two operand vector operations.
Includes: v_add_f32, v_mul_f32, v_and_b32, v_or_b32, v_xor_b32,
v_lshrrev_b32, v_lshlrev_b32, v_fmac_f32, v_fmaak_f32, v_fmamk_f32,
v_add_nc_u32, v_cndmask_b32, v_add_f16, v_mul_f16
"""
import unittest
from test.amd.hw.helpers import *
class TestBasicArithmetic(unittest.TestCase):
"""Tests for basic arithmetic VOP2 instructions."""
def test_v_add_f32(self):
"""V_ADD_F32 adds two floats."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
v_mov_b32_e32(v[1], 2.0),
v_add_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
def test_v_mul_f32(self):
"""V_MUL_F32 multiplies two floats."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_mul_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 8.0, places=5)
def test_v_add_f32_dpp_row_shl(self):
"""V_ADD_F32 DPP row_shl swizzles src0 before the add."""
instructions = [
v_cvt_f32_u32_e32(v[0], v[255]),
v_add_f32_e32(v[1], DPP, v[0], vsrc0=v[0], dpp=0x101, row_mask=0xf, bank_mask=0xf, bc=1),
]
st = run_program(instructions, n_lanes=16)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 1.0, places=5)
self.assertAlmostEqual(i2f(st.vgpr[1][1]), 3.0, places=5)
self.assertAlmostEqual(i2f(st.vgpr[14][1]), 29.0, places=5)
def test_v_fmac_f32(self):
"""V_FMAC_F32: d = d + a*b using inline constants."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_mov_b32_e32(v[2], 1.0),
v_fmac_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmaak_f32(self):
"""V_FMAAK_F32: d = a * b + K using inline constants."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 4.0),
v_fmaak_f32_e32(v[2], v[0], v[1], literal=0x3f800000),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmamk_f32_basic(self):
"""V_FMAMK_F32: d = a * K + b."""
instructions = [
v_mov_b32_e32(v[0], 2.0),
v_mov_b32_e32(v[1], 1.0),
v_fmamk_f32_e32(v[2], v[0], v[1], literal=0x40800000),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 9.0, places=5)
def test_v_fmamk_f32_small_constant(self):
"""V_FMAMK_F32 with small constant."""
instructions = [
v_mov_b32_e32(v[0], 4.0),
v_mov_b32_e32(v[1], 1.0),
v_fmamk_f32_e32(v[2], v[0], v[1], literal=f2i(0.5)),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 3.0, places=5)
class TestBitManipulation(unittest.TestCase):
"""Tests for bit manipulation VOP2 instructions."""
def test_v_and_b32(self):
"""V_AND_B32 bitwise and."""
instructions = [
s_mov_b32(s[0], 0xff),
s_mov_b32(s[1], 0x0f),
v_mov_b32_e32(v[0], s[0]),
v_and_b32_e32(v[1], s[1], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x0f)
def test_v_and_b32_quadrant(self):
"""V_AND_B32 for quadrant extraction (n & 3)."""
instructions = [
s_mov_b32(s[0], 15915),
v_mov_b32_e32(v[0], s[0]),
v_and_b32_e32(v[1], 3, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 15915 & 3)
def test_v_lshrrev_b32(self):
"""V_LSHRREV_B32 logical shift right."""
instructions = [
s_mov_b32(s[0], 0xff00),
v_mov_b32_e32(v[0], s[0]),
v_lshrrev_b32_e32(v[1], 8, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xff)
def test_v_lshlrev_b32(self):
"""V_LSHLREV_B32 logical shift left."""
instructions = [
s_mov_b32(s[0], 0xff),
v_mov_b32_e32(v[0], s[0]),
v_lshlrev_b32_e32(v[1], 8, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0xff00)
def test_v_xor_b32(self):
"""V_XOR_B32 bitwise xor (used in sin for sign)."""
instructions = [
s_mov_b32(s[0], 0x80000000),
s_mov_b32(s[1], f2i(1.0)),
v_mov_b32_e32(v[0], s[1]),
v_xor_b32_e32(v[1], s[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), -1.0, places=5)
def test_v_xor_b32_sign_flip(self):
"""V_XOR_B32 for sign flip pattern."""
instructions = [
s_mov_b32(s[0], 0x80000000),
v_mov_b32_e32(v[0], -2.0),
v_xor_b32_e32(v[1], s[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][1]), 2.0, places=5)
class TestSpecialValues(unittest.TestCase):
"""Tests for special float values - inf, nan, zero handling."""
def test_v_mul_f32_zero_times_inf(self):
"""V_MUL_F32: 0 * inf = NaN."""
import math
instructions = [
v_mov_b32_e32(v[0], 0),
s_mov_b32(s[0], 0x7f800000),
v_mov_b32_e32(v[1], s[0]),
v_mul_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
def test_v_add_f32_inf_minus_inf(self):
"""V_ADD_F32: inf + (-inf) = NaN."""
import math
instructions = [
s_mov_b32(s[0], 0x7f800000),
s_mov_b32(s[1], 0xff800000),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f32_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertTrue(math.isnan(i2f(st.vgpr[0][2])))
class TestF16Ops(unittest.TestCase):
"""Tests for 16-bit VOP2 operations."""
def test_v_add_f16_basic(self):
"""V_ADD_F16 adds two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4200, f"Expected 0x4200 (f16 3.0), got 0x{result:04x}")
def test_v_add_f16_negative(self):
"""V_ADD_F16 with negative values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0xc000), # f16 -2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0xbc00, f"Expected 0xbc00 (f16 -1.0), got 0x{result:04x}")
def test_v_mul_f16_basic(self):
"""V_MUL_F16 multiplies two f16 values."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mul_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4600, f"Expected 0x4600 (f16 6.0), got 0x{result:04x}")
def test_v_mul_f16_by_zero(self):
"""V_MUL_F16 by zero."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0),
v_mul_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x0000, f"Expected 0x0000 (f16 0.0), got 0x{result:04x}")
def test_v_fmac_f16_basic(self):
"""V_FMAC_F16: d = d + a*b."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
s_mov_b32(s[2], 0x3c00), # f16 1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
v_fmac_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
def test_v_max_f16_basic(self):
"""V_MAX_F16 returns the maximum of two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_max_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected 0x4000 (f16 2.0), got 0x{result:04x}")
def test_v_min_f16_basic(self):
"""V_MIN_F16 returns the minimum of two f16 values."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_min_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected 0x3c00 (f16 1.0), got 0x{result:04x}")
def test_v_fmaak_f16_basic(self):
"""V_FMAAK_F16: d = a * b + K."""
instructions = [
s_mov_b32(s[0], 0x4000), # f16 2.0
s_mov_b32(s[1], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_fmaak_f16_e32(v[2], v[0], v[1], literal=0x3c00), # + f16 1.0
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected 0x4700 (f16 7.0), got 0x{result:04x}")
class TestHiHalfOps(unittest.TestCase):
"""Tests for VOP2 16-bit operations with hi-half operands."""
def test_v_add_f16_src0_hi_fold(self):
"""V_ADD_F16 with src0 hi-half fold (same register, different halves)."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # lo=f16(1.0), hi=f16(2.0)
v_mov_b32_e32(v[0], s[0]),
VOP3(VOP3Op.V_ADD_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_add_f16_src0_hi_different_reg(self):
"""V_ADD_F16 with src0 hi-half from different register."""
instructions = [
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
s_mov_b32(s[1], 0x00003c00), # hi=0, lo=f16(1.0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
VOP3(VOP3Op.V_ADD_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_mul_f16_src0_hi(self):
"""V_MUL_F16 with src0 from high half."""
instructions = [
s_mov_b32(s[0], 0x40000000), # hi=f16(2.0), lo=0
s_mov_b32(s[1], 0x00004200), # hi=0, lo=f16(3.0)
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
VOP3(VOP3Op.V_MUL_F16, vdst=v[2], src0=v[0], src1=v[1], opsel=0b0001),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
def test_v_mul_f16_hi_half(self):
"""V_MUL_F16 reading from high half."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # lo=1.0, hi=2.0
v_mov_b32_e32(v[0], s[0]),
VOP3(VOP3Op.V_MUL_F16, vdst=v[1], src0=v[0], src1=v[0], opsel=0b0011),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
self.assertEqual(result, 0x4400, f"Expected f16(4.0)=0x4400, got 0x{result:04x}")
def test_v_fma_f16_hi_dest(self):
"""V_FMA_F16 writing to high half with opsel.
Uses V_FMA_F16 (not V_FMAC_F16) because it has explicit src2 operand
which makes opsel handling clearer.
"""
instructions = [
s_mov_b32(s[0], 0x3c000000), # hi=f16(1.0), lo=0
s_mov_b32(s[1], 0x4000), # f16(2.0) in lo
s_mov_b32(s[2], 0x4200), # f16(3.0) in lo
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
# V_FMA_F16: dst = src0 * src1 + src2
# opsel=0b1100: bit2=src2 hi, bit3=dst hi
# So: v[0].hi = v[1].lo * v[2].lo + v[0].hi = 2.0 * 3.0 + 1.0 = 7.0
VOP3(VOP3Op.V_FMA_F16, vdst=v[0], src0=v[1], src1=v[2], src2=v[0], opsel=0b1100),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][0] >> 16) & 0xffff
# 2.0 * 3.0 + 1.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(hi, 0x4700, f"Expected f16(7.0)=0x4700 in hi, got 0x{hi:04x}")
def test_v_add_f16_multilane(self):
"""V_ADD_F16 with multiple lanes."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_add_f16_e32(v[2], v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
for lane in range(4):
result = st.vgpr[lane][2] & 0xffff
self.assertEqual(result, 0x4200, f"Lane {lane}: expected 0x4200, got 0x{result:04x}")
class TestVop2F16HiHalf(unittest.TestCase):
"""Regression tests for VOP2 f16 hi-half operand handling.
These test the bugs where:
1. VOP2 vsrc1 >= 384 (v[128]+) wasn't extracting hi 16 bits
2. VOP2 vdst >= 384 (v[128]+) wasn't preserving lo 16 bits
"""
def test_v_add_f16_e32_vsrc1_hi_half(self):
"""V_ADD_F16_E32 with vsrc1 from hi-half (v[128]+).
When vsrc1 >= 384 (representing v[128]+), the hardware reads from the hi 16 bits
of v[vsrc1-128]. The emulator must extract bits [31:16] from the actual VGPR.
Regression test for: VOP2 f16 vsrc1 hi-half extraction bug.
"""
instructions = [
# v[0] = 0x4000_3c00: hi=f16(2.0), lo=f16(1.0)
s_mov_b32(s[0], 0x40003c00),
v_mov_b32_e32(v[0], s[0]),
# v_add_f16_e32 v[1], v[0], v[128] (vsrc1=v[128] reads hi of v[0])
# In VOP2 encoding, vsrc1=384 means v[128], which maps to v[0].hi
# v[1] = v[0].lo + v[0].hi = 1.0 + 2.0 = 3.0
VOP2(VOP2Op.V_ADD_F16, vdst=v[1], src0=v[0], vsrc1=v[128]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
# 1.0 + 2.0 = 3.0, f16 3.0 = 0x4200
self.assertEqual(result, 0x4200, f"Expected f16(3.0)=0x4200, got 0x{result:04x}")
def test_v_mul_f16_e32_vsrc1_hi_half(self):
"""V_MUL_F16_E32 with vsrc1 from hi-half.
Regression test for: VOP2 f16 vsrc1 hi-half extraction bug.
"""
instructions = [
# v[0] = 0x4200_4000: hi=f16(3.0), lo=f16(2.0)
s_mov_b32(s[0], 0x42004000),
v_mov_b32_e32(v[0], s[0]),
# v_mul_f16_e32 v[1], v[0], v[128] (vsrc1=v[128] reads hi of v[0])
# v[1] = v[0].lo * v[0].hi = 2.0 * 3.0 = 6.0
VOP2(VOP2Op.V_MUL_F16, vdst=v[1], src0=v[0], vsrc1=v[128]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
# 2.0 * 3.0 = 6.0, f16 6.0 = 0x4600
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
def test_v_add_f16_e32_vdst_hi_half(self):
"""V_ADD_F16_E32 writing to hi-half destination (v[128]+).
When vdst >= 384 (representing v[128]+), the hardware writes to bits [31:16]
of v[vdst-128] while preserving bits [15:0]. The emulator must merge the result.
Regression test for: VOP2 f16 vdst hi-half write bug.
"""
instructions = [
# v[0] = 0x0000_BEEF: lo has marker value
s_mov_b32(s[0], 0x0000BEEF),
v_mov_b32_e32(v[0], s[0]),
# v[1] = f16(1.0), v[2] = f16(2.0)
s_mov_b32(s[1], 0x3c00),
s_mov_b32(s[2], 0x4000),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
# v_add_f16_e32 v[128], v[1], v[2] (vdst=v[128] writes hi of v[0])
# v[0].hi = 1.0 + 2.0 = 3.0, v[0].lo preserved = 0xBEEF
VOP2(VOP2Op.V_ADD_F16, vdst=v[128], src0=v[1], vsrc1=v[2]),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][0] >> 16) & 0xffff
lo = st.vgpr[0][0] & 0xffff
# hi = 3.0 = 0x4200, lo preserved = 0xBEEF
self.assertEqual(hi, 0x4200, f"Expected hi=f16(3.0)=0x4200, got 0x{hi:04x}")
self.assertEqual(lo, 0xBEEF, f"Expected lo preserved=0xBEEF, got 0x{lo:04x}")
def test_v_mul_f16_e32_vdst_hi_half(self):
"""V_MUL_F16_E32 writing to hi-half destination.
Regression test for: VOP2 f16 vdst hi-half write bug.
"""
instructions = [
# v[0] = 0x0000_DEAD: lo has marker value
s_mov_b32(s[0], 0x0000DEAD),
v_mov_b32_e32(v[0], s[0]),
# v[1] = f16(2.0), v[2] = f16(4.0)
s_mov_b32(s[1], 0x4000),
s_mov_b32(s[2], 0x4400),
v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]),
# v_mul_f16_e32 v[128], v[1], v[2] (vdst=v[128] writes hi of v[0])
# v[0].hi = 2.0 * 4.0 = 8.0, v[0].lo preserved = 0xDEAD
VOP2(VOP2Op.V_MUL_F16, vdst=v[128], src0=v[1], vsrc1=v[2]),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][0] >> 16) & 0xffff
lo = st.vgpr[0][0] & 0xffff
# hi = 8.0 = 0x4800, lo preserved = 0xDEAD
self.assertEqual(hi, 0x4800, f"Expected hi=f16(8.0)=0x4800, got 0x{hi:04x}")
self.assertEqual(lo, 0xDEAD, f"Expected lo preserved=0xDEAD, got 0x{lo:04x}")
def test_v_add_f16_e32_both_hi_half(self):
"""V_ADD_F16_E32 with both vsrc1 and vdst as hi-half (different underlying regs).
Tests the combination of both fixes: reading vsrc1 from hi-half AND
writing result to hi-half destination, using different underlying VGPRs.
Regression test for: VOP2 f16 hi-half bugs (combined).
"""
instructions = [
# v[0] = 0x4000_xxxx: hi=f16(2.0) for vsrc1
s_mov_b32(s[0], 0x40000000),
v_mov_b32_e32(v[0], s[0]),
# v[1] = 0x0000_3c00: lo=f16(1.0) for src0
s_mov_b32(s[1], 0x00003c00),
v_mov_b32_e32(v[1], s[1]),
# v[2] = 0x0000_CAFE: lo=marker for vdst preservation
s_mov_b32(s[2], 0x0000CAFE),
v_mov_b32_e32(v[2], s[2]),
# v_add_f16_e32 v[130], v[1], v[128]
# src0 = v[1].lo = 1.0
# vsrc1 = v[128] reads v[0].hi = 2.0
# result = 1.0 + 2.0 = 3.0
# vdst = v[130] writes to v[2].hi, preserving v[2].lo
VOP2(VOP2Op.V_ADD_F16, vdst=v[130], src0=v[1], vsrc1=v[128]),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][2] >> 16) & 0xffff
lo = st.vgpr[0][2] & 0xffff
# hi = 3.0 = 0x4200, lo preserved = 0xCAFE
self.assertEqual(hi, 0x4200, f"Expected hi=f16(3.0)=0x4200, got 0x{hi:04x}")
self.assertEqual(lo, 0xCAFE, f"Expected lo preserved=0xCAFE, got 0x{lo:04x}")
def test_v_fmac_f16_e32_vsrc1_hi_half(self):
"""V_FMAC_F16_E32 with vsrc1 from hi-half.
V_FMAC_F16: vdst = vdst + src0 * vsrc1
Regression test for: VOP2 f16 vsrc1 hi-half extraction bug.
"""
instructions = [
# v[0] = 0x4000_3c00: hi=f16(2.0), lo=f16(1.0)
s_mov_b32(s[0], 0x40003c00),
v_mov_b32_e32(v[0], s[0]),
# v[1] = f16(3.0) = 0x4200
s_mov_b32(s[1], 0x4200),
v_mov_b32_e32(v[1], s[1]),
# v_fmac_f16_e32 v[1], v[0], v[128]
# vdst = v[1] = 3.0 + v[0].lo * v[0].hi = 3.0 + 1.0 * 2.0 = 5.0
VOP2(VOP2Op.V_FMAC_F16, vdst=v[1], src0=v[0], vsrc1=v[128]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xffff
# 3.0 + 1.0 * 2.0 = 5.0, f16 5.0 = 0x4500
self.assertEqual(result, 0x4500, f"Expected f16(5.0)=0x4500, got 0x{result:04x}")
def test_v_fmac_f16_e32_vdst_hi_half(self):
"""V_FMAC_F16_E32 writing to hi-half destination.
V_FMAC_F16: vdst.h = vdst.h + src0 * vsrc1
When vdst is v[128]+, the accumulator D0 must also read from the hi-half.
This tests the bug where D0 was read from lo-half instead of hi-half.
Regression test for: VOP2 FMAC hi-half D0 accumulator read bug.
"""
instructions = [
# v[0] = 0x3800_DEAD: hi=f16(0.5), lo=marker (0xDEAD)
s_mov_b32(s[0], 0x3800DEAD),
v_mov_b32_e32(v[0], s[0]),
# v[1] = f16(2.0) = 0x4000
s_mov_b32(s[1], 0x4000),
v_mov_b32_e32(v[1], s[1]),
# v[2] = f16(3.0) = 0x4200
s_mov_b32(s[2], 0x4200),
v_mov_b32_e32(v[2], s[2]),
# v_fmac_f16_e32 v[128], v[1], v[2]
# vdst = v[128] means v[0].hi
# D0 = v[0].hi = 0.5
# result = D0 + src0 * vsrc1 = 0.5 + 2.0 * 3.0 = 6.5
# v[0].hi = 6.5, v[0].lo preserved = 0xDEAD
VOP2(VOP2Op.V_FMAC_F16, vdst=v[128], src0=v[1], vsrc1=v[2]),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][0] >> 16) & 0xffff
lo = st.vgpr[0][0] & 0xffff
# hi = 6.5 = 0x4680, lo preserved = 0xDEAD
self.assertEqual(hi, 0x4680, f"Expected hi=f16(6.5)=0x4680, got 0x{hi:04x}")
self.assertEqual(lo, 0xDEAD, f"Expected lo preserved=0xDEAD, got 0x{lo:04x}")
def test_v_mul_f16_e32_src0_hi_half(self):
"""V_MUL_F16_E32 with src0 from hi-half (src0 >= v[128]).
When src0 >= 384 (representing v[128]+), the hardware reads from the hi 16 bits
of v[src0-128]. The emulator must extract bits [31:16] from the actual VGPR.
Regression test for: VOP2 f16 src0 hi-half extraction bug.
"""
instructions = [
# v[0] = 0x4000_3c00: hi=f16(2.0), lo=f16(1.0)
s_mov_b32(s[0], 0x40003c00),
v_mov_b32_e32(v[0], s[0]),
# v[1] = f16(3.0) = 0x4200
s_mov_b32(s[1], 0x4200),
v_mov_b32_e32(v[1], s[1]),
# v_mul_f16_e32 v[2], v[128], v[1]
# src0 = v[128] reads from v[0].hi = 2.0
# result = 2.0 * 3.0 = 6.0
VOP2(VOP2Op.V_MUL_F16, vdst=v[2], src0=v[128], vsrc1=v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 * 3.0 = 6.0, f16 6.0 = 0x4600
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
def test_v_add_f16_e32_src0_hi_half(self):
"""V_ADD_F16_E32 with src0 from hi-half (src0 >= v[128]).
Regression test for: VOP2 f16 src0 hi-half extraction bug.
"""
instructions = [
# v[0] = 0x4000_3c00: hi=f16(2.0), lo=f16(1.0)
s_mov_b32(s[0], 0x40003c00),
v_mov_b32_e32(v[0], s[0]),
# v[1] = f16(5.0) = 0x4500
s_mov_b32(s[1], 0x4500),
v_mov_b32_e32(v[1], s[1]),
# v_add_f16_e32 v[2], v[128], v[1]
# src0 = v[128] reads from v[0].hi = 2.0
# result = 2.0 + 5.0 = 7.0
VOP2(VOP2Op.V_ADD_F16, vdst=v[2], src0=v[128], vsrc1=v[1]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
# 2.0 + 5.0 = 7.0, f16 7.0 = 0x4700
self.assertEqual(result, 0x4700, f"Expected f16(7.0)=0x4700, got 0x{result:04x}")
class TestF16InlineConstants(unittest.TestCase):
"""Regression tests for VOP2 F16 inline float constants.
For 16-bit VOP2 operations (v_add_f16, v_mul_f16, etc.), inline float constants
like 1.0, 2.0 must use F16 encoding (0x3c00, 0x4000) not F32 encoding (0x3f800000).
The emulator's rsrc() function needs bits=16 to select F16_INLINE constants.
Regression test for: VOP2 16-bit inline constant using F32 instead of F16.
"""
def test_v_add_f16_inline_constant_1_0(self):
"""V_ADD_F16_E32 with inline constant 1.0 should use F16 encoding."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 1.0
v_mov_b32_e32(v[0], s[0]),
# v_add_f16_e32 v[1], 1.0, v[0] -- 1.0 must be F16 0x3c00, not F32 0x3f800000
v_add_f16_e32(v[1], 1.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xFFFF
# 1.0 + 1.0 = 2.0, f16 2.0 = 0x4000
self.assertEqual(result, 0x4000, f"Expected f16(2.0)=0x4000, got 0x{result:04x}")
def test_v_add_f16_inline_constant_2_0(self):
"""V_ADD_F16_E32 with inline constant 2.0."""
instructions = [
s_mov_b32(s[0], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_add_f16_e32(v[1], 2.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xFFFF
# 2.0 + 3.0 = 5.0, f16 5.0 = 0x4500
self.assertEqual(result, 0x4500, f"Expected f16(5.0)=0x4500, got 0x{result:04x}")
def test_v_mul_f16_inline_constant(self):
"""V_MUL_F16_E32 with inline constant 2.0."""
instructions = [
s_mov_b32(s[0], 0x4200), # f16 3.0
v_mov_b32_e32(v[0], s[0]),
v_mul_f16_e32(v[1], 2.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][1] & 0xFFFF
# 2.0 * 3.0 = 6.0, f16 6.0 = 0x4600
self.assertEqual(result, 0x4600, f"Expected f16(6.0)=0x4600, got 0x{result:04x}")
class TestCndmask(unittest.TestCase):
"""Tests for V_CNDMASK_B32 and V_CNDMASK_B16."""
def test_v_cndmask_b16_select_src0(self):
"""V_CNDMASK_B16 selects src0 when VCC bit is 0."""
instructions = [
s_mov_b32(VCC_LO, 0), # VCC = 0
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cndmask_b16(v[2], v[0], v[1], VCC),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x3c00, f"Expected src0=0x3c00, got 0x{result:04x}")
def test_v_cndmask_b16_select_src1(self):
"""V_CNDMASK_B16 selects src1 when VCC bit is 1."""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1
s_mov_b32(s[0], 0x3c00), # f16 1.0
s_mov_b32(s[1], 0x4000), # f16 2.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cndmask_b16(v[2], v[0], v[1], VCC),
]
st = run_program(instructions, n_lanes=1)
result = st.vgpr[0][2] & 0xffff
self.assertEqual(result, 0x4000, f"Expected src1=0x4000, got 0x{result:04x}")
def test_v_cndmask_b16_write_hi(self):
"""V_CNDMASK_B16 can write to high 16 bits with opsel."""
instructions = [
s_mov_b32(s[0], 0x3c003800), # src0: hi=1.0, lo=0.5
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x4000c000), # src1: hi=2.0, lo=-2.0
v_mov_b32_e32(v[1], s[1]),
s_mov_b32(s[2], 0xDEAD0000), # v2 initial: hi=0xDEAD, lo=0
v_mov_b32_e32(v[2], s[2]),
s_mov_b32(VCC_LO, 0), # vcc = 0, select src0
# opsel=0b1011: bit0=src0 hi, bit1=src1 hi, bit3=dst hi
VOP3(VOP3Op.V_CNDMASK_B16, vdst=v[2], src0=v[0], src1=v[1], src2=SrcEnum.VCC_LO, opsel=0b1011),
]
st = run_program(instructions, n_lanes=1)
hi = (st.vgpr[0][2] >> 16) & 0xffff
lo = st.vgpr[0][2] & 0xffff
# vcc=0 selects src0.h = 1.0 = 0x3c00, writes to hi
self.assertEqual(hi, 0x3c00, f"Expected hi=0x3c00 (1.0), got 0x{hi:04x}")
self.assertEqual(lo, 0x0000, f"Expected lo preserved as 0, got 0x{lo:04x}")
class TestSpecialFloatValues(unittest.TestCase):
"""Tests for special float value handling in VOP2 instructions."""
def test_neg_zero_add(self):
"""-0.0 + 0.0 = +0.0 (IEEE 754)."""
neg_zero = 0x80000000
instructions = [
s_mov_b32(s[0], neg_zero),
v_mov_b32_e32(v[0], s[0]),
v_add_f32_e32(v[1], 0.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
def test_neg_zero_mul(self):
"""-0.0 * -1.0 = +0.0."""
neg_zero = 0x80000000
instructions = [
s_mov_b32(s[0], neg_zero),
v_mov_b32_e32(v[0], s[0]),
v_mul_f32_e32(v[1], -1.0, v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000, "Should be +0.0")
def test_inf_minus_inf(self):
"""+inf - inf = NaN."""
import math
pos_inf = 0x7f800000
neg_inf = 0xff800000
instructions = [
s_mov_b32(s[0], pos_inf),
s_mov_b32(s[1], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_sub_f32_e32(v[2], v[0], v[1]), # inf - (-inf) = inf
v_add_f32_e32(v[3], v[0], v[1]), # inf + (-inf) = NaN
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], pos_inf, "inf - (-inf) = inf")
self.assertTrue(math.isnan(i2f(st.vgpr[0][3])), "inf + (-inf) = NaN")
def test_denormal_f32_mul_ftz(self):
"""Denormal * normal - RDNA3 flushes denormals to zero (FTZ mode)."""
smallest_denorm = 0x00000001 # Smallest positive denormal
instructions = [
s_mov_b32(s[0], smallest_denorm),
v_mov_b32_e32(v[0], s[0]),
v_mul_f32_e32(v[1], 2.0, v[0]), # Denormal input gets flushed to 0
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x00000000)
class TestCarryOps(unittest.TestCase):
"""Tests for VOP2 carry instructions (v_add_co_ci_u32, v_sub_co_ci_u32, v_subrev_co_ci_u32)."""
def test_v_subrev_co_ci_u32_no_borrow(self):
"""V_SUBREV_CO_CI_U32: D0 = S1 - S0 - VCC_IN, when VCC_IN=0."""
instructions = [
s_mov_b32(VCC_LO, 0), # VCC = 0 (no borrow in)
v_mov_b32_e32(v[0], 5), # S0 = 5
v_mov_b32_e32(v[1], 10), # S1 = 10
v_subrev_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 10 - 5 - 0 = 5
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vcc, 0) # No borrow out
def test_v_subrev_co_ci_u32_with_borrow(self):
"""V_SUBREV_CO_CI_U32: D0 = S1 - S0 - VCC_IN, when VCC_IN=1."""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (borrow in)
v_mov_b32_e32(v[0], 5), # S0 = 5
v_mov_b32_e32(v[1], 10), # S1 = 10
v_subrev_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 10 - 5 - 1 = 4
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 4)
self.assertEqual(st.vcc, 0) # No borrow out
def test_v_subrev_co_ci_u32_generates_borrow(self):
"""V_SUBREV_CO_CI_U32: generates borrow when S0 + VCC_IN > S1."""
instructions = [
s_mov_b32(VCC_LO, 0), # VCC = 0
v_mov_b32_e32(v[0], 10), # S0 = 10
v_mov_b32_e32(v[1], 5), # S1 = 5
v_subrev_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 5 - 10 - 0 = -5 (underflow)
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0xFFFFFFFB) # -5 as unsigned
self.assertEqual(st.vcc, 1) # Borrow out
def test_v_add_co_ci_u32_no_carry(self):
"""V_ADD_CO_CI_U32: D0 = S0 + S1 + VCC_IN, when VCC_IN=0."""
instructions = [
s_mov_b32(VCC_LO, 0), # VCC = 0 (no carry in)
v_mov_b32_e32(v[0], 5), # S0 = 5
v_mov_b32_e32(v[1], 10), # S1 = 10
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 5 + 10 + 0 = 15
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 15)
self.assertEqual(st.vcc, 0) # No carry out
def test_v_add_co_ci_u32_with_carry(self):
"""V_ADD_CO_CI_U32: D0 = S0 + S1 + VCC_IN, when VCC_IN=1."""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
v_mov_b32_e32(v[0], 5), # S0 = 5
v_mov_b32_e32(v[1], 10), # S1 = 10
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 5 + 10 + 1 = 16
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 16)
self.assertEqual(st.vcc, 0) # No carry out
def test_v_add_co_ci_u32_generates_carry(self):
"""V_ADD_CO_CI_U32: generates carry when overflow occurs."""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
s_mov_b32(s[0], 0xFFFFFFFF), # max u32
v_mov_b32_e32(v[0], s[0]), # S0 = 0xFFFFFFFF
v_mov_b32_e32(v[1], 0), # S1 = 0
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 0xFFFFFFFF + 0 + 1 = 0 (overflow)
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0) # Overflowed to 0
self.assertEqual(st.vcc, 1) # Carry out
def test_v_add_co_ci_u32_clears_carry(self):
"""V_ADD_CO_CI_U32: VCC must be updated even when no carry is generated.
This tests the case where VCC=1 going in (carry-in consumed) but the addition
does not overflow, so VCC must be cleared to 0.
Regression test for: VCC not being written by v_add_co_ci_u32_e32.
"""
instructions = [
s_mov_b32(VCC_LO, 1), # VCC = 1 (carry in)
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 3) # 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out - VCC must be cleared
def test_v_add_co_ci_u32_multilane_clears_vcc(self):
"""V_ADD_CO_CI_U32 with multiple lanes: VCC bits must be updated per-lane.
When VCC has multiple bits set (one per active lane), and the addition doesn't
overflow for any lane, all VCC bits must be cleared.
Regression test for: VCC not being written by v_add_co_ci_u32_e32 in multi-lane case.
"""
instructions = [
s_mov_b32(VCC_LO, 0b11), # VCC = 0b11 (lanes 0,1 have carry-in)
v_mov_b32_e32(v[0], 1), # S0 = 1 for all lanes
v_mov_b32_e32(v[1], 1), # S1 = 1 for all lanes
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 1 = 3 (no overflow)
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vgpr[0][2], 3) # lane 0: 1 + 1 + 1 = 3
self.assertEqual(st.vgpr[1][2], 3) # lane 1: 1 + 1 + 1 = 3
self.assertEqual(st.vcc, 0) # No carry out for any lane - all VCC bits must be cleared
def test_v_add_co_ci_u32_preserves_inactive_vcc_bits(self):
"""V_ADD_CO_CI_U32: VCC carry-out overwrites entire VCC register.
VOP2 carry instructions write ALL VCC bits based on carry-out, clearing
bits for lanes that don't overflow regardless of EXEC mask.
Note: This differs from VOPC which only writes active lane bits.
"""
instructions = [
s_mov_b32(VCC_LO, 0x00010000), # VCC bit 16 set
v_mov_b32_e32(v[0], 1), # S0 = 1
v_mov_b32_e32(v[1], 1), # S1 = 1
v_add_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 1 + 1 + 0 = 2 (no carry)
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vgpr[0][2], 2) # lane 0: 1 + 1 + 0 = 2
# VCC should be completely cleared (all lanes have no carry-out)
self.assertEqual(st.vcc, 0)
def test_v_add_co_ci_u32_all_lanes_same_result(self):
"""V_ADD_CO_CI_U32: all active lanes should produce the same result.
When the same constant inputs are used across all lanes, each lane should
compute the same result and write to its own VGPR slot.
Regression test for: VGPR writes not happening for all lanes.
"""
instructions = [
s_mov_b32(VCC_LO, 0), # No carry-in
v_mov_b32_e32(v[0], 3), # inline constant 3
v_mov_b32_e32(v[1], 5), # value 5
v_add_co_ci_u32_e32(v[1], 3, v[1]), # v[1] = 3 + v[1] + 0 = 3 + 5 = 8
]
st = run_program(instructions, n_lanes=4)
# All 4 lanes should have v[1] = 8
for lane in range(4):
self.assertEqual(st.vgpr[lane][1], 8, f"lane {lane} should have v[1]=8")
def test_v_sub_co_ci_u32_no_borrow(self):
"""V_SUB_CO_CI_U32: D0 = S0 - S1 - VCC_IN, when VCC_IN=0."""
instructions = [
s_mov_b32(VCC_LO, 0), # VCC = 0 (no borrow in)
v_mov_b32_e32(v[0], 10), # S0 = 10
v_mov_b32_e32(v[1], 5), # S1 = 5
v_sub_co_ci_u32_e32(v[2], v[0], v[1]), # D0 = 10 - 5 - 0 = 5
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 5)
self.assertEqual(st.vcc, 0) # No borrow out
def test_v_sub_co_ci_u32_vop3sd_separate_carry_regs(self):
"""VOP3SD V_SUB_CO_CI_U32: carry-in from src2, carry-out to sdst (separate registers).
This tests the VOP3SD encoding where src2 specifies the carry-in register
independently from sdst (carry-out). The bug was reading carry-in from sdst
instead of src2.
Computation: D0 = S0 - S1 - carry_in = 0 - 0 - 1 = -1 = 0xFFFFFFFF
"""
instructions = [
s_mov_b32(s[6], 1), # carry-in = 1 (in s[6])
s_mov_b32(s[10], 0), # carry-out dest = 0 initially (in s[10])
# VOP3SD: v_sub_co_ci_u32(vdst, sdst, src0, src1, src2)
# src2 is carry-in (s[6]=1), sdst is carry-out (s[10])
v_sub_co_ci_u32(v[0], s[10], 0, 0, s[6]), # D0 = 0 - 0 - 1 = -1
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0xFFFFFFFF) # -1 as unsigned
self.assertEqual(st.sgpr[10], 1) # Borrow out to s[10]
def test_v_add_co_ci_u32_vop3sd_separate_carry_regs(self):
"""VOP3SD V_ADD_CO_CI_U32: carry-in from src2, carry-out to sdst (separate registers).
This tests the VOP3SD encoding where src2 specifies the carry-in register
independently from sdst (carry-out).
Computation: D0 = S0 + S1 + carry_in = 5 + 10 + 1 = 16
"""
instructions = [
s_mov_b32(s[6], 1), # carry-in = 1 (in s[6])
s_mov_b32(s[10], 0), # carry-out dest = 0 initially (in s[10])
# VOP3SD: v_add_co_ci_u32(vdst, sdst, src0, src1, src2)
v_add_co_ci_u32(v[0], s[10], 5, 10, s[6]), # D0 = 5 + 10 + 1 = 16
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 16)
self.assertEqual(st.sgpr[10], 0) # No carry out
def test_v_add_co_ci_u32_vop3sd_null_sdst(self):
"""VOP3SD V_ADD_CO_CI_U32 with sdst=NULL: carry output is discarded.
When sdst=NULL (register 124), the carry-out should NOT be written anywhere.
We verify this by checking that VCC (which we set to a sentinel value) is unchanged.
"""
instructions = [
s_mov_b32(VCC_LO, 0xDEADBEEF), # Sentinel value in VCC
s_mov_b32(s[6], 0), # carry-in = 0
# VOP3SD with NULL sdst: carry-out should be discarded
# Uses 0xFFFFFFFF + 1 + 0 = 0 with carry-out=1, but carry should not be written
v_add_co_ci_u32(v[0], NULL, 0xFFFFFFFF, 1, s[6]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][0], 0) # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
self.assertEqual(st.vcc, 0xDEADBEEF) # VCC unchanged - carry was discarded
if __name__ == '__main__':
unittest.main()

3656
tinygrad_repo/test/amd/hw/test_vop3.py Normal file
View File

File diff suppressed because it is too large Load Diff

1036
tinygrad_repo/test/amd/hw/test_vop3p.py Normal file
View File

File diff suppressed because it is too large Load Diff

964
tinygrad_repo/test/amd/hw/test_vopc.py Normal file
View File

@@ -0,0 +1,964 @@
"""Tests for VOPC instructions - vector compare operations.
Includes: v_cmp_class_f32, v_cmp_class_f16, v_cmp_eq_*, v_cmp_lt_*, v_cmp_gt_*
"""
import unittest
from test.amd.hw.helpers import *
VCC = 106 # SGPR index for VCC_LO
class TestCmpClass(unittest.TestCase):
"""Tests for V_CMP_CLASS_F32 float classification."""
def test_cmp_class_quiet_nan(self):
"""V_CMP_CLASS_F32 detects quiet NaN."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
def test_cmp_class_signaling_nan(self):
"""V_CMP_CLASS_F32 detects signaling NaN."""
signal_nan = 0x7f800001
instructions = [
s_mov_b32(s[0], signal_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
def test_cmp_class_positive_inf(self):
"""V_CMP_CLASS_F32 detects +inf."""
pos_inf = 0x7f800000
instructions = [
s_mov_b32(s[0], pos_inf),
s_mov_b32(s[1], 0b1000000000), # bit 9 = +inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect +inf")
def test_cmp_class_negative_inf(self):
"""V_CMP_CLASS_F32 detects -inf."""
neg_inf = 0xff800000
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000100), # bit 2 = -inf
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect -inf")
def test_cmp_class_normal_positive(self):
"""V_CMP_CLASS_F32 detects positive normal."""
instructions = [
v_mov_b32_e32(v[0], 1.0),
s_mov_b32(s[1], 0b0100000000), # bit 8 = positive normal
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
def test_cmp_class_normal_negative(self):
"""V_CMP_CLASS_F32 detects negative normal."""
instructions = [
v_mov_b32_e32(v[0], -1.0),
v_mov_b32_e32(v[1], 0b0000001000), # bit 3 = negative normal
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
def test_cmp_class_quiet_nan_not_signaling(self):
"""Quiet NaN does not match signaling NaN mask."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000001), # bit 0 = signaling NaN only
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Quiet NaN should not match signaling mask")
def test_cmp_class_signaling_nan_not_quiet(self):
"""Signaling NaN does not match quiet NaN mask."""
signal_nan = 0x7f800001
instructions = [
s_mov_b32(s[0], signal_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0b0000000010), # bit 1 = quiet NaN only
v_cmp_class_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")
def test_v_cmp_lg_f32_nan(self):
"""v_cmp_lg_f32 is ordered not-equal (<>): NaN <> x should be False per IEEE 754."""
quiet_nan = 0x7fc00000
one_f32 = 0x3f800000 # 1.0f
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], one_f32),
v_mov_b32_e32(v[1], s[1]),
v_cmp_lg_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "v_cmp_lg_f32(NaN, 1.0) should be 0")
def test_v_cmp_neq_f32_nan(self):
"""v_cmp_neq_f32 is unordered not-equal (!=): NaN != x should be True per IEEE 754."""
quiet_nan = 0x7fc00000
one_f32 = 0x3f800000 # 1.0f
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], one_f32),
v_mov_b32_e32(v[1], s[1]),
v_cmp_neq_f32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "v_cmp_neq_f32(NaN, 1.0) should be 1")
def test_v_cmp_sets_vcc_bits(self):
"""V_CMP_EQ sets VCC bits based on per-lane comparison."""
instructions = [
s_mov_b32(s[0], 5),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
class TestCmpClassF16(unittest.TestCase):
"""Tests for V_CMP_CLASS_F16 float classification.
Class bit mapping:
bit 0 = signaling NaN
bit 1 = quiet NaN
bit 2 = -infinity
bit 3 = -normal
bit 4 = -denormal
bit 5 = -zero
bit 6 = +zero
bit 7 = +denormal
bit 8 = +normal
bit 9 = +infinity
"""
def test_cmp_class_f16_positive_zero(self):
"""V_CMP_CLASS_F16: +zero matches bit 6."""
instructions = [
v_mov_b32_e32(v[0], 0x0000), # f16 +0.0
v_mov_b32_e32(v[1], 0x40), # bit 6 = +zero
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive zero")
def test_cmp_class_f16_negative_zero(self):
"""V_CMP_CLASS_F16: -zero matches bit 5."""
instructions = [
s_mov_b32(s[0], 0x8000), # f16 -0.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x20), # bit 5 = -zero
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative zero")
def test_cmp_class_f16_positive_normal(self):
"""V_CMP_CLASS_F16: +1.0 (normal) matches bit 8."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 +1.0
s_mov_b32(s[1], 0x100), # bit 8 = +normal
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive normal")
def test_cmp_class_f16_negative_normal(self):
"""V_CMP_CLASS_F16: -1.0 (normal) matches bit 3."""
instructions = [
s_mov_b32(s[0], 0xbc00), # f16 -1.0
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x08), # bit 3 = -normal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative normal")
def test_cmp_class_f16_positive_infinity(self):
"""V_CMP_CLASS_F16: +inf matches bit 9."""
instructions = [
s_mov_b32(s[0], 0x7c00), # f16 +inf
s_mov_b32(s[1], 0x200), # bit 9 = +inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive infinity")
def test_cmp_class_f16_negative_infinity(self):
"""V_CMP_CLASS_F16: -inf matches bit 2."""
instructions = [
s_mov_b32(s[0], 0xfc00), # f16 -inf
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x04), # bit 2 = -inf
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative infinity")
def test_cmp_class_f16_quiet_nan(self):
"""V_CMP_CLASS_F16: quiet NaN matches bit 1."""
instructions = [
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x02), # bit 1 = quiet NaN
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect quiet NaN")
def test_cmp_class_f16_signaling_nan(self):
"""V_CMP_CLASS_F16: signaling NaN matches bit 0."""
instructions = [
s_mov_b32(s[0], 0x7c01), # f16 signaling NaN
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x01), # bit 0 = signaling NaN
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect signaling NaN")
def test_cmp_class_f16_positive_denormal(self):
"""V_CMP_CLASS_F16: positive denormal matches bit 7."""
instructions = [
v_mov_b32_e32(v[0], 1), # f16 +denormal (0x0001)
v_mov_b32_e32(v[1], 0x80), # bit 7 = +denormal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect positive denormal")
def test_cmp_class_f16_negative_denormal(self):
"""V_CMP_CLASS_F16: negative denormal matches bit 4."""
instructions = [
s_mov_b32(s[0], 0x8001), # f16 -denormal
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], 0x10), # bit 4 = -denormal
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Should detect negative denormal")
def test_cmp_class_f16_combined_mask_zeros(self):
"""V_CMP_CLASS_F16: mask 0x60 covers both +zero and -zero."""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
v_mov_b32_e32(v[1], 0x60), # bits 5 and 6 (+-zero)
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x60")
def test_cmp_class_f16_combined_mask_1f8(self):
"""V_CMP_CLASS_F16: mask 0x1f8 covers -normal,-denorm,-zero,+zero,+denorm,+normal.
This is the exact mask used in the f16 sin kernel at PC=46.
"""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
s_mov_b32(s[0], 0x1f8),
v_mov_b32_e32(v[1], s[0]), # mask 0x1f8
v_cmp_class_f16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with mask 0x1f8")
def test_cmp_class_f16_vop3_encoding(self):
"""V_CMP_CLASS_F16 in VOP3 encoding (v_cmp_class_f16_e64)."""
instructions = [
v_mov_b32_e32(v[0], 0), # f16 +0.0
s_mov_b32(s[0], 0x1f8), # class mask
v_cmp_class_f16_e64(VCC_LO, v[0], s[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +zero with VOP3 encoding")
def test_cmp_class_f16_vop3_normal_positive(self):
"""V_CMP_CLASS_F16 VOP3 encoding with +1.0 (normal)."""
instructions = [
s_mov_b32(s[0], 0x3c00), # f16 +1.0
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "VCC should be 1 for +1.0 (normal) with mask 0x1f8")
def test_cmp_class_f16_vop3_nan_fails_mask(self):
"""V_CMP_CLASS_F16 VOP3: NaN should NOT match mask 0x1f8 (no NaN bits set)."""
instructions = [
s_mov_b32(s[0], 0x7e00), # f16 quiet NaN
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for NaN with mask 0x1f8 (no NaN bits)")
def test_cmp_class_f16_vop3_inf_fails_mask(self):
"""V_CMP_CLASS_F16 VOP3: +inf should NOT match mask 0x1f8 (no inf bits set)."""
instructions = [
s_mov_b32(s[0], 0x7c00), # f16 +inf
v_mov_b32_e32(v[0], s[0]),
s_mov_b32(s[1], 0x1f8), # class mask
v_cmp_class_f16_e64(VCC_LO, v[0], s[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "VCC should be 0 for +inf with mask 0x1f8 (no inf bits)")
class TestCmpInt(unittest.TestCase):
"""Tests for integer comparison operations."""
def test_v_cmp_eq_u32(self):
"""V_CMP_EQ_U32 sets VCC bits based on per-lane comparison."""
instructions = [
s_mov_b32(s[0], 5),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should match")
def test_v_cmp_ne_u32_with_zero(self):
"""V_CMP_NE_U32: compare with zero, used for int->bool cast."""
instructions = [
v_mov_b32_e32(v[1], 0),
v_cmp_eq_u32_e32(1, v[255]), # vcc = (lane == 1)
v_cndmask_b32_e64(v[1], v[1], 1, VCC_LO), # v1[lane1] = 1
v_cmp_ne_u32_e32(0, v[1]), # vcc = (0 != v1)
v_cndmask_b32_e64(v[0], 0, 1, VCC_LO), # v0 = vcc ? 1 : 0
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vgpr[0][0], 0, "lane 0: 0 != 0 should be false")
self.assertEqual(st.vgpr[1][0], 1, "lane 1: 0 != 1 should be true")
self.assertEqual(st.vcc & 0x3, 0x2, "VCC should be 0b10")
def test_v_cmp_ne_u32_all_nonzero(self):
"""V_CMP_NE_U32: all lanes have nonzero values."""
instructions = [
v_mov_b32_e32(v[1], 5),
v_cmp_ne_u32_e32(0, v[1]),
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.vcc & 0xf, 0xf, "All lanes should be != 0")
def test_cmp_eq_u16_opsel_lo_lo(self):
"""V_CMP_EQ_U16 comparing lo halves."""
instructions = [
s_mov_b32(s[0], 0x12340005), # lo=5, hi=0x1234
s_mov_b32(s[1], 0xABCD0005), # lo=5, hi=0xABCD
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[1]),
v_cmp_eq_u16_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Lo halves should be equal")
def test_cmp_eq_u16_opsel_hi_hi(self):
"""V_CMP_EQ_U16 comparing hi halves with VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x00051234), # hi=5, lo=0x1234
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x0005ABCD), # hi=5, lo=0xABCD
v_mov_b32_e32(v[1], s[2]),
v_cmp_eq_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "Hi halves should be equal: 5==5")
def test_cmp_eq_u16_opsel_hi_hi_equal(self):
"""V_CMP_EQ_U16 VOP3 with opsel=3 compares hi halves (equal case)."""
instructions = [
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x12340009), # lo=9, hi=0x1234
v_mov_b32_e32(v[1], s[2]),
v_cmp_eq_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "hi==hi should be true: 0x1234==0x1234")
def test_cmp_gt_u16_opsel_hi(self):
"""V_CMP_GT_U16 VOP3 with opsel=3 compares hi halves."""
instructions = [
s_mov_b32(s[2], 0x99990005), # lo=5, hi=0x9999
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x12340005), # lo=5, hi=0x1234
v_mov_b32_e32(v[1], s[2]),
v_cmp_gt_u16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "hi>hi should be true: 0x9999>0x1234")
class TestCmpFloat(unittest.TestCase):
"""Tests for float comparison operations."""
def test_v_cmp_lt_f16_vsrc1_hi(self):
"""V_CMP_LT_F16 with both operands from high half using VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
v_mov_b32_e32(v[1], s[2]),
v_cmp_lt_f16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "1.0 < 2.0 should be true")
def test_v_cmp_gt_f16_vsrc1_hi(self):
"""V_CMP_GT_F16 with both operands from high half using VOP3 opsel."""
instructions = [
s_mov_b32(s[2], 0x40000000), # hi=2.0 (f16), lo=0
v_mov_b32_e32(v[0], s[2]),
s_mov_b32(s[2], 0x3c000000), # hi=1.0 (f16), lo=0
v_mov_b32_e32(v[1], s[2]),
v_cmp_gt_f16_e64(vdst=s[0], src0=v[0], src1=v[1], opsel=3),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[0] & 1, 1, "2.0 > 1.0 should be true")
def test_v_cmp_eq_f16_vsrc1_hi_equal(self):
"""v_cmp_eq_f16 with equal low and high halves."""
instructions = [
s_mov_b32(s[0], 0x42004200), # hi=3.0 (0x4200), lo=3.0 (0x4200)
v_mov_b32_e32(v[0], s[0]),
v_cmp_eq_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (3.0 == 3.0)")
def test_v_cmp_neq_f16_vsrc1_hi(self):
"""v_cmp_neq_f16 with different low and high halves."""
instructions = [
s_mov_b32(s[0], 0x40003c00), # hi=2.0 (0x4000), lo=1.0 (0x3c00)
v_mov_b32_e32(v[0], s[0]),
v_cmp_lg_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")
def test_v_cmp_nge_f16_inf_self(self):
"""v_cmp_nge_f16 comparing -inf with itself (unordered less than).
Regression test: -inf < -inf should be false (IEEE 754).
"""
instructions = [
s_mov_b32(s[0], 0xFC00FC00), # both halves = -inf (0xFC00)
v_mov_b32_e32(v[0], s[0]),
v_cmp_nge_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "Expected vcc=0 (-inf >= -inf)")
def test_v_cmp_f16_multilane(self):
"""v_cmp_lt_f16 with vsrc1=v128 across multiple lanes."""
instructions = [
# Lane 0: v0 = 0x40003c00 (hi=2.0, lo=1.0) -> 1.0 < 2.0 = true
# Lane 1: v0 = 0x3c004000 (hi=1.0, lo=2.0) -> 2.0 < 1.0 = false
v_mov_b32_e32(v[0], 0x40003c00), # default
v_cmp_eq_u32_e32(1, v[255]), # vcc = (lane == 1)
v_cndmask_b32_e64(v[0], v[0], 0x3c004000, SrcEnum.VCC_LO),
v_cmp_lt_f16_e32(v[0], v[0].h),
]
st = run_program(instructions, n_lanes=2)
self.assertEqual(st.vcc & 1, 1, "Lane 0: expected vcc=1 (1.0 < 2.0)")
self.assertEqual((st.vcc >> 1) & 1, 0, "Lane 1: expected vcc=0 (2.0 < 1.0)")
class TestVOP3VOPCModifiers(unittest.TestCase):
"""Tests for VOP3 VOPC with abs/neg modifiers."""
def test_v_cmp_ge_f32_abs_both(self):
"""v_cmp_ge_f32 with abs on both sources: abs(0.0) >= abs(-1.0) = false.
Regression test: int16 mod operation uses v_cmp_ge_f32 with abs modifiers.
"""
instructions = [
v_mov_b32_e32(v[0], 0.0),
v_mov_b32_e32(v[1], -1.0),
# abs=0b11 means abs(src0) and abs(src1)
v_cmp_ge_f32_e64(VCC_LO, v[0], v[1], abs=0b11),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "abs(0.0) >= abs(-1.0) should be false")
def test_v_cmp_ge_f32_abs_negative_divisor(self):
"""v_cmp_ge_f32 with abs: remainder check for negative divisor.
Tests the exact comparison used in int16 mod: abs(rem_f) >= abs(div_f).
For 1 % -1: rem_f = 0.0, div_f = -1.0, so abs(0.0) >= abs(-1.0) = false.
"""
instructions = [
v_mov_b32_e32(v[0], 0.0), # remainder as float
v_mov_b32_e32(v[1], -1.0), # divisor as float
v_cmp_ge_f32_e64(VCC_LO, v[0], v[1], abs=0b11),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "abs(0.0) >= abs(-1.0) should be false")
def test_v_cmp_ge_f32_abs_small_remainder(self):
"""v_cmp_ge_f32 with abs: abs(-0.5) >= abs(-3.0) = false."""
instructions = [
v_mov_b32_e32(v[0], -0.5),
v_mov_b32_e32(v[1], -3.0),
v_cmp_ge_f32_e64(VCC_LO, v[0], v[1], abs=0b11),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "abs(-0.5) >= abs(-3.0) should be false")
def test_v_cmp_ge_f32_abs_equal(self):
"""v_cmp_ge_f32 with abs: abs(-1.0) >= abs(1.0) = true."""
instructions = [
v_mov_b32_e32(v[0], -1.0),
v_mov_b32_e32(v[1], 1.0),
v_cmp_ge_f32_e64(VCC_LO, v[0], v[1], abs=0b11),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "abs(-1.0) >= abs(1.0) should be true")
class TestVOP3VOPC64Bit(unittest.TestCase):
"""Tests for VOP3 VOPC with 64-bit operands."""
def test_v_cmp_lt_f64_basic(self):
"""v_cmp_lt_f64: 0.0 < 1.0 = true."""
zero_f64 = f2i64(0.0)
one_f64 = f2i64(1.0)
instructions = [
s_mov_b32(s[0], zero_f64 & 0xffffffff),
s_mov_b32(s[1], zero_f64 >> 32),
s_mov_b32(s[2], one_f64 & 0xffffffff),
s_mov_b32(s[3], one_f64 >> 32),
v_cmp_lt_f64_e64(VCC_LO, s[0:1], s[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "0.0 < 1.0 should be true")
def test_v_cmp_lt_f64_negative(self):
"""v_cmp_lt_f64: -1.0 < 0.0 = true."""
neg_one_f64 = f2i64(-1.0)
zero_f64 = f2i64(0.0)
instructions = [
s_mov_b32(s[0], neg_one_f64 & 0xffffffff),
s_mov_b32(s[1], neg_one_f64 >> 32),
s_mov_b32(s[2], zero_f64 & 0xffffffff),
s_mov_b32(s[3], zero_f64 >> 32),
v_cmp_lt_f64_e64(VCC_LO, s[0:1], s[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "-1.0 < 0.0 should be true")
def test_v_cmp_lt_i64_signed(self):
"""v_cmp_lt_i64: 0 < -1 (signed) = false."""
instructions = [
s_mov_b32(s[0], 0),
s_mov_b32(s[1], 0), # s[0:1] = 0
s_mov_b32(s[2], 0xffffffff),
s_mov_b32(s[3], 0xffffffff), # s[2:3] = -1
v_cmp_lt_i64_e64(VCC_LO, s[0:1], s[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "0 < -1 (signed) should be false")
def test_v_cmp_lt_u64_unsigned(self):
"""v_cmp_lt_u64: 0 < 0xFFFFFFFFFFFFFFFF (unsigned) = true."""
instructions = [
s_mov_b32(s[0], 0),
s_mov_b32(s[1], 0), # s[0:1] = 0
s_mov_b32(s[2], 0xffffffff),
s_mov_b32(s[3], 0xffffffff), # s[2:3] = max uint64
v_cmp_lt_u64_e64(VCC_LO, s[0:1], s[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "0 < max_uint64 should be true")
class TestVOPCF64(unittest.TestCase):
"""Tests for VOPC (E32 encoding) with 64-bit float operands. Regression test for f64 compare bug."""
def test_v_cmp_lt_f64_e32_true(self):
"""v_cmp_lt_f64_e32: 2.0 < 3.0 = true."""
lo0, hi0 = f2i64(2.0) & 0xffffffff, f2i64(2.0) >> 32
lo1, hi1 = f2i64(3.0) & 0xffffffff, f2i64(3.0) >> 32
instructions = [
s_mov_b32(s[0], lo0), s_mov_b32(s[1], hi0),
s_mov_b32(s[2], lo1), s_mov_b32(s[3], hi1),
v_mov_b32_e32(v[0], s[0]), v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]), v_mov_b32_e32(v[3], s[3]),
v_cmp_lt_f64_e32(v[0:1], v[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "2.0 < 3.0 should be true")
def test_v_cmp_lt_f64_e32_false(self):
"""v_cmp_lt_f64_e32: 3.0 < 2.0 = false."""
lo0, hi0 = f2i64(3.0) & 0xffffffff, f2i64(3.0) >> 32
lo1, hi1 = f2i64(2.0) & 0xffffffff, f2i64(2.0) >> 32
instructions = [
s_mov_b32(s[0], lo0), s_mov_b32(s[1], hi0),
s_mov_b32(s[2], lo1), s_mov_b32(s[3], hi1),
v_mov_b32_e32(v[0], s[0]), v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]), v_mov_b32_e32(v[3], s[3]),
v_cmp_lt_f64_e32(v[0:1], v[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "3.0 < 2.0 should be false")
def test_v_cmp_nlt_f64_e32_true(self):
"""v_cmp_nlt_f64_e32: !(3.0 < 2.0) = true."""
lo0, hi0 = f2i64(3.0) & 0xffffffff, f2i64(3.0) >> 32
lo1, hi1 = f2i64(2.0) & 0xffffffff, f2i64(2.0) >> 32
instructions = [
s_mov_b32(s[0], lo0), s_mov_b32(s[1], hi0),
s_mov_b32(s[2], lo1), s_mov_b32(s[3], hi1),
v_mov_b32_e32(v[0], s[0]), v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]), v_mov_b32_e32(v[3], s[3]),
v_cmp_nlt_f64_e32(v[0:1], v[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(3.0 < 2.0) should be true")
def test_v_cmp_nlt_f64_e32_false(self):
"""v_cmp_nlt_f64_e32: !(2.0 < 3.0) = false."""
lo0, hi0 = f2i64(2.0) & 0xffffffff, f2i64(2.0) >> 32
lo1, hi1 = f2i64(3.0) & 0xffffffff, f2i64(3.0) >> 32
instructions = [
s_mov_b32(s[0], lo0), s_mov_b32(s[1], hi0),
s_mov_b32(s[2], lo1), s_mov_b32(s[3], hi1),
v_mov_b32_e32(v[0], s[0]), v_mov_b32_e32(v[1], s[1]),
v_mov_b32_e32(v[2], s[2]), v_mov_b32_e32(v[3], s[3]),
v_cmp_nlt_f64_e32(v[0:1], v[2:3]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(2.0 < 3.0) should be false")
class TestCmpxExec(unittest.TestCase):
"""Tests for V_CMPX instructions that modify EXEC mask."""
def test_v_cmpx_ngt_f32_e64_all_true(self):
"""V_CMPX_NGT_F32_E64: all lanes pass (literal <= all values)."""
# 131072.0 = 0x48000000
# All values > 131072, so !(131072 > val) = true for all
instructions = [
s_mov_b32(EXEC_LO, 0x7), # 3 lanes active
v_mov_b32_e32(v[0], f2i(200000.0)), # lane 0
v_cmp_eq_u32_e32(1, v[255]),
v_cndmask_b32_e64(v[1], v[0], f2i(300000.0), VCC_LO), # lane 1
v_cmp_eq_u32_e32(2, v[255]),
v_cndmask_b32_e64(v[1], v[1], f2i(400000.0), VCC_LO), # lane 2
# Now v[1] has: lane0=200000, lane1=300000, lane2=400000
# Compare: !(131072.0 > v[1]) i.e., 131072.0 <= v[1]
v_cmpx_ngt_f32_e64(EXEC_LO, f2i(131072.0), v[1]),
]
st = run_program(instructions, n_lanes=3)
# All values > 131072, so all lanes should remain active
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0x7, 0x7, "All 3 lanes should remain active")
def test_v_cmpx_ngt_f32_e64_some_false(self):
"""V_CMPX_NGT_F32_E64: some lanes fail (literal > some values)."""
instructions = [
s_mov_b32(EXEC_LO, 0x7), # 3 lanes active
v_mov_b32_e32(v[0], f2i(100000.0)), # lane 0: 131072 > 100000 = true, so !(true) = false
v_cmp_eq_u32_e32(1, v[255]),
v_cndmask_b32_e64(v[1], v[0], f2i(200000.0), VCC_LO), # lane 1: 131072 > 200000 = false, so !(false) = true
v_cmp_eq_u32_e32(2, v[255]),
v_cndmask_b32_e64(v[1], v[1], f2i(150000.0), VCC_LO), # lane 2: 131072 > 150000 = false, so !(false) = true
v_cmpx_ngt_f32_e64(EXEC_LO, f2i(131072.0), v[1]),
]
st = run_program(instructions, n_lanes=3)
# lane 0: fail (100000 < 131072), lanes 1,2: pass
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0x7, 0x6, "Lanes 1,2 should be active, lane 0 inactive")
def test_v_cmpx_ngt_f32_e64_all_false(self):
"""V_CMPX_NGT_F32_E64: all lanes fail (literal > all values)."""
instructions = [
s_mov_b32(EXEC_LO, 0x7), # 3 lanes active
v_mov_b32_e32(v[0], f2i(100.0)), # all lanes have 100.0
# 131072 > 100 = true, so !(true) = false for all
v_cmpx_ngt_f32_e64(EXEC_LO, f2i(131072.0), v[0]),
]
st = run_program(instructions, n_lanes=3)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0x7, 0x0, "All lanes should be inactive")
def test_v_cmpx_ngt_f32_e64_large_values(self):
"""V_CMPX_NGT_F32_E64: test with values that trigger Payne-Hanek in sin().
This is a regression test for the sin(859240.0) bug.
Values 859240, 1000000, 100594688 should all pass !(131072 > val).
"""
instructions = [
s_mov_b32(EXEC_LO, 0x7), # 3 lanes active
v_mov_b32_e32(v[0], f2i(859240.0)), # lane 0
v_cmp_eq_u32_e32(1, v[255]),
v_cndmask_b32_e64(v[1], v[0], f2i(1000000.0), VCC_LO), # lane 1
v_cmp_eq_u32_e32(2, v[255]),
v_cndmask_b32_e64(v[1], v[1], f2i(100594688.0), VCC_LO), # lane 2
v_cmpx_ngt_f32_e64(EXEC_LO, f2i(131072.0), v[1]),
]
st = run_program(instructions, n_lanes=3)
# All values > 131072, so !(131072 > val) = true for all
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0x7, 0x7, "All 3 lanes should remain active")
class TestVCCBehavior(unittest.TestCase):
"""Tests for VCC condition code behavior."""
def test_vcc_all_lanes_true(self):
"""VCC should have all bits set when all lanes compare true."""
instructions = [
v_mov_b32_e32(v[0], 5),
v_mov_b32_e32(v[1], 5),
v_cmp_eq_u32_e32(v[0], v[1]),
]
st = run_program(instructions, n_lanes=32)
self.assertEqual(st.vcc, 0xFFFFFFFF, "All 32 lanes should be true")
def test_vcc_lane_dependent(self):
"""VCC should differ per lane based on lane_id comparison."""
instructions = [
v_mov_b32_e32(v[0], 16),
v_cmp_lt_u32_e32(v[255], v[0]), # lanes 0-15 are < 16
]
st = run_program(instructions, n_lanes=32)
self.assertEqual(st.vcc & 0xFFFF, 0xFFFF, "Lanes 0-15 should be true")
self.assertEqual(st.vcc >> 16, 0x0000, "Lanes 16-31 should be false")
class TestCmpNge(unittest.TestCase):
"""Tests for V_CMP_NGE (not-greater-or-equal) with NaN semantics.
NGE = !(a >= b). With NaN inputs:
- If either input is NaN, a >= b is false, so !(false) = true
- This differs from a < b which returns false for NaN inputs
"""
def test_v_cmp_nge_f32_normal_values(self):
"""v_cmp_nge_f32: basic comparison with normal floats."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(2.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 2.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= 2.0) should be true")
def test_v_cmp_nge_f32_equal_values(self):
"""v_cmp_nge_f32: equal values should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(1.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_greater_value(self):
"""v_cmp_nge_f32: greater value should return false."""
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(2.0 >= 1.0) = !(true) = false
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 0, "!(2.0 >= 1.0) should be false")
def test_v_cmp_nge_f32_neg_inf(self):
"""v_cmp_nge_f32: -inf compared to normal value."""
neg_inf = 0xff800000 # -inf
instructions = [
s_mov_b32(s[0], neg_inf),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(-inf >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(-inf >= 1.0) should be true")
def test_v_cmp_nge_f32_clears_inactive_vcc_bits(self):
"""v_cmp_nge_f32 with partial EXEC clears inactive VCC bits (hardware behavior)."""
neg_inf = 0xff800000 # -inf
instructions = [
# Set VCC to all 1s first
s_mov_b32(VCC_LO, 0xFFFFFFFF),
# Set EXEC to only lane 0
s_mov_b32(EXEC_LO, 0x00000001),
# v0 = 1.0 for lane 0
v_mov_b32_e32(v[0], f2i(1.0)),
# Compare: !(-inf >= 1.0) = true for lane 0
v_cmp_nge_f32_e32(neg_inf, v[0]),
]
st = run_program(instructions, n_lanes=16)
# Hardware clears inactive lane bits, only active lane results remain
# Lane 0 result = 1 (true), lanes 1-15 = 0 (cleared)
self.assertEqual(st.vcc, 0x00000001, "VCC should only have active lane results")
def test_v_cmp_nge_f32_nan_src0(self):
"""v_cmp_nge_f32: NaN in src0 should return true (NaN >= x is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], f2i(1.0)),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= 1.0) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= 1.0) should be true")
def test_v_cmp_nge_f32_nan_src1(self):
"""v_cmp_nge_f32: NaN in src1 should return true (x >= NaN is false)."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], f2i(1.0)),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(1.0 >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(1.0 >= NaN) should be true")
def test_v_cmp_nge_f32_both_nan(self):
"""v_cmp_nge_f32: both NaN should return true."""
quiet_nan = 0x7fc00000
instructions = [
s_mov_b32(s[0], quiet_nan),
v_mov_b32_e32(v[0], s[0]),
v_mov_b32_e32(v[1], s[0]),
v_cmp_nge_f32_e32(v[0], v[1]), # !(NaN >= NaN) = !(false) = true
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vcc & 1, 1, "!(NaN >= NaN) should be true")
class TestCmpxPartialWavefront(unittest.TestCase):
"""Tests for V_CMPX with partial wavefronts (fewer than 32 active lanes).
Regression tests for bug where v_cmpx incorrectly set EXEC bits for inactive
lanes when the wavefront had fewer than 32 lanes. This caused garbage data
from uninitialized lanes to corrupt memory writes.
"""
def test_v_cmpx_eq_u32_partial_wave_3_lanes(self):
"""V_CMPX_EQ_U32 with 3 active lanes should only affect those 3 lanes.
With n_lanes=3, initial EXEC=0x7. After v_cmpx comparing lane_id == 1,
only lane 1 should pass, so EXEC should become 0x2 (not have bits 3-31 set).
"""
instructions = [
v_cmpx_eq_u32_e32(1, v[255]), # EXEC = lanes where lane_id == 1
]
st = run_program(instructions, n_lanes=3)
# Only lane 1 should be active (bit 1 set)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x2,
"Only lane 1 should be active after v_cmpx_eq_u32 with 3 lanes")
def test_v_cmpx_eq_u32_partial_wave_5_lanes(self):
"""V_CMPX_EQ_U32 with 5 active lanes."""
instructions = [
v_cmpx_eq_u32_e32(3, v[255]), # EXEC = lanes where lane_id == 3
]
st = run_program(instructions, n_lanes=5)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x8,
"Only lane 3 should be active after v_cmpx_eq_u32 with 5 lanes")
def test_v_cmpx_lt_u32_partial_wave(self):
"""V_CMPX_LT_U32 with partial wavefront."""
# VOPC: src0 < vsrc1, so we need v_cmpx_gt_u32 to get lane_id < 2
instructions = [
v_cmpx_gt_u32_e32(2, v[255]), # EXEC = lanes where 2 > lane_id (i.e., lane_id < 2)
]
st = run_program(instructions, n_lanes=4)
# Lanes 0,1 should be active (bits 0,1 set = 0x3)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x3,
"Only lanes 0,1 should be active after v_cmpx_gt_u32(2, lane_id) with 4 lanes")
def test_v_cmpx_ge_u32_partial_wave(self):
"""V_CMPX_GE_U32 with partial wavefront."""
# VOPC: src0 >= vsrc1, so v_cmpx_le_u32(1, lane_id) gives lane_id >= 2? No.
# v_cmpx_le_u32(src0, vsrc1) = src0 <= vsrc1 = 1 <= lane_id
instructions = [
v_cmpx_le_u32_e32(2, v[255]), # EXEC = lanes where 2 <= lane_id (i.e., lane_id >= 2)
]
st = run_program(instructions, n_lanes=4)
# Lanes 2,3 should be active (bits 2,3 set = 0xC)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0xC,
"Only lanes 2,3 should be active after v_cmpx_le_u32(2, lane_id) with 4 lanes")
def test_v_cmpx_ne_u32_partial_wave_all_pass(self):
"""V_CMPX_NE_U32 where all active lanes pass."""
instructions = [
v_cmpx_ne_u32_e32(99, v[255]), # EXEC = lanes where lane_id != 99
]
st = run_program(instructions, n_lanes=3)
# All 3 lanes should remain active (bits 0,1,2 set = 0x7)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x7,
"All 3 lanes should remain active when all pass")
def test_v_cmpx_eq_u32_partial_wave_none_pass(self):
"""V_CMPX_EQ_U32 where no active lanes pass."""
instructions = [
v_cmpx_eq_u32_e32(99, v[255]), # EXEC = lanes where lane_id == 99
]
st = run_program(instructions, n_lanes=3)
# No lanes should be active
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x0,
"No lanes should be active when none pass")
def test_v_cmpx_f32_partial_wave(self):
"""V_CMPX_GT_F32 with partial wavefront - float comparison."""
instructions = [
v_cvt_f32_u32_e32(v[0], v[255]), # v[0] = float(lane_id)
v_mov_b32_e32(v[1], f2i(0.5)), # v[1] = 0.5
v_cmpx_gt_f32_e32(v[0], v[1]), # EXEC = lanes where v[0] > 0.5
]
st = run_program(instructions, n_lanes=4)
# Lanes 1,2,3 have values > 0.5, lane 0 has 0.0
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0xE,
"Lanes 1,2,3 should be active (float > 0.5)")
def test_v_cmpx_e64_partial_wave(self):
"""V_CMPX_EQ_U32_E64 (VOP3 encoding) with partial wavefront."""
instructions = [
v_cmpx_eq_u32_e64(EXEC_LO, v[255], 2), # EXEC = lanes where lane_id == 2
]
st = run_program(instructions, n_lanes=4)
self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x4,
"Only lane 2 should be active after v_cmpx_eq_u32_e64")
if __name__ == '__main__':
unittest.main()

201
tinygrad_repo/test/amd/hw/test_vopd.py Normal file
View File

@@ -0,0 +1,201 @@
"""Tests for VOPD instructions - dual-issue vector operations.
VOPD executes two operations simultaneously. Key behavior:
- Both ops read their sources BEFORE either writes (dual-issue semantics)
- This means if X writes to a register that Y reads, Y sees the OLD value
- Op X can use ops 0-15 (FMAC, MUL, ADD, MOV, etc.)
- Op Y can use ops 0-18 (includes ADD_NC_U32, LSHLREV, AND)
"""
import unittest
from test.amd.hw.helpers import run_program, v, v_mov_b32_e32
from tinygrad.runtime.autogen.amd.rdna3.ins import VOPD, VOPD_LIT, VOPDOp
class TestVOPDBasic(unittest.TestCase):
"""Basic VOPD functionality tests."""
def test_vopd_dual_mov(self):
"""VOPD with two MOV operations to different registers."""
instructions = [
v_mov_b32_e32(v[0], 0x12345678),
v_mov_b32_e32(v[1], 0xDEADBEEF),
# X: v[2] = v[0], Y: v[3] = v[1]
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, v[2], v[3], v[0], v[1], v[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x12345678)
self.assertEqual(st.vgpr[0][3], 0xDEADBEEF)
def test_vopd_mov_and_add(self):
"""VOPD with MOV (X) and ADD_NC_U32 (Y) - ADD_NC_U32 can only be Y op."""
instructions = [
v_mov_b32_e32(v[0], 10),
v_mov_b32_e32(v[1], 5),
# X: v[2] = 100 (literal), Y: v[3] = v[0] + v[1] = 15
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_ADD_NC_U32, v[2], v[3], 100, v[0], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.vgpr[0][2], 100)
self.assertEqual(st.vgpr[0][3], 15)
class TestVOPDReadBeforeWrite(unittest.TestCase):
"""Tests for VOPD dual-issue read-before-write semantics.
In VOPD, both X and Y operations read their sources BEFORE either writes.
This is critical when X's destination is Y's source.
"""
def test_vopd_x_writes_y_reads_same_reg(self):
"""VOPD where X writes to a register that Y reads.
X: v[2] = 0 (overwrites v[2])
Y: v[1] = v[2] + v[0] (srcy0=v[2], vsrcy1=v[0])
If reads happen before writes: v[1] = OLD_v[2] + v[0] = 0xFFFFFFFF + 1 = 0
If writes happen before reads: v[1] = 0 + v[0] = 0 + 1 = 1
Hardware does reads-before-writes, so v[1] should be 0.
"""
instructions = [
v_mov_b32_e32(v[0], 1), # v[0] = 1
v_mov_b32_e32(v[1], 0x99999999), # v[1] = placeholder (will be overwritten)
v_mov_b32_e32(v[2], 0xFFFFFFFF), # v[2] = 0xFFFFFFFF
# X: v[2] = 0 (literal), srcx0=0, vsrcx1=v[0] (unused for MOV)
# Y: v[1] = srcy0 + vsrcy1 = v[2] + v[0] (should read OLD v[2] = 0xFFFFFFFF)
# vdsty encoding: (vdsty << 1) | ((vdstx & 1) ^ 1) where vdsty field = 0, vdstx = v[2]
# So vdsty_reg = (0 << 1) | ((2 & 1) ^ 1) = 0 | 1 = 1 = v[1]
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_ADD_NC_U32, v[2], v[0], 0, v[2], v[0], v[0]),
]
st = run_program(instructions, n_lanes=1)
# X should have written 0 to v[2]
self.assertEqual(st.vgpr[0][2], 0, "X should write 0 to v[2]")
# Y should have read OLD v[2] (0xFFFFFFFF) and added v[0] (1)
# 0xFFFFFFFF + 1 = 0 (wrap around)
self.assertEqual(st.vgpr[0][1], 0, "Y should read OLD v[2]=0xFFFFFFFF, compute 0xFFFFFFFF+1=0")
def test_vopd_x_writes_y_reads_same_reg_v2(self):
"""VOPD where X writes to a register that Y reads - cleaner test case.
X: v[2] = 0 (MOV)
Y: v[1] = v[2] + v[2] (ADD_NC_U32 with both sources from v[2])
If reads happen before writes: v[1] = OLD_v[2] + OLD_v[2] = 100 + 100 = 200
If writes happen before reads: v[1] = 0 + 0 = 0
Hardware does reads-before-writes, so v[1] should be 200.
"""
instructions = [
v_mov_b32_e32(v[0], 0x88888888), # v[0] = unused placeholder
v_mov_b32_e32(v[1], 0x99999999), # v[1] = placeholder (will be overwritten)
v_mov_b32_e32(v[2], 100), # v[2] = 100
# X: v[2] = 0 (literal)
# Y: v[1] = srcy0 + vsrcy1 = v[2] + v[2] (should read OLD v[2] = 100)
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_ADD_NC_U32, v[2], v[0], 0, v[2], v[0], v[2]),
]
st = run_program(instructions, n_lanes=1)
# X should have written 0 to v[2]
self.assertEqual(st.vgpr[0][2], 0, "X should write 0 to v[2]")
# Y should have read OLD v[2] (100) twice and added them
self.assertEqual(st.vgpr[0][1], 200, "Y should read OLD v[2]=100 twice, compute 100+100=200")
class TestVOPDLiterals(unittest.TestCase):
"""Tests for VOPD instructions that use SIMM32 literals (FMAAK, FMAMK)."""
def test_vopd_fmaak_f32(self):
"""VOPD V_DUAL_FMAAK_F32: D = S0 * S1 + SIMM32 (literal addend).
Tests that the 32-bit literal (SIMM32) is correctly passed to the instruction.
fma(2.0, 3.0, 10.0) = 2*3 + 10 = 16.0
"""
from test.amd.hw.helpers import f2i, i2f
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)), # v[0] = 2.0
v_mov_b32_e32(v[1], f2i(3.0)), # v[1] = 3.0
# VOPD args: opx, opy, vdstx, vdsty, srcx0, srcy0, vsrcx1, vsrcy1
# X: v[2] = fma(srcx0, vsrcx1, SIMM32) = v[0]*v[1]+10.0 = 2*3+10 = 16
# Y: v[3] = srcy0 (MOV) = v[0] = 2.0
VOPD_LIT(VOPDOp.V_DUAL_FMAAK_F32, VOPDOp.V_DUAL_MOV_B32, v[2], v[3], v[0], v[0], v[1], v[0], literal=f2i(10.0)),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 16.0, places=5, msg="fma(2.0, 3.0, 10.0) should be 16.0")
def test_vopd_fmamk_f32(self):
"""VOPD V_DUAL_FMAMK_F32: D = S0 * SIMM32 + S1 (literal multiplier).
Tests that the 32-bit literal (SIMM32) is correctly used as the multiplier.
fma(2.0, 5.0, 3.0) = 2*5 + 3 = 13.0
"""
from test.amd.hw.helpers import f2i, i2f
instructions = [
v_mov_b32_e32(v[0], f2i(2.0)), # v[0] = 2.0
v_mov_b32_e32(v[1], f2i(3.0)), # v[1] = 3.0
# X: v[2] = fma(srcx0, SIMM32, vsrcx1) = v[0]*5.0+v[1] = 2*5+3 = 13
# Y: v[3] = srcy0 (MOV) = v[0] = 2.0
VOPD_LIT(VOPDOp.V_DUAL_FMAMK_F32, VOPDOp.V_DUAL_MOV_B32, v[2], v[3], v[0], v[0], v[1], v[0], literal=f2i(5.0)),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][2]), 13.0, places=5, msg="fma(2.0, 5.0, 3.0) should be 13.0")
class TestVOPDDot2Acc(unittest.TestCase):
"""Tests for V_DUAL_DOT2ACC_F32_F16 - packed f16 dot product accumulate."""
def test_vopd_dot2acc_f32_f16_basic(self):
"""V_DUAL_DOT2ACC_F32_F16: D += lo(S0)*lo(S1) + hi(S0)*hi(S1).
S0 = pack(1.0h, 2.0h), S1 = pack(3.0h, 4.0h), D = 10.0f
result = 10.0 + 1.0*3.0 + 2.0*4.0 = 10.0 + 3.0 + 8.0 = 21.0
"""
from test.amd.hw.helpers import f2i, i2f, f32_to_f16
pk_s0 = f32_to_f16(1.0) | (f32_to_f16(2.0) << 16) # lo=1.0h, hi=2.0h
pk_s1 = f32_to_f16(3.0) | (f32_to_f16(4.0) << 16) # lo=3.0h, hi=4.0h
instructions = [
v_mov_b32_e32(v[0], pk_s0),
v_mov_b32_e32(v[1], pk_s1),
v_mov_b32_e32(v[3], f2i(10.0)), # accumulator in v[3] (vdsty with vdstx=v[4])
# X: v[4] = MOV v[0] (don't care), Y: v[3] += dot2(v[0], v[1])
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][3]), 21.0, places=2, msg="10.0 + 1.0*3.0 + 2.0*4.0 = 21.0")
def test_vopd_dot2acc_f32_f16_zero_accum(self):
"""V_DUAL_DOT2ACC_F32_F16 with zero accumulator — pure dot product.
S0 = pack(0.5h, -1.0h), S1 = pack(2.0h, 3.0h), D = 0.0f
result = 0.0 + 0.5*2.0 + (-1.0)*3.0 = 1.0 - 3.0 = -2.0
"""
from test.amd.hw.helpers import f2i, i2f, f32_to_f16
pk_s0 = f32_to_f16(0.5) | (f32_to_f16(-1.0) << 16)
pk_s1 = f32_to_f16(2.0) | (f32_to_f16(3.0) << 16)
instructions = [
v_mov_b32_e32(v[0], pk_s0),
v_mov_b32_e32(v[1], pk_s1),
v_mov_b32_e32(v[3], f2i(0.0)), # zero accumulator in v[3]
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
]
st = run_program(instructions, n_lanes=1)
self.assertAlmostEqual(i2f(st.vgpr[0][3]), -2.0, places=2, msg="0.5*2.0 + (-1.0)*3.0 = -2.0")
class TestVOPDMultilane(unittest.TestCase):
"""Tests for VOPD with multiple lanes."""
def test_vopd_multilane_mov_add(self):
"""VOPD MOV and ADD with multiple active lanes - no register conflict."""
instructions = [
v_mov_b32_e32(v[0], 5),
v_mov_b32_e32(v[1], 10),
# X: v[2] = 100 (constant), Y: v[1] = v[0] + v[1] = 5 + 10 = 15
# vdsty_reg = (vdsty << 1) | ((vdstx.offset & 1) ^ 1) = (0 << 1) | ((258 & 1) ^ 1) = 0 | 1 = 1
VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_ADD_NC_U32, v[2], v[0], 100, v[0], v[2], v[1]),
]
st = run_program(instructions, n_lanes=4)
for lane in range(4):
self.assertEqual(st.vgpr[lane][2], 100, f"Lane {lane}: v[2] should be 100")
self.assertEqual(st.vgpr[lane][1], 15, f"Lane {lane}: v[1] should be 15 (5+10)")
if __name__ == '__main__':
unittest.main()

212
tinygrad_repo/test/amd/test_custom_kernel.py Normal file
View File

@@ -0,0 +1,212 @@
import unittest
import functools
import numpy as np
from tinygrad import Tensor, Device, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
from tinygrad.renderer import Estimates
from tinygrad.dtype import AddrSpace
from tinygrad.helpers import getenv
from tinygrad.runtime.autogen.amd.rdna3.ins import *
import tinygrad.runtime.autogen.amd.rdna3.ins as r3
import tinygrad.runtime.autogen.amd.rdna4.ins as r4
from tinygrad.renderer.amd.dsl import s, v, NULL
from test.amd.helpers import TARGET_TO_ARCH
from extra.gemm.amd_asm_matmul import Kernel
def custom_add_one(A:UOp) -> UOp:
A = A.flatten()
assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
threads = UOp.special(A.numel(), "lidx0")
insts = [
s_load_b64(s[0:1], s[0:1], soffset=NULL),
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset
global_load_b32(v[1], v[0], saddr=s[0:1]),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_mov_b32_e32(v[2], 1.0),
v_add_f32_e32(v[1], v[1], v[2]),
global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
s_endpgm(),
]
sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
def custom_add_var(A:UOp, B:UOp) -> UOp:
A,B = A.flatten(), B.flatten()
assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
threads = UOp.special(A.numel(), "lidx0")
var = UOp.variable("var", 0, 10)
insts = [
s_load_b128(s[4:7], s[0:1]),
s_load_b32(s[8], s[0:1], offset=0x10), # all threads load the same variable
s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
v_lshlrev_b32_e32(v[0], 2, v[0]), # element offset, different per thread
global_load_b32(v[1], v[0], saddr=s[6:7]),
s_waitcnt_vmcnt(sdst=NULL, simm16=0),
v_add_nc_u32_e32(v[1], s[8], v[1]),
global_store_b32(addr=v[0], data=v[1], saddr=s[4:5]),
s_endpgm(),
]
sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.numel()}"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
def custom_wave_sync(A:UOp, arch:str) -> UOp:
# 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
# s_sleep yields the SIMD so waves from different WGs interleave, causing barrier packet reordering
threads = UOp.special(128, "lidx0")
wg = UOp.special(1024, "gidx0")
insts = []
for _ in range(4):
insts.append(s_sleep(4))
insts += [s_barrier()] if arch == "rdna3" else [r4.s_barrier_signal(), r4.s_barrier_wait()]
insts += [s_nop(0)]*4
insts.append(s_endpgm())
sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
def custom_lds_sync(A:UOp, arch:str) -> UOp:
A = A.flatten()
num_threads = A.shape[0]
threads = UOp.special(num_threads, "lidx0")
wg = UOp.special(1, "gidx0")
lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds') # 128 * 4 bytes
isa = r4 if arch == "rdna4" else r3
wait_kmcnt = [isa.s_wait_kmcnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt_lgkmcnt(sdst=NULL, simm16=0)]
wait_dscnt = [isa.s_wait_dscnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt_lgkmcnt(sdst=NULL, simm16=0)]
barrier = [isa.s_barrier_signal(ssrc0=-1), isa.s_barrier_wait(simm16=-1)] if arch == "rdna4" else [isa.s_barrier()]
global_store = [isa.global_store_b32(vaddr=v[6:7], saddr=s[0:1], vsrc=v[5])] if arch == "rdna4" \
else [isa.global_store_b32(addr=v[6], data=v[5], saddr=s[0:1])]
insts = [
isa.s_load_b64(s[0:1], s[0:1], soffset=NULL),
*wait_kmcnt,
isa.v_lshlrev_b32_e32(v[1], 2, v[0]),
# lds[thread_idx] = thread_idx
isa.ds_store_b32(addr=v[1], data0=v[0]),
*wait_dscnt,
*barrier,
# out[threaed_idx] = thread_idx == num_threads ? -1 : lds[thread_idx + 1]
isa.v_add_nc_u32_e32(v[2], 4, v[1]),
isa.v_cmp_gt_u32_e32(num_threads-1, v[0]),
isa.ds_load_b32(vdst=v[3], addr=v[2]),
*wait_dscnt,
isa.v_mov_b32_e32(v[4], -1),
isa.v_cndmask_b32_e32(v[5], v[4], v[3]),
isa.v_lshlrev_b32_e32(v[6], 2, v[0]),
*global_store,
isa.s_endpgm(),
]
sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
def custom_handwritten(A:UOp) -> UOp:
A = A.flatten()
threads = UOp.special(128, "lidx0")
wg = UOp.special(1, "gidx0")
lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds') # 128 * 4 bytes
pipes = {getenv("PIPE", "")} if getenv("PIPE", "") else {"SALU", "VALU", "TRANSCENDENTAL", "WMMA"}
k = Kernel()
# wrap in loop to filter out icache misses
LOOP_N, UNROLL_N = 8, 5
k.emit(r4.s_mov_b32(s[1], LOOP_N))
k.label("loop")
if "SALU" in pipes:
for i in range(UNROLL_N):
k.emit(r4.s_mov_b32(s[20+i], i))
k.emit(r4.s_min_i32(s[30+i], i))
k.emit(r4.s_mov_b32(s[40+i], i))
k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32))
if "VALU" in pipes:
for i in range(UNROLL_N):
k.emit(r4.v_mov_b32_e32(v[20+i], i))
k.emit(r4.v_lshlrev_b64_e32(v[30+2*i:31+2*i], 2, v[12+i:13+i]))
k.emit(r4.v_mad_co_u64_u32(v[40+2*i:41+2*i], NULL, v[12+i], v[13+i], v[14+i:15+i]))
if "TRANSCENDENTAL" in pipes:
# transcendental VALU runs on the TFU, it can run regular VALU at the same time
for i in range(UNROLL_N):
k.emit(r4.v_mov_b32_e32(v[20+i], i))
k.emit(r4.v_s_rcp_f32(s[20+i], s[12+i]))
k.emit(r4.v_rcp_f32_e32(v[30+i], v[12+i]))
k.emit(r4.v_s_exp_f32(s[30+i], s[12+i]))
if "WMMA" in pipes:
base = 30
for i in range(UNROLL_N):
a = base + i*40
b, cd = a + 4, a + 8
k.emit(r4.v_wmma_f32_16x16x16_f16(v[cd:cd+7], v[a:a+3], v[b:b+3], v[cd:cd+7]))
a = base + i*40 + 16
b, cd = a + 2, a + 4
k.emit(r4.v_wmma_i32_16x16x16_iu8(v[cd:cd+7], v[a:a+1], v[b:b+1], v[cd:cd+7]))
k.emit(r4.s_add_co_i32(s[1], s[1], -1))
k.emit(r4.s_cmp_eq_i32(s[1], 0))
k.emit(r4.s_cbranch_scc0(), target="loop")
k.emit(r4.s_endpgm())
insts = k.finalize()
sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
def custom_data_deps(A:UOp) -> UOp:
A = A.flatten()
threads = UOp.special(A.numel(), "lidx0")
k = Kernel()
k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
k.emit(s_waitcnt_lgkmcnt(sdst=NULL, simm16=0))
k.emit(v_lshlrev_b32_e32(v[0], 2, v[0]))
k.emit(global_load_b32(v[1], v[0], saddr=s[0:1]))
k.emit(s_waitcnt_vmcnt(sdst=NULL, simm16=0))
k.emit(v_add_f32_e32(v[1], 1.0, v[1]))
k.emit(global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]))
k.emit(s_endpgm())
insts = k.finalize()
sink = UOp.sink(A.base, threads, arg=KernelInfo("custom_data_deps"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
class TestCustomKernel(unittest.TestCase):
def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
def test_simple(self):
if self.arch != "rdna3": self.skipTest("only rdna3")
a = Tensor.full((16, 16), 1.).contiguous().realize()
a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
linear = compile_linear(a.schedule_linear())
est = estimate_uop(linear.src[-1])
self.assertEqual(est.ops, a.numel())
self.assertEqual(est.mem, a.nbytes()*2)
run_linear(linear)
self.assertTrue((a.numpy() == 2.).all())
def test_variable(self):
if self.arch != "rdna3": self.skipTest("only rdna3")
b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
a = Tensor.zeros_like(b).contiguous().realize()
a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
linear = a.schedule_linear()
for i in range(4):
run_linear(linear, var_vals={"var":i})
self.assertTrue((a.numpy() == 1+i).all())
def test_lds_sync(self):
if self.arch not in ("rdna3", "rdna4"): self.skipTest("only rdna3/rdna4")
a = Tensor.empty(128, dtype=dtypes.int32).contiguous().realize()
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_lds_sync, arch=self.arch))[0]
a.realize()
ref = Tensor.arange(1, 129, dtype=dtypes.int32)
ref[127] = -1
self.assertListEqual(a.tolist(), ref.tolist())
def test_handwritten(self):
if self.arch != "rdna4": self.skipTest("only tested on rdna4")
a = Tensor.empty(1024, dtype=dtypes.int32).contiguous().realize()
a = Tensor.custom_kernel(a, fxn=custom_handwritten)[0]
a.realize()
def test_data_deps(self):
if self.arch != "rdna3": self.skipTest("only tested on rdna3")
a = Tensor(np.full(32, 5.0, dtype=np.float32)).realize()
a = Tensor.custom_kernel(a, fxn=custom_data_deps)[0]
a.realize()
self.assertTrue((a.numpy() == 6.0).all())
if __name__ == "__main__":
unittest.main()

170
tinygrad_repo/test/amd/test_dsl2.py Normal file
View File

@@ -0,0 +1,170 @@
import unittest
from tinygrad.renderer.amd.dsl import *
from tinygrad.renderer.amd.dsl import VDSTYField
from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op
from tinygrad.runtime.autogen.amd.rdna3.ins import VOP1
class TestRegisters(unittest.TestCase):
def test_vgpr_single(self):
self.assertEqual(repr(v[5]), "v[5]")
self.assertEqual(v[5].offset, 261) # 256 + 5
self.assertEqual(v[5].sz, 1)
def test_sgpr_single(self):
self.assertEqual(repr(s[10]), "s[10]")
self.assertEqual(s[10].offset, 10)
def test_vgpr_range(self):
self.assertEqual(repr(v[0:3]), "v[0:3]")
self.assertEqual(v[0:3].offset, 256)
self.assertEqual(v[0:3].sz, 4)
def test_sgpr_range(self):
self.assertEqual(repr(s[4:5]), "s[4:5]")
self.assertEqual(s[4:5].sz, 2)
def test_ttmp_reslice(self):
# ttmp is src[108:123], so ttmp[0] should be src[108]
self.assertEqual(ttmp[0].offset, 108)
self.assertEqual(ttmp[1].offset, 109)
# ttmp[0:1] is 2 elements (inclusive slicing)
self.assertEqual(ttmp[0:1].offset, 108)
self.assertEqual(ttmp[0:1].sz, 2)
# ttmp[0:1][0] should be src[108]
self.assertEqual(ttmp[0:1][0].offset, 108)
def test_special_regs(self):
self.assertEqual(NULL.offset, 124)
self.assertEqual(M0.offset, 125)
self.assertEqual(EXEC_LO.offset, 126)
self.assertEqual(EXEC_HI.offset, 127)
# Check repr round-trips
self.assertEqual(repr(NULL), "NULL")
self.assertEqual(repr(M0), "M0")
self.assertEqual(repr(EXEC_LO), "EXEC_LO")
self.assertEqual(repr(EXEC), "EXEC")
def test_vcc(self):
self.assertEqual(VCC.offset, 106)
self.assertEqual(VCC.sz, 2)
self.assertEqual(VCC_LO.offset, 106)
self.assertEqual(VCC_HI.offset, 107)
# Check repr round-trips
self.assertEqual(repr(VCC_LO), "VCC_LO")
self.assertEqual(repr(VCC_HI), "VCC_HI")
self.assertEqual(repr(VCC), "VCC")
def test_float_constants(self):
self.assertEqual(src[240].offset, 240)
self.assertEqual(repr(src[240]), "0.5")
self.assertEqual(repr(src[242]), "1.0")
self.assertEqual(repr(src[243]), "-1.0")
def test_int_constants(self):
self.assertEqual(repr(src[128]), "0")
self.assertEqual(repr(src[129]), "1")
self.assertEqual(repr(src[192]), "64")
self.assertEqual(repr(src[193]), "-1")
self.assertEqual(repr(src[208]), "-16")
class TestEnumBitField(unittest.TestCase):
def test_enum_name(self):
self.assertEqual(VOP1Op.V_MOV_B32_E32.name, "V_MOV_B32_E32")
def test_enum_value(self):
self.assertEqual(VOP1Op.V_MOV_B32_E32.value, 1)
def test_enum_comparison(self):
self.assertEqual(VOP1Op.V_MOV_B32_E32, VOP1Op.V_MOV_B32_E32)
self.assertNotEqual(VOP1Op.V_NOP_E32, VOP1Op.V_MOV_B32_E32)
def test_enum_different_types(self):
# VOP1Op and VOP2Op are different enums, even if same value
self.assertNotEqual(VOP1Op.V_MOV_B32_E32, VOP2Op.V_CNDMASK_B32_E32)
def test_wrong_enum_type_raises(self):
# Passing VOP2Op to VOP1 should raise
with self.assertRaises(RuntimeError):
VOP1(VOP2Op.V_CNDMASK_B32_E32, v[5], v[6])
class TestVOP1(unittest.TestCase):
def test_class_setup(self):
self.assertEqual(VOP1._size(), 4)
field_names = [n for n, _ in VOP1._fields]
self.assertIn('encoding', field_names)
self.assertIn('op', field_names)
self.assertIn('vdst', field_names)
self.assertIn('src0', field_names)
def test_encoding_vgpr_vgpr(self):
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
raw = i._raw
# Check each field
self.assertEqual((raw >> 25) & 0x7f, 0b0111111) # encoding
self.assertEqual((raw >> 17) & 0xff, 5) # vdst (just VGPR index)
self.assertEqual((raw >> 9) & 0xff, 1) # op
self.assertEqual(raw & 0x1ff, 262) # src0 (256 + 6)
def test_encoding_vgpr_sgpr(self):
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], s[10])
raw = i._raw
self.assertEqual((raw >> 17) & 0xff, 5) # vdst (just VGPR index)
self.assertEqual(raw & 0x1ff, 10) # src0 (SGPR encoded)
def test_to_bytes(self):
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
b = i.to_bytes()
self.assertEqual(len(b), 4)
self.assertEqual(int.from_bytes(b, 'little'), i._raw)
def test_from_bytes(self):
i1 = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
i2 = VOP1.from_bytes(i1.to_bytes())
self.assertEqual(i1._raw, i2._raw)
def test_repr(self):
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
self.assertEqual(repr(i), "v_mov_b32_e32(v[5], v[6])")
def test_repr_sgpr_src(self):
i = VOP1(VOP1Op.V_MOV_B32_E32, v[5], s[10])
self.assertEqual(repr(i), "v_mov_b32_e32(v[5], s[10])")
def test_kwargs(self):
i1 = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
i2 = VOP1(op=VOP1Op.V_MOV_B32_E32, vdst=v[5], src0=v[6])
self.assertEqual(i1._raw, i2._raw)
def test_kwargs_partial(self):
i1 = VOP1(VOP1Op.V_MOV_B32_E32, v[5], v[6])
i2 = VOP1(VOP1Op.V_MOV_B32_E32, src0=v[6], vdst=v[5])
self.assertEqual(i1._raw, i2._raw)
class TestVDSTYField(unittest.TestCase):
def test_encode_even_vgpr(self):
f = VDSTYField(6, 0) # 7-bit field
self.assertEqual(f.encode(v[0]), 0)
self.assertEqual(f.encode(v[2]), 1)
self.assertEqual(f.encode(v[4]), 2)
self.assertEqual(f.encode(v[254]), 127)
def test_encode_non_vgpr_raises(self):
f = VDSTYField(6, 0)
with self.assertRaises(ValueError) as ctx:
f.encode(s[0])
self.assertIn("VGPR", str(ctx.exception))
def test_encode_non_reg_raises(self):
f = VDSTYField(6, 0)
with self.assertRaises(TypeError) as ctx:
f.encode(42)
self.assertIn("Reg", str(ctx.exception))
def test_decode_returns_raw(self):
f = VDSTYField(6, 0)
# decode returns raw value, actual vdsty computed with vdstx context
self.assertEqual(f.decode(0), 0)
self.assertEqual(f.decode(127), 127)
if __name__ == "__main__":
unittest.main()

396
tinygrad_repo/test/amd/test_emu2_pcode.py Normal file
View File

@@ -0,0 +1,396 @@
"""Tests for the pcode parser."""
import unittest
from collections import defaultdict
from tinygrad.helpers import DEBUG
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, Ops
from test.mockgpu.amd.emu import parse_pcode
from test.mockgpu.amd.pcode import parse_expr
from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE
from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp, GLOBALOp
def _srcs():
"""Create minimal source variables for pcode parsing."""
def u32(v=0): return UOp.const(dtypes.uint32, v)
return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}
class TestBasicParsing(unittest.TestCase):
"""Test basic pcode parsing for common instruction patterns."""
def test_v_add_f32(self):
"""Test parsing V_ADD_F32 pcode."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], _srcs())
self.assertEqual(len(assigns), 1)
dest, _ = assigns[0]
self.assertTrue(dest.startswith('D0'))
def test_v_lshlrev_b32(self):
"""Test parsing V_LSHLREV_B32 pcode."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_LSHLREV_B32_E32], _srcs())
self.assertEqual(len(assigns), 1)
def test_s_cselect_b32(self):
"""Test parsing S_CSELECT_B32 pcode with ternary."""
_, assigns = parse_pcode(PCODE[SOP2Op.S_CSELECT_B32], _srcs())
self.assertEqual(len(assigns), 1)
def test_v_add_co_ci_u32(self):
"""Test parsing V_ADD_CO_CI_U32 with carry."""
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_CO_CI_U32_E32], _srcs())
self.assertGreaterEqual(len(assigns), 1)
class TestWithSources(unittest.TestCase):
"""Test pcode parsing with actual source operand values."""
def test_v_add_f32_with_sources(self):
"""Test V_ADD_F32 with actual float constants."""
s0 = UOp.const(dtypes.uint32, 0x3f800000) # 1.0f
s1 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f
_, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be an ADD operation
self.assertEqual(val.op, Ops.ADD)
def test_v_mul_f32_with_sources(self):
"""Test V_MUL_F32 with actual float constants."""
s0 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f
s1 = UOp.const(dtypes.uint32, 0x40400000) # 3.0f
_, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertEqual(val.op, Ops.MUL)
class TestParseExpr(unittest.TestCase):
"""Test the parse_expr function directly."""
def test_integer_literals(self):
"""Test parsing integer literals."""
self.assertEqual(parse_expr('0', {}).arg, 0)
self.assertEqual(parse_expr('42', {}).arg, 42)
self.assertEqual(parse_expr('42U', {}).arg, 42)
def test_negative_integers(self):
"""Test parsing negative integer literals."""
result = parse_expr('-1', {})
self.assertEqual(result.arg, -1)
self.assertEqual(result.dtype, dtypes.int)
def test_float_literals(self):
"""Test parsing float literals."""
result = parse_expr('1.0F', {})
self.assertEqual(result.arg, 1.0)
self.assertEqual(result.dtype, dtypes.float32)
def test_hex_literals(self):
"""Test parsing hex literals."""
result = parse_expr('0xFF', {})
self.assertEqual(result.arg, 255)
def test_variable_lookup(self):
"""Test variable lookup in parse_expr."""
vrs = {'x': UOp.const(dtypes.uint32, 42)}
result = parse_expr('x', vrs)
self.assertEqual(result.arg, 42)
def test_binary_ops(self):
"""Test parsing binary operations."""
vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}
# Addition
result = parse_expr('a + b', vrs)
self.assertEqual(result.op, Ops.ADD)
# Subtraction with constant folding
result = parse_expr('10 - 5', {})
self.assertEqual(result.op, Ops.CONST)
self.assertEqual(result.arg, 5)
def test_ternary(self):
"""Test parsing ternary expressions."""
vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
result = parse_expr('cond ? a : b', vrs)
self.assertEqual(result.op, Ops.WHERE)
class TestForLoopParsing(unittest.TestCase):
"""Test for loop parsing (CLZ/CTZ patterns)."""
def test_clz_pcode_exists(self):
"""Verify CLZ pcode is available."""
pcode = PCODE.get(VOP1Op.V_CLZ_I32_U32_E32)
self.assertIsNotNone(pcode)
assert pcode is not None
self.assertIn('for', pcode.lower())
def test_clz_parsing(self):
"""Test CLZ pcode parsing produces correct structure."""
pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF) # All ones - CLZ should be 0
_vrs, assigns = parse_pcode(pcode, {'S0': S0})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be a nested WHERE structure
self.assertEqual(val.op, Ops.WHERE)
def test_clz_with_zero(self):
"""Test CLZ with input 0 - should return -1."""
pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
S0 = UOp.const(dtypes.uint32, 0)
_vrs, assigns = parse_pcode(pcode, {'S0': S0})
# Check that the innermost value (default) is -1 (may be wrapped in CAST)
val = assigns[0][1]
# Traverse to innermost WHERE
while val.op == Ops.WHERE:
val = val.src[2] # false branch
# Unwrap CAST if present
while val.op == Ops.CAST:
val = val.src[0]
self.assertEqual(val.arg, -1)
def test_ctz_parsing(self):
"""Test CTZ pcode parsing."""
pcode = PCODE.get(VOP1Op.V_CTZ_I32_B32_E32)
if pcode is None:
self.skipTest("V_CTZ_I32_B32_E32 pcode not available")
S0 = UOp.const(dtypes.uint32, 1) # LSB set - CTZ should be 0
_vrs, assigns = parse_pcode(pcode, {'S0': S0})
self.assertEqual(len(assigns), 1)
class TestDSPcodePatterns(unittest.TestCase):
"""Test DS instruction pcode patterns."""
def test_global_atomic_add_f32_parsing(self):
"""Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
vmem = UOp.param(2, dtypes.uint32.ptr(1024))
srcs = {
'ADDR': UOp.const(dtypes.uint64, 0),
'DATA': UOp.const(dtypes.uint32, 0x3f800000),
'_vmem': vmem,
}
_, assigns = parse_pcode(PCODE[GLOBALOp.GLOBAL_ATOMIC_ADD_F32], srcs)
mem_write = next(val for dest, val in assigns if dest == 'MEM[ADDR].f32')
self.assertEqual(mem_write[1].op, Ops.ADD) # type: ignore[index]
self.assertEqual(mem_write[1].dtype, dtypes.float32) # type: ignore[index]
def test_ds_load_b32_pcode(self):
"""Test DS_LOAD_B32 pcode is parseable."""
pcode = PCODE.get(DSOp.DS_LOAD_B32)
self.assertIsNotNone(pcode)
assert pcode is not None
self.assertIn('RETURN_DATA', pcode)
self.assertIn('MEM[', pcode)
def test_ds_store_b32_pcode(self):
"""Test DS_STORE_B32 pcode is parseable."""
pcode = PCODE.get(DSOp.DS_STORE_B32)
self.assertIsNotNone(pcode)
assert pcode is not None
self.assertIn('MEM[', pcode)
self.assertIn('DATA', pcode)
def test_mem_read_parsing(self):
"""Test MEM[addr].type read expression parsing."""
# Create a mock LDS buffer
lds = UOp.param(3, dtypes.uint32.ptr(16384))
addr = UOp.const(dtypes.uint32, 0)
vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}
result = parse_expr('MEM[ADDR + OFFSET].b32', vrs)
# Should be an INDEX operation into LDS
self.assertIsNotNone(result)
def test_ds_store_2addr_b32_parsing(self):
"""Test DS_STORE_2ADDR_B32 pcode parsing produces MEM writes."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
self.assertIsNotNone(pcode)
assert pcode is not None
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Should have 2 MEM write assignments
self.assertEqual(len(assigns), 2)
for dest, val in assigns:
self.assertTrue(dest.startswith('MEM['))
# val should be (addr, write_val) tuple
self.assertIsInstance(val, tuple)
self.assertEqual(len(val), 2) # type: ignore[arg-type]
def test_ds_load_2addr_b32_parsing(self):
"""Test DS_LOAD_2ADDR_B32 pcode parsing produces RETURN_DATA assignments."""
pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
self.assertIsNotNone(pcode)
assert pcode is not None
lds = UOp.param(3, dtypes.uint32.ptr(16384))
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'_lds': lds,
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Should have 2 RETURN_DATA assignments
self.assertEqual(len(assigns), 2)
self.assertEqual(assigns[0][0], 'RETURN_DATA[31:0]')
self.assertEqual(assigns[1][0], 'RETURN_DATA[63:32]')
def test_ds_store_address_calculation(self):
"""Test DS_STORE_2ADDR_B32 calculates correct addresses (offset * 4)."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
assert pcode is not None
srcs = {
'ADDR': UOp.const(dtypes.uint32, 100),
'OFFSET0': UOp.const(dtypes.uint32, 2),
'OFFSET1': UOp.const(dtypes.uint32, 5),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
# assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
self.assertEqual(assigns[0][1][0].simplify().arg, 108) # type: ignore[index]
self.assertEqual(assigns[1][1][0].simplify().arg, 120) # type: ignore[index]
def test_ds_store_data_values(self):
"""Test DS_STORE_2ADDR_B32 uses correct data values."""
pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
assert pcode is not None
srcs = {
'ADDR': UOp.const(dtypes.uint32, 0),
'OFFSET0': UOp.const(dtypes.uint32, 0),
'OFFSET1': UOp.const(dtypes.uint32, 1),
'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
}
srcs['laneId'] = UOp.const(dtypes.uint32, 0)
_, assigns = parse_pcode(pcode, srcs)
# assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
# DATA[31:0] should preserve the value
self.assertEqual(assigns[0][1][1].simplify().arg, 0xAAAAAAAA) # type: ignore[index]
self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB) # type: ignore[index]
class TestConditionalParsing(unittest.TestCase):
"""Test conditional (if/elsif/else) pcode parsing."""
def test_ternary_in_assignment(self):
"""Test parsing ternary expression (which becomes WHERE)."""
# S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32
pcode = PCODE[SOP2Op.S_CSELECT_B32]
s0 = UOp.const(dtypes.uint32, 10)
s1 = UOp.const(dtypes.uint32, 20)
scc = UOp.const(dtypes.uint32, 1)
_vrs, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
self.assertEqual(len(assigns), 1)
dest, val = assigns[0]
self.assertTrue(dest.startswith('D0'))
# Result should be a WHERE (ternary becomes WHERE)
self.assertEqual(val.op, Ops.WHERE)
class TestConcatWidthParsing(unittest.TestCase):
"""Test that bit extracts keep the right width for concat/unary ops."""
def test_permlanex16_altrow_concat(self):
for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(dtypes.uint32, row)})
self.assertEqual(parsed.simplify().arg, expected)
def test_permlane64_altlane_concat(self):
for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(dtypes.uint32, lane)})
self.assertEqual(parsed.simplify().arg, expected)
def test_permlane64_wave64_pcode_indices(self):
vgpr = UOp.param(0, dtypes.uint32.ptr(256))
srcs = {
'SRC0': UOp.const(dtypes.uint32, 0),
'VDST': UOp.const(dtypes.uint32, 1),
'EXEC_LO': UOp.const(dtypes.uint32, 0xFFFFFFFF),
'EXEC': UOp.const(dtypes.uint64, 0xFFFFFFFFFFFFFFFF),
'_vgpr': vgpr,
'_wave_size': 64,
'S0': UOp.const(dtypes.uint32, 0),
'S1': UOp.const(dtypes.uint32, 0),
'S2': UOp.const(dtypes.uint32, 0),
}
def load_idx(v: UOp) -> int:
simp = v.simplify()
self.assertEqual(simp.op, Ops.LOAD)
self.assertEqual(simp.src[0].op, Ops.INDEX)
idx = simp.src[0].src[1].simplify()
self.assertEqual(idx.op, Ops.CONST)
return idx.arg
_, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
self.assertEqual(len(assigns), 64)
for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
self.assertEqual(assigns[lane][1][0].simplify().arg, dst_idx) # type: ignore[index]
self.assertEqual(load_idx(assigns[lane][1][1]), src_idx) # type: ignore[index]
class TestAllPcode(unittest.TestCase):
"""Test that all pcode from all architectures can be parsed."""
def _make_srcs(self):
"""Create dummy source variables for pcode parsing."""
u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
lds = UOp.param(3, dtypes.uint32.ptr(16384))
return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
'SIMM16': u32(), 'SIMM32': u32(), 'OFFSET': u32(), 'OFFSET0': u32(), 'OFFSET1': u32(), 'offset1': u32(),
'ADDR': u32(), 'ADDR_BASE': u32(), 'TADDR': u32(), 'DATA': u32(), 'DATA0': u32(), 'DATA1': u32(), 'DATA2': u32(),
'VDATA': u32(), 'VDATA0': u32(), 'VDATA1': u32(), 'VDATA2': u32(), 'VDATA3': u32(),
'OPSEL': u32(), 'OPSEL_HI': u32(), 'NEG': u32(), 'NEG_HI': u32(), 'CLAMP': u32(),
'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(),
'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
'_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}
def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float):
"""Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs."""
srcs = self._make_srcs()
passed, skipped, errors = 0, 0, defaultdict(list)
for op, pcode in pcode_dict.items():
try:
parse_pcode(pcode, srcs)
passed += 1
except RuntimeError as e:
skipped += 1
errors[str(e)].append(op.name)
except Exception as e: self.fail(f"[{arch}] {op.name}: {e}\nPcode: {pcode[:200]}")
total = len(pcode_dict)
pct = 100 * passed / total
print(f"{arch}: {passed}/{total} ({pct:.1f}%) parsed, {skipped} skipped")
if DEBUG >= 2:
for err, ops in sorted(errors.items(), key=lambda x: -len(x[1])):
print(f" {err}: {', '.join(ops[:5])}{'...' if len(ops) > 5 else ''} ({len(ops)})")
self.assertGreaterEqual(pct, min_pct, f"[{arch}] {pct:.1f}% < {min_pct}% threshold")
def test_parse_all_cdna_pcode(self):
from tinygrad.runtime.autogen.amd.cdna.str_pcode import PCODE as CDNA_PCODE
self._parse_all_pcode(CDNA_PCODE, "CDNA", min_pct=60)
def test_parse_all_rdna3_pcode(self):
from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE as RDNA3_PCODE
self._parse_all_pcode(RDNA3_PCODE, "RDNA3", min_pct=90)
def test_parse_all_rdna4_pcode(self):
from tinygrad.runtime.autogen.amd.rdna4.str_pcode import PCODE as RDNA4_PCODE
self._parse_all_pcode(RDNA4_PCODE, "RDNA4", min_pct=65)
if __name__ == "__main__":
unittest.main()

229
tinygrad_repo/test/amd/test_formats.py Normal file
View File

@@ -0,0 +1,229 @@
#!/usr/bin/env python3
"""Test DS and other compute-relevant instruction formats.
Note: Graphics-only formats (EXP, MUBUF, MTBUF, MIMG) are not supported - use GLOBAL/FLAT for memory access in compute.
"""
import unittest
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from tinygrad.renderer.amd.dsl import VCC_HI, EXEC_LO, NULL
OFF = NULL # OFF is alias for NULL
from tinygrad.renderer.amd import detect_format
class TestDS(unittest.TestCase):
"""Test DS (data share / LDS) instructions."""
def test_ds_store_b32(self):
# ds_store_b32 v0, v1
# GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_load_b32(self):
# ds_load_b32 v0, v1
# GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
inst = ds_load_b32(vdst=v[0], addr=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
def test_ds_store_b32_offset(self):
# ds_store_b32 v0, v1 offset:64
# GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_load_b64(self):
# ds_load_b64 v[0:1], v2
# GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
inst = ds_load_b64(vdst=v[0:1], addr=v[2])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
def test_ds_add_u32(self):
# ds_add_u32 v0, v1
# GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
inst = ds_add_u32(addr=v[0], data0=v[1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
def test_ds_store_b32_gds(self):
# ds_store_b32 v0, v1 gds
# GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
class TestVOP3(unittest.TestCase):
"""Test VOP3 (3-operand vector) instructions."""
def test_v_fma_f32(self):
# v_fma_f32 v0, v1, v2, v3
# GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
def test_v_mad_f32(self):
# v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
# Use v_fma_f32 with vdst == src2
inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
def test_v_add3_u32(self):
# v_add3_u32 v0, v1, v2, v3
# GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
class TestFLAT(unittest.TestCase):
"""Test FLAT/GLOBAL/SCRATCH memory instructions."""
def test_global_load_b32(self):
# global_load_b32 v0, v[1:2], off (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
inst = global_load_b32(vdst=v[0], addr=v[1:2], saddr=OFF)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
def test_global_store_b32(self):
# global_store_b32 v[0:1], v2, off (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
inst = global_store_b32(addr=v[0:1], data=v[2], saddr=OFF)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
def test_global_load_b32_saddr(self):
# global_load_b32 v0, v1, s[0:1] (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:1])
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
def test_global_load_b32_offset(self):
# global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
# GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
inst = global_load_b32(vdst=v[0], addr=v[1:2], saddr=OFF, offset=256)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
def test_global_load_b64(self):
# global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
# GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
inst = global_load_b64(vdst=v[0:1], addr=v[2:3], saddr=OFF)
self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
class TestSMEM(unittest.TestCase):
"""Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
def test_smem_dlc_bit_position(self):
# s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
# GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
inst = s_load_b32(sdata=s[5], sbase=s[2:3], soffset=s[0], dlc=1)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_glc_bit_position(self):
# s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
# GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
inst = s_load_b32(sdata=s[5], sbase=s[2:3], soffset=s[0], glc=1)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_glc_dlc_combined(self):
# s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
# GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
inst = s_load_b32(sdata=s[5], sbase=s[2:3], soffset=s[0], glc=1, dlc=1)
self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
def test_smem_disasm_roundtrip_dlc(self):
# Test that disassembly/reassembly preserves DLC bit correctly
data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
decoded = SMEM.from_bytes(data)
self.assertEqual(decoded.to_bytes(), data)
def test_smem_disasm_roundtrip_glc_dlc(self):
# Test that disassembly/reassembly preserves GLC+DLC bits correctly
data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
decoded = SMEM.from_bytes(data)
self.assertEqual(decoded.to_bytes(), data)
class TestVOP3Literal(unittest.TestCase):
"""Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
def test_vop3_with_literal(self):
# v_add3_u32 v5, vcc_hi, 0xaf123456, v255
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=VCC_HI, src1=0xaf123456, src2=v[255])
expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
self.assertEqual(inst.to_bytes(), expected)
def test_vop3_literal_null_operand(self):
# v_add3_u32 v5, null, exec_lo, 0xaf123456
# GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=EXEC_LO, src2=0xaf123456)
expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
self.assertEqual(inst.to_bytes(), expected)
def test_vop3p_with_literal(self):
# Test VOP3P literal encoding (also uses Inst64)
inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=0.5, src1=0x12345678, src2=v[0])
self.assertEqual(len(inst.to_bytes()), 12) # 8 bytes + 4 byte literal
class TestDetectFormat(unittest.TestCase):
"""Test detect_format uses encoding from autogen classes."""
def test_detect_sopp(self):
self.assertEqual(detect_format(s_endpgm().to_bytes()), SOPP)
self.assertEqual(detect_format(s_nop(0).to_bytes()), SOPP)
self.assertEqual(detect_format(s_barrier().to_bytes()), SOPP)
def test_detect_sop1(self):
self.assertEqual(detect_format(s_mov_b32(s[0], 0).to_bytes()), SOP1)
self.assertEqual(detect_format(s_mov_b64(s[0:1], 0).to_bytes()), SOP1)
def test_detect_sop2(self):
self.assertEqual(detect_format(s_add_u32(s[0], s[1], s[2]).to_bytes()), SOP2)
self.assertEqual(detect_format(s_mul_i32(s[0], s[1], s[2]).to_bytes()), SOP2)
def test_detect_sopc(self):
self.assertEqual(detect_format(s_cmp_eq_i32(s[0], s[1]).to_bytes()), SOPC)
def test_detect_sopk(self):
self.assertEqual(detect_format(s_movk_i32(s[0], 0x1234).to_bytes()), SOPK)
def test_detect_vop1(self):
self.assertEqual(detect_format(v_mov_b32_e32(v[0], 0).to_bytes()), VOP1)
self.assertEqual(detect_format(v_rcp_f32_e32(v[0], v[1]).to_bytes()), VOP1)
def test_detect_vop2(self):
self.assertEqual(detect_format(v_add_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
self.assertEqual(detect_format(v_mul_f32_e32(v[0], v[1], v[2]).to_bytes()), VOP2)
def test_detect_vopc(self):
self.assertEqual(detect_format(v_cmp_eq_f32_e32(v[0], v[1]).to_bytes()), VOPC)
self.assertEqual(detect_format(v_cmp_lt_i32_e32(v[0], v[1]).to_bytes()), VOPC)
def test_detect_vop3(self):
self.assertEqual(detect_format(v_add_f32_e64(v[0], v[1], v[2]).to_bytes()), VOP3)
self.assertEqual(detect_format(v_fma_f32(v[0], v[1], v[2], v[3]).to_bytes()), VOP3)
def test_detect_vop3p(self):
self.assertEqual(detect_format(VOP3P(VOP3POp.V_PK_ADD_F16, v[0], v[1], v[2], v[3]).to_bytes()), VOP3P)
def test_detect_smem(self):
self.assertEqual(detect_format(s_load_b32(sdata=s[0], sbase=s[2:3], offset=0).to_bytes()), SMEM)
self.assertEqual(detect_format(s_load_b64(sdata=s[0:1], sbase=s[2:3], soffset=s[5]).to_bytes()), SMEM)
def test_detect_ds(self):
self.assertEqual(detect_format(ds_load_b32(v[0], v[1]).to_bytes()), DS)
self.assertEqual(detect_format(ds_store_b32(v[0], v[1]).to_bytes()), DS)
def test_detect_flat(self):
self.assertEqual(detect_format(global_load_b32(vdst=v[0], addr=v[1:2], saddr=NULL).to_bytes()), GLOBAL)
self.assertEqual(detect_format(global_store_b32(addr=v[0:1], data=v[2], saddr=NULL).to_bytes()), GLOBAL)
def test_detect_vopd(self):
inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=0, srcy0=0)
self.assertEqual(detect_format(inst.to_bytes()), VOPD)
def test_detect_vinterp(self):
inst = VINTERP(VINTERPOp.V_INTERP_P10_F32, vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
self.assertEqual(detect_format(inst.to_bytes()), VINTERP)
if __name__ == "__main__":
unittest.main()

197
tinygrad_repo/test/amd/test_handwritten.py Normal file
View File

@@ -0,0 +1,197 @@
# do not change these tests. we need to fix bugs to make them pass
# the Inst constructor should be looking at the types of the fields to correctly set the value
import unittest, struct
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from tinygrad.renderer.amd.dsl import Inst
from test.amd.test_roundtrip import compile_asm
from test.amd.disasm import disasm
class IntegrationTestBase(unittest.TestCase):
inst: Inst
arch: str
def tearDown(self):
if not hasattr(self, 'inst'): return
b = self.inst.to_bytes()
st = disasm(self.inst)
# Test that the instruction can be compiled by LLVM and produces the same bytes
desc = f"{st:25s} {self.inst} {b!r}"
self.assertEqual(b, compile_asm(st, arch=self.arch), desc)
print(desc)
class TestIntegration(IntegrationTestBase):
arch: str = "rdna3"
def test_wmma(self):
self.inst = v_wmma_f32_16x16x16_f16(v[0:7], v[184:191], v[136:143], v[0:7])
def test_load_b128(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
def test_load_b128_wrong_size(self):
# this should have to be 4 regs on the loaded to
with self.assertRaises(TypeError):
self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
def test_mov_b32(self):
self.inst = s_mov_b32(s[80], s[0])
def test_mov_b64(self):
self.inst = s_mov_b64(s[80:81], s[0:1])
def test_mov_b32_wrong(self):
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80:81], s[0:1])
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80:81], s[0])
with self.assertRaises(Exception):
self.inst = s_mov_b32(s[80], s[0:1])
def test_mov_b64_wrong(self):
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80], s[0])
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80], s[0:1])
with self.assertRaises(Exception):
self.inst = s_mov_b64(s[80:81], s[0])
def test_load_b128_no_0(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL)
def test_load_b128_s(self):
self.inst = s_load_b128(s[4:7], s[0:1], s[8], 0)
def test_load_b128_v(self):
with self.assertRaises(TypeError):
self.inst = s_load_b128(s[4:7], s[0:1], v[8], 0)
def test_load_b128_off(self):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, 3)
def test_simple_stos(self):
self.inst = s_mov_b32(s[0], s[1])
def test_simple_wrong(self):
with self.assertRaises(TypeError):
self.inst = s_mov_b32(v[0], s[1])
def test_simple_vtov(self):
self.inst = v_mov_b32_e32(v[0], v[1])
def test_simple_stov(self):
self.inst = v_mov_b32_e32(v[0], s[2])
def test_simple_float_to_v(self):
self.inst = v_mov_b32_e32(v[0], 1.0)
def test_simple_v_to_float(self):
with self.assertRaises(TypeError):
self.inst = v_mov_b32_e32(1, v[0])
def test_invalid_field(self):
with self.assertRaises(TypeError):
self.inst = s_load_b128(s[4:7], s[0:1], NULL, ioffset=0x8)
def test_simple_int_to_v(self):
self.inst = v_mov_b32_e32(v[0], 1)
def test_three_add(self):
self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3])
def test_three_add_v(self):
self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3])
def test_three_add_const(self):
self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3])
def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07)
def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7)
def test_vmad(self):
self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2])
def test_large_imm(self):
self.inst = v_mov_b32_e32(v[0], 0x1234)
def test_dual_mov(self):
self.inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], srcy0=v[4])
def test_dual_mul(self):
self.inst = v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
def test_simple_int_to_s(self):
self.inst = s_mov_b32(s[0], 3)
def test_complex_int_to_s(self):
self.inst = s_mov_b32(s[0], 0x235646)
def test_simple_float_to_s(self):
self.inst = s_mov_b32(s[0], 1.0)
def test_complex_float_to_s(self):
self.inst = s_mov_b32(s[0], 1337.0)
int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
self.assertEqual(self.inst, int_inst)
class TestIntegrationCDNA(IntegrationTestBase):
arch = "cdna"
def test_mfma(self):
from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x16_f16
self.inst = v_mfma_f32_16x16x16_f16(v[0:3], v[0:1], v[0:1], 0)
def test_mfma_fp8(self):
from tinygrad.runtime.autogen.amd.cdna.ins import v_mfma_f32_16x16x128_f8f6f4
self.inst = v_mfma_f32_16x16x128_f8f6f4(v[0:3], v[0:5], v[0:5], 1, cbsz=2, blgp=2)
class TestRegisterSliceSyntax(unittest.TestCase):
"""
Issue: Register slice syntax should use AMD assembly convention (inclusive end).
In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
The DSL should match this convention so that:
- s[4:7] gives 4 registers
- Disassembler output can be copied directly back into DSL code
Fix: Change _RegFactory.__getitem__ to use inclusive end:
key.stop - key.start + 1 (instead of key.stop - key.start)
"""
def test_register_slice_count(self):
# s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
reg = s[4:7]
self.assertEqual(reg.sz, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
def test_register_slice_roundtrip(self):
# Round-trip: DSL -> disasm -> DSL should preserve register count
reg = s[4:7] # 4 registers in AMD convention
inst = s_load_b128(reg, s[0:1], NULL, 0)
d = disasm(inst)
# Disasm shows s[4:7] - user should be able to copy this back
self.assertIn("s[4:7]", d)
# And s[4:7] in DSL should give the same 4 registers
reg_from_disasm = s[4:7]
self.assertEqual(reg_from_disasm.sz, 4, "s[4:7] from disasm should give 4 registers")
class TestInstructionEquality(unittest.TestCase):
"""
Issue: No __eq__ method - instruction comparison requires repr() workaround.
Two identical instructions should compare equal with ==, but currently:
inst1 == inst2 returns False
The test_handwritten.py works around this with:
self.assertEqual(repr(self.inst), repr(reasm))
"""
def test_identical_instructions_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[1])
self.assertEqual(inst1, inst2, "identical instructions should be equal")
def test_different_instructions_not_equal(self):
inst1 = v_mov_b32_e32(v[0], v[1])
inst2 = v_mov_b32_e32(v[0], v[2])
self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
if __name__ == "__main__":
unittest.main()

117
tinygrad_repo/test/amd/test_integration.py Normal file
View File

@@ -0,0 +1,117 @@
#!/usr/bin/env python3
"""Integration test: round-trip RDNA3 assembly through LLVM toolchain."""
import unittest
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from test.amd.helpers import llvm_assemble, llvm_disasm
def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
def assemble_and_disassemble(instructions: list, mcpu: str = "gfx1100", mattr: str = "+real-true16,+wavefrontsize32") -> list[str]:
"""Assemble instructions with our DSL, then disassemble with LLVM."""
code_bytes = b''.join(inst.to_bytes() for inst in instructions)
return llvm_disasm(code_bytes, mcpu, mattr)
class TestIntegration(unittest.TestCase):
"""Test our DSL output matches LLVM disassembly."""
def test_simple_sop1(self):
"""Test SOP1 instructions round-trip."""
instructions = [s_mov_b32(s[0], s[1]), s_mov_b32(s[2], 0), s_not_b32(s[3], s[4])]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_mov_b32', disasm[0])
self.assertIn('s_mov_b32', disasm[1])
self.assertIn('s_not_b32', disasm[2])
def test_simple_sop2(self):
"""Test SOP2 instructions round-trip."""
instructions = [s_add_u32(s[0], s[1], s[2]), s_sub_u32(s[3], s[4], 10), s_and_b32(s[5], s[6], s[7])]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_add_u32', disasm[0])
self.assertIn('s_sub_u32', disasm[1])
self.assertIn('s_and_b32', disasm[2])
def test_simple_vop2(self):
"""Test VOP2 instructions round-trip."""
instructions = [v_add_f32_e32(v[0], v[1], v[2]), v_mul_f32_e32(v[3], 1.0, v[4]), v_and_b32_e32(v[5], 10, v[6])]
disasm = assemble_and_disassemble(instructions)
self.assertIn('v_add_f32', disasm[0])
self.assertIn('v_mul_f32', disasm[1])
def test_control_flow(self):
"""Test control flow instructions."""
instructions = [s_waitcnt(simm16=waitcnt(lgkmcnt=0)), s_endpgm()]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_waitcnt', disasm[0])
self.assertIn('s_endpgm', disasm[1])
def test_memory_ops(self):
"""Test memory instructions."""
instructions = [s_load_b32(s[0], s[0:1], NULL), s_waitcnt(simm16=waitcnt(lgkmcnt=0)), global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
s_endpgm()]
disasm = assemble_and_disassemble(instructions)
self.assertIn('s_load_b32', disasm[0])
self.assertIn('s_waitcnt', disasm[1])
self.assertIn('global_store_b32', disasm[2])
def test_full_kernel(self):
"""Test a complete kernel similar to tinygrad output."""
instructions = [v_mov_b32_e32(v[0], s[0]), v_mov_b32_e32(v[1], s[1]), global_load_b32(vdst=v[2], addr=v[0:1], saddr=OFF),
s_waitcnt(simm16=waitcnt(vmcnt=0)), v_add_f32_e32(v[2], 1.0, v[2]), global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
s_endpgm()]
disasm = assemble_and_disassemble(instructions)
self.assertTrue(any('global_load' in d for d in disasm))
self.assertTrue(any('v_add_f32' in d for d in disasm))
self.assertTrue(any('global_store' in d for d in disasm))
self.assertTrue(any('s_endpgm' in d for d in disasm))
def test_bytes_roundtrip(self):
"""Test that our bytes match what LLVM assembler produces."""
inst = s_mov_b32(s[0], s[1])
our_bytes = inst.to_bytes()
llvm_bytes = llvm_assemble(["s_mov_b32 s0, s1"], "gfx1100", "+real-true16,+wavefrontsize32")[0]
self.assertEqual(our_bytes, llvm_bytes, f"Bytes mismatch: ours={our_bytes.hex()} LLVM={llvm_bytes.hex()}")
class TestTinygradIntegration(unittest.TestCase):
"""Test that we can parse tinygrad kernel disassembly."""
def _get_kernel_code(self, op_fn) -> bytes:
from tinygrad import Tensor
from tinygrad.helpers import Target
from tinygrad.codegen import to_program
from tinygrad.renderer.llvmir import AMDLLVMRenderer
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.uop.ops import Ops
result = op_fn(Tensor)
linear = result.schedule_linear()
sink_items = [call for call in linear.src if call.src[0].op == Ops.SINK]
assert len(sink_items) > 0, "No SINK in schedule"
renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100'))
prg = to_program(sink_items[0].src[0], renderer)
lib = renderer.compiler.compile(prg.src[3].arg)
return next(s.content for s in elf_loader(lib)[1] if s.name == ".text")
def test_simple_add_kernel(self):
"""Generate a simple add kernel from tinygrad and verify disassembly."""
code = self._get_kernel_code(lambda T: T([1.0, 2.0, 3.0, 4.0]).realize() + T([5.0, 6.0, 7.0, 8.0]).realize())
instrs = llvm_disasm(code, "gfx1100", "+real-true16,+wavefrontsize32")
self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
def test_matmul_kernel(self):
"""Generate a matmul kernel and verify disassembly has expected patterns."""
code = self._get_kernel_code(lambda T: T.rand(4, 4).realize() @ T.rand(4, 4).realize())
instrs = llvm_disasm(code, "gfx1100", "+real-true16,+wavefrontsize32")
has_mul = any('mul' in i.lower() for i in instrs)
has_add = any('add' in i.lower() for i in instrs)
self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
def test_disasm_to_bytes_roundtrip(self):
"""Verify s_endpgm encoding matches between our DSL and LLVM."""
our_bytes = s_endpgm().to_bytes()
llvm_bytes = llvm_assemble(["s_endpgm"], "gfx1100", "+real-true16,+wavefrontsize32")[0]
self.assertEqual(our_bytes, llvm_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} LLVM={llvm_bytes.hex()}")
if __name__ == "__main__":
unittest.main()

156
tinygrad_repo/test/amd/test_llvm.py Normal file
View File

@@ -0,0 +1,156 @@
#!/usr/bin/env python3
"""Test AMD assembler/disassembler against LLVM test vectors.
Only compute-relevant instruction formats are tested. Graphics-only formats not supported:
- MUBUF/MTBUF: buffer instructions with resource descriptors (use GLOBAL/FLAT instead)
- MIMG: image/texture instructions
- EXP/VEXPORT: export instructions for pixel/vertex output
- VIMAGE/VSAMPLE: image sampling instructions (RDNA4)
- VBUFFER: buffer instructions (RDNA4)
"""
import unittest, re, functools
from tinygrad.helpers import fetch
from test.amd.disasm import disasm
from tinygrad.renderer.amd import decode_inst, detect_format
from test.amd.helpers import llvm_assemble, llvm_filter_valid_asm, get_target, get_mattr
LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/llvmorg-21.1.0/llvm/test/MC/AMDGPU"
# RDNA3 (gfx11) test files for compute instructions
# Excluded: gfx11_asm_mubuf.s, gfx11_asm_mtbuf.s, gfx11_asm_mimg.s, gfx11_asm_mubuf_alias.s, gfx11_asm_mtbuf_alias.s (graphics-only)
RDNA_FILES = ['gfx11_asm_sop1.s', 'gfx11_asm_sop2.s', 'gfx11_asm_sopp.s', 'gfx11_asm_sopk.s', 'gfx11_asm_sopc.s',
'gfx11_asm_vop1.s', 'gfx11_asm_vop2.s', 'gfx11_asm_vopc.s', 'gfx11_asm_vop3.s', 'gfx11_asm_vop3p.s', 'gfx11_asm_vinterp.s',
'gfx11_asm_vopd.s', 'gfx11_asm_vopcx.s', 'gfx11_asm_vop3_from_vop1.s', 'gfx11_asm_vop3_from_vop2.s', 'gfx11_asm_vop3_from_vopc.s',
'gfx11_asm_vop3_from_vopcx.s', 'gfx11_asm_ds.s', 'gfx11_asm_smem.s', 'gfx11_asm_flat.s',
'gfx11_asm_wmma.s', 'gfx11_asm_vop3_features.s', 'gfx11_asm_vop3p_features.s', 'gfx11_asm_vopd_features.s',
'gfx11_asm_vop3_alias.s', 'gfx11_asm_vop3p_alias.s', 'gfx11_asm_vopc_alias.s', 'gfx11_asm_vopcx_alias.s', 'gfx11_asm_vinterp_alias.s',
'gfx11_asm_smem_alias.s']
# CDNA (gfx9/gfx90a/gfx942/gfx950) test files for compute instructions
# Excluded: gfx9_asm_mubuf.s, gfx9_asm_mtbuf.s, gfx90a_ldst_acc.s (has MIMG mixed in)
# Exclude gfx90a: 'gfx90a_asm_features.s', 'mai-gfx90a.s',
# Exclude gfx950: 'gfx950_asm_features.s' (disasm error)
CDNA_FILES = ['gfx9_asm_sop1.s', 'gfx9_asm_sop2.s', 'gfx9_asm_sopp.s', 'gfx9_asm_sopk.s', 'gfx9_asm_sopc.s',
'gfx9_asm_vop1.s', 'gfx9_asm_vop2.s', 'gfx9_asm_vopc.s', 'gfx9_asm_vop3.s', 'gfx9_asm_vop3p.s',
'gfx9_asm_ds.s', 'gfx9_asm_flat.s', 'gfx9_asm_smem.s',
'flat-scratch-gfx942.s', 'gfx942_asm_features.s', 'mai-gfx942.s',
'gfx950_asm_vop1.s', 'gfx950_asm_read_tr.s', 'mai-gfx950.s']
# RDNA4 (gfx12) test files for compute instructions
# Excluded: gfx12_asm_vbuffer_mubuf.s, gfx12_asm_vbuffer_mtbuf.s, gfx12_asm_exp.s (graphics-only)
RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx12_asm_sopk.s', 'gfx12_asm_sopc.s',
'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
'gfx12_asm_wmma_w32.s']
def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
tests = []
for block in text.split('\n\n'):
asm_text, encoding = None, None
for line in block.split('\n'):
line = line.strip()
if not line or line.startswith(('.', ';')): continue
if not line.startswith('//'):
asm_text = line.split('//')[0].strip() or asm_text
if m := re.search(pattern + r'[^:]*:.*?(?:encoding:\s*)?\[(0x[0-9a-f,x\s]+)\]', line, re.I):
encoding = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
if asm_text and encoding:
try: tests.append((asm_text, bytes.fromhex(encoding)))
except ValueError: pass
return tests
def _get_tests_uncached(f: str, arch: str) -> list[tuple[str, bytes]]:
text = fetch(f"{LLVM_BASE}/{f}").read_bytes().decode('utf-8', errors='ignore')
if arch == "rdna3":
# Match GFX11 and W32 only (wavefront32 mode)
tests = _parse_llvm_tests(text, r'(?:GFX11|W32)')
elif arch == "rdna4":
# Match GFX12 (but not GFX1250) and W32 only (wavefront32 mode)
tests = _parse_llvm_tests(text, r'(?:GFX12(?!50)|W32)')
elif 'gfx90a' in f or 'gfx942' in f or 'gfx950' in f:
tests = _parse_llvm_tests(text, r'(?:GFX90A|GFX942|GFX950)')
else:
tests = _parse_llvm_tests(text, r'(?:VI9|GFX9|CHECK)')
# Exclude v_interp_* (graphics-only, not on CDNA)
if arch == "cdna": tests = [(asm, data) for asm, data in tests if not asm.startswith('v_interp_')]
# Filter out tests where original ASM isn't valid on target (e.g., gfx9 tests with gfx942/gfx950 constraints)
if arch == "cdna" and not ('gfx942' in f or 'gfx950' in f or 'gfx90a' in f):
tests = llvm_filter_valid_asm(tests, get_target(arch), get_mattr(arch))
return tests
@functools.cache
def _get_tests(f: str, arch: str) -> list[tuple[str, bytes]]: return _get_tests_uncached(f, arch)
def _make_test(f: str, arch: str, test_type: str):
def test(self):
tests = _get_tests(f, arch)
name = f"{arch}_{test_type}_{f}"
mcpu = "gfx942" if arch == "cdna" and "gfx942" in f else get_target(arch)
if test_type == "roundtrip":
passed, skipped = 0, 0
for _, data in tests:
try:
decoded = detect_format(data, arch).from_bytes(data)
self.assertEqual(decoded.to_bytes()[:len(data)], data)
passed += 1
except ValueError: skipped += 1 # skip invalid opcodes not in enum
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "repr":
# Test that eval(repr(inst)) reproduces the instruction
if arch == "rdna3": import tinygrad.runtime.autogen.amd.rdna3.ins as ins # type: ignore[no-redef]
elif arch == "rdna4": import tinygrad.runtime.autogen.amd.rdna4.ins as ins # type: ignore[no-redef]
elif arch == "cdna": import tinygrad.runtime.autogen.amd.cdna.ins as ins # type: ignore[no-redef]
ns = {k: getattr(ins, k) for k in dir(ins) if not k.startswith('_')}
passed, skipped = 0, 0
for _, data in tests:
try:
decoded = detect_format(data, arch).from_bytes(data)
if decoded.to_bytes()[:len(data)] != data:
skipped += 1
continue # skip if binary roundtrip fails
r = repr(decoded)
try:
decoded2 = eval(r, ns) # noqa: S307
if decoded == decoded2: passed += 1
else: skipped += 1
except Exception: skipped += 1
except ValueError: skipped += 1
print(f"{name}: {passed} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
elif test_type == "disasm":
to_test = []
for _, data in tests:
try:
decoded = decode_inst(data, arch)
enc = decoded.to_bytes()[:len(data)]
# Skip if roundtrip fails, disasm fails, or op_name is missing (disasm starts with space)
if enc == data and (d := disasm(decoded)) and not d.startswith(' '): to_test.append((enc, d))
except Exception: pass
skipped = len(tests) - len(to_test)
print(f"{name}: {len(to_test)} passed, {skipped} skipped")
self.assertEqual(skipped, 0, f"{name}: {skipped} tests skipped, expected 0")
# Compare disasm->reassemble with original encoding (filter reserved bit cases where LLVM can't reproduce)
llvm_bytes = llvm_assemble([t[1] for t in to_test], mcpu, get_mattr(arch))
valid = [(enc, d, llvm) for (enc, d), llvm in zip(to_test, llvm_bytes) if llvm == enc]
print(f"{name}: {len(valid)}/{len(to_test)} matched LLVM encoding")
for enc, _, llvm in valid: self.assertEqual(llvm, enc)
return test
class TestLLVM(unittest.TestCase): pass
for f in RDNA_FILES:
setattr(TestLLVM, f"test_rdna3_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "roundtrip"))
setattr(TestLLVM, f"test_rdna3_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "disasm"))
setattr(TestLLVM, f"test_rdna3_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna3", "repr"))
for f in CDNA_FILES:
setattr(TestLLVM, f"test_cdna_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "roundtrip"))
setattr(TestLLVM, f"test_cdna_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "disasm"))
setattr(TestLLVM, f"test_cdna_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "cdna", "repr"))
for f in RDNA4_FILES:
setattr(TestLLVM, f"test_rdna4_roundtrip_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "roundtrip"))
setattr(TestLLVM, f"test_rdna4_disasm_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "disasm"))
setattr(TestLLVM, f"test_rdna4_repr_{f.replace('.s', '').replace('-', '_')}", _make_test(f, "rdna4", "repr"))
if __name__ == "__main__":
unittest.main()

50
tinygrad_repo/test/amd/test_mockgpu_invalid.py Normal file
View File

@@ -0,0 +1,50 @@
#!/usr/bin/env python3
"""Test that invalid instructions raise exceptions through the mock GPU stack."""
import unittest, subprocess, os, sys, time
class TestMockGPUInvalidInstruction(unittest.TestCase):
def test_unsupported_instruction_raises(self):
"""Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
test_code = '''
import struct
from tinygrad import Device, Tensor
from tinygrad.engine.realize import compile_linear
from tinygrad.runtime.ops_amd import AMDProgram
dev = Device["AMD"]
a = Tensor([1.0]).realize()
b = a + 1
linear = compile_linear(b.schedule_linear())
lib = bytearray(linear.src[-1].src[0].src[4].arg)
# Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
# VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
found = False
for i in range(0, len(lib) - 4, 4):
if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
lib[i:i+4] = struct.pack("<I", 0x7E008400)
found = True
break
assert found, "s_endpgm not found"
patched_prg = AMDProgram(dev, "patched", bytes(lib))
b.uop.buffer.allocate()
patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
dev.synchronize()
'''
env = os.environ.copy()
env["DEV"] = "MOCKKFD+AMD"
env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
st = time.perf_counter()
result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
elapsed = time.perf_counter() - st
self.assertNotEqual(result.returncode, 0, "should have raised")
self.assertTrue("Error" in result.stderr, f"expected an error in stderr, got: {result.stderr[:500]}")
# Should exit immediately, not wait for the full timeout
self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
if __name__ == "__main__":
unittest.main()

53
tinygrad_repo/test/amd/test_pdf.py Normal file
View File

@@ -0,0 +1,53 @@
#!/usr/bin/env python3
"""Test PDF pseudocode extraction from generate.py."""
import unittest
from tinygrad.renderer.amd.generate import extract_pdf_text, extract_pcode, parse_xml, ARCHS, FIXES
EXPECTED_PAGES = {"rdna3": 655, "rdna4": 711, "cdna": 610}
class TestPcodePDF(unittest.TestCase):
pages: dict
enums: dict
pcode: dict
@classmethod
def setUpClass(cls):
cls.pages = {arch: extract_pdf_text(cfg["pdf"]) for arch, cfg in ARCHS.items()}
cls.enums = {}
for arch, cfg in ARCHS.items():
_, enums, _, _, _, _ = parse_xml(cfg["xml"])
for fmt, ops in FIXES.get(arch, {}).items(): enums.setdefault(fmt, {}).update(ops)
cls.enums[arch] = enums
cls.pcode = {arch: extract_pcode(cls.pages[arch], {n: op for ops in cls.enums[arch].values() for op, n in ops.items()}) for arch in ARCHS}
def test_page_counts(self):
for name, exp in EXPECTED_PAGES.items():
self.assertEqual(len(self.pages[name]), exp, f"{name} page count")
def test_pcode_extracted(self):
"""Check we extracted a reasonable number of pcode entries."""
for name in ARCHS:
self.assertGreater(len(self.pcode[name]), 500, f"{name} pcode count too low")
def test_pcode_rdna3_tricky(self):
"""Test specific pseudocode patterns that are tricky to extract correctly."""
pcode = self.pcode['rdna3']
# BUFFER_ATOMIC_MAX_U64: should have 4 statements (not truncated)
self.assertEqual(pcode[('BUFFER_ATOMIC_MAX_U64', 72)],
'tmp = MEM[ADDR].u64;\nsrc = DATA.u64;\nMEM[ADDR].u64 = src >= tmp ? src : tmp;\nRETURN_DATA.u64 = tmp')
# GLOBAL_STORE_B128: should have 4 MEM stores (not truncated)
self.assertEqual(pcode[('GLOBAL_STORE_B128', 29)],
'MEM[ADDR].b32 = VDATA[31 : 0];\nMEM[ADDR + 4U].b32 = VDATA[63 : 32];\n'
'MEM[ADDR + 8U].b32 = VDATA[95 : 64];\nMEM[ADDR + 12U].b32 = VDATA[127 : 96]')
# S_CMOVK_I32: should have full if/endif block
self.assertEqual(pcode[('S_CMOVK_I32', 2)],
"if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
def test_pcode_no_examples(self):
"""Pseudocode should not contain example lines with '=>'."""
for name in ARCHS:
for (op_name, opcode), code in self.pcode[name].items():
self.assertNotIn('=>', code, f"{name} {op_name} contains example line with '=>'")
if __name__ == "__main__":
unittest.main()

66
tinygrad_repo/test/amd/test_rdna3_asm.py Normal file
View File

@@ -0,0 +1,66 @@
#!/usr/bin/env python3
import unittest
from tinygrad.runtime.autogen.amd.rdna3.ins import *
from test.amd.helpers import llvm_assemble
from test.amd.disasm import disasm
def _asm(asm: str) -> bytes: return llvm_assemble([asm], 'gfx1100', '+real-true16,+wavefrontsize32')[0]
class TestRDNA3Asm(unittest.TestCase):
def test_full_program(self):
"""Test the full program from rdna3fun.py matches LLVM output."""
program = [
v_bfe_u32(v[1], v[0], 10, 10),
s_load_b128(s[4:7], s[0:1], NULL),
v_and_b32_e32(v[0], 0x3FF, v[0]),
s_mulk_i32(s[3], 0x87),
v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
v_mul_u32_u24_e32(v[0], 45, v[0]),
v_ashrrev_i32_e32(v[2], 31, v[1]),
v_add3_u32(v[0], v[0], s[3], v[1]),
v_lshlrev_b64(v[2:3], 2, v[1:2]),
v_ashrrev_i32_e32(v[1], 31, v[0]),
v_lshlrev_b64(v[0:1], 2, v[0:1]),
s_waitcnt(0xfc07), # lgkmcnt(0)
v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
v_add_co_ci_u32_e32(v[3], s[7], v[3]),
v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
global_load_b32(vdst=v[2], addr=v[2:3], saddr=OFF),
v_add_co_ci_u32_e32(v[1], s[5], v[1]),
s_waitcnt(0x03f7), # vmcnt(0)
global_store_b32(addr=v[0:1], data=v[2], saddr=OFF),
s_endpgm(),
]
asm_lines = [
"v_bfe_u32 v1, v0, 10, 10", "s_load_b128 s[4:7], s[0:1], null", "v_and_b32_e32 v0, 0x3FF, v0",
"s_mulk_i32 s3, 0x87", "v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]", "v_mul_u32_u24_e32 v0, 45, v0",
"v_ashrrev_i32_e32 v2, 31, v1", "v_add3_u32 v0, v0, s3, v1", "v_lshlrev_b64 v[2:3], 2, v[1:2]",
"v_ashrrev_i32_e32 v1, 31, v0", "v_lshlrev_b64 v[0:1], 2, v[0:1]", "s_waitcnt lgkmcnt(0)",
"v_add_co_u32 v2, vcc_lo, s6, v2", "v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo",
"v_add_co_u32 v0, vcc_lo, s4, v0", "global_load_b32 v2, v[2:3], off",
"v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo", "s_waitcnt vmcnt(0)",
"global_store_b32 v[0:1], v2, off", "s_endpgm",
]
expected = llvm_assemble(asm_lines, 'gfx1100', '+real-true16,+wavefrontsize32')
for inst, rt in zip(program, asm_lines): print(f"{disasm(inst):50s} {rt}")
for inst, exp in zip(program, expected): self.assertEqual(inst.to_bytes(), exp)
def test_sop2_s_add_u32(self):
inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
self.assertEqual(inst.to_bytes(), _asm("s_add_u32 s3, s0, s1"))
def test_vop2_v_and_b32_inline_const(self):
inst = v_and_b32_e32(v[0], 10, v[0])
self.assertEqual(inst.to_bytes(), _asm("v_and_b32_e32 v0, 10, v0"))
def test_sopp_s_endpgm(self):
inst = s_endpgm()
self.assertEqual(inst.to_bytes(), _asm("s_endpgm"))
def test_sop1_s_mov_b32(self):
inst = s_mov_b32(s[0], s[1])
self.assertEqual(inst.to_bytes(), _asm("s_mov_b32 s0, s1"))
if __name__ == "__main__":
unittest.main()

279
tinygrad_repo/test/amd/test_roundtrip.py Normal file
View File

@@ -0,0 +1,279 @@
#!/usr/bin/env python3
"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
import unittest, io, sys, re
from dataclasses import dataclass
from tinygrad import Device
from tinygrad.renderer.amd import detect_format
from test.amd.helpers import llvm_assemble, llvm_disasm, get_target, get_mattr
from test.amd.disasm import disasm
def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
"""Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
old_stdout = sys.stdout
sys.stdout = io.StringIO()
compiler.disassemble(lib)
output = sys.stdout.getvalue()
sys.stdout = old_stdout
results = []
for line in output.splitlines():
if '//' not in line: continue
instr = line.split('//')[0].strip()
if not instr: continue
comment = line.split('//')[1].strip()
if ':' not in comment: continue
hex_str = comment.split(':')[1].strip().split()[0]
try:
machine_bytes = bytes.fromhex(hex_str)[::-1] # big-endian to little-endian
results.append((instr, machine_bytes))
except ValueError:
continue
return results
def compile_asm(instr: str, arch: str = 'rdna3') -> bytes:
"""Compile a single instruction using LLVM."""
return llvm_assemble([instr], get_target(arch), get_mattr(arch))[0]
def compile_asm_batch(instrs: list[str], arch: str = 'rdna3') -> list[bytes]:
"""Compile multiple instructions with a single LLVM emission."""
return llvm_assemble(instrs, get_target(arch), get_mattr(arch))
def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str]:
"""Compile instructions with LLVM and get LLVM's disassembly."""
if not instrs: return []
mcpu, mattr = get_target(arch), get_mattr(arch)
code = b''.join(llvm_assemble(instrs, mcpu, mattr))
return llvm_disasm(code, mcpu, mattr)[:len(instrs)]
@dataclass
class KernelSnapshot:
code: bytes
src: str
global_size: tuple[int, int, int]
local_size: tuple[int, int, int]
buf_idxs: list[int] # indices into shared buffer pool
buf_sizes: list[int] # sizes for each buffer index
def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]:
"""Compile a tinygrad operation and extract all kernels with their buffer mappings."""
from tinygrad import Tensor
from tinygrad.uop.ops import Ops
from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi
from tinygrad.runtime.support.elf import elf_loader
out = op_fn(Tensor)
linear = compile_linear(out.schedule_linear())
kernels = []
buf_pool: dict[int, int] = {} # buffer id -> size
buf_data: dict[int, bytes] = {} # buffer id -> initial data from COPY
for call in linear.src:
ast = call.src[0]
for bufs, _ in unwrap_multi(call, resolve_params(call, ())):
if ast.op is Ops.COPY:
# Handle COPY: extract source data to initialize destination buffer
if len(bufs) >= 2:
dst_buf, src_buf = bufs[0], bufs[1]
dst_id = id(dst_buf)
if dst_id not in buf_pool:
buf_pool[dst_id] = dst_buf.nbytes
# Get source data if it's from numpy/CPU
if hasattr(src_buf, 'base') and src_buf.base is not None and src_buf.base.is_allocated():
src_data = bytes(src_buf.base._buf)
buf_data[dst_id] = src_data
elif ast.op is Ops.PROGRAM:
info = ast.arg
if len(ast.src) > 4 and ast.src[4].op is Ops.BINARY:
lib = bytes(ast.src[4].arg)
_, sections, _ = elf_loader(lib)
for sec in sections:
if sec.name == '.text':
buf_idxs = []
buf_sizes = []
for b in bufs:
buf_id = id(b)
if buf_id not in buf_pool:
buf_pool[buf_id] = b.nbytes
buf_idxs.append(buf_id)
buf_sizes.append(b.nbytes)
kernels.append(KernelSnapshot(
code=bytes(sec.content),
src=ast.src[3].arg,
global_size=tuple(info.global_size),
local_size=tuple(info.local_size),
buf_idxs=buf_idxs,
buf_sizes=buf_sizes
))
if not kernels: raise RuntimeError("No kernel found")
return kernels, buf_pool, buf_data
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
class TestTinygradKernelRoundtrip(unittest.TestCase):
"""Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
arch = 'rdna3'
def _test_kernel_roundtrip(self, op_fn):
"""Generate kernel from op_fn, test:
1. decode -> reencode matches original bytes
2. disasm() -> LLVM asm -> bytes matches original (validates disasm correctness)
3. our disasm() matches LLVM's disassembly string (informational)
"""
arch = self.arch
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
from tinygrad.helpers import DEV
kernels, _, _ = get_kernels_from_tinygrad(op_fn)
# rendered source can be C or llvmir
compiler = (AMDLLVMCompiler if DEV.renderer == "LLVM" else HIPCompiler)(get_target(arch))
# First pass: decode all instructions and collect info
decoded_instrs: list[tuple] = [] # list of (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err)
for ki, kernel in enumerate(kernels):
offset = 0
code = next((s.content for s in elf_loader(compiler.compile(kernel.src))[1] if s.name == ".text"))
while offset < len(code):
remaining = code[offset:]
fmt = detect_format(remaining, arch)
base_size = fmt._size()
if len(remaining) < base_size:
break
try:
decoded = fmt.from_bytes(remaining) # pass all remaining bytes so from_bytes can read literal
size = decoded.size() # actual size including literal
orig_bytes = remaining[:size]
reencoded = decoded.to_bytes()
our_disasm = disasm(decoded)
decode_ok = reencoded == orig_bytes
decode_err: str | None = None if decode_ok else f"orig={orig_bytes.hex()} reenc={reencoded.hex()}"
decoded_instrs.append((ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err))
except Exception as e:
decoded_instrs.append((ki, offset, remaining[:base_size], None, None, False, str(e)))
size = base_size
offset += size
# Collect disasm strings for batched LLVM calls - skip unknown opcodes (op_X) that LLVM can't compile
asm_test_instrs: list[tuple[int, str, bytes]] = [] # (idx, our_disasm, orig_bytes) for asm test
disasm_test_instrs: list[tuple[int, str]] = [] # (idx, our_disasm) for disasm comparison test
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
if our_disasm is None: continue
# Skip unknown opcodes and malformed instructions
if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm): continue
asm_test_instrs.append((idx, our_disasm, orig_bytes))
disasm_test_instrs.append((idx, our_disasm))
# Batch compile for asm test (our disasm -> LLVM asm -> bytes)
asm_llvm_results = compile_asm_batch([d for _, d, _ in asm_test_instrs], arch)
asm_llvm_map = {idx: (result, orig) for (idx, _, orig), result in zip(asm_test_instrs, asm_llvm_results)}
# Batch compile+disasm for disasm comparison test
disasm_llvm_results = compile_and_disasm_batch([d for _, d in disasm_test_instrs], arch)
disasm_llvm_map = {idx: result for (idx, _), result in zip(disasm_test_instrs, disasm_llvm_results)}
# Now evaluate results
decode_passed, decode_failed, decode_skipped = 0, 0, 0
asm_passed, asm_failed, asm_skipped = 0, 0, 0
disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
decode_failures: list[str] = []
asm_failures: list[str] = []
disasm_failures: list[str] = []
for idx, (ki, offset, orig_bytes, decoded, our_disasm, decode_ok, decode_err) in enumerate(decoded_instrs):
# Decode test
if decode_ok:
decode_passed += 1
elif decode_err == "no format":
decode_skipped += 1
else:
decode_failed += 1
decode_failures.append(f"K{ki}@{offset}: {our_disasm}: {decode_err}")
# Asm test: our disasm -> LLVM asm -> compare bytes with original
if our_disasm is None:
asm_skipped += 1
elif idx in asm_llvm_map:
llvm_bytes, orig = asm_llvm_map[idx]
if llvm_bytes == orig[:len(llvm_bytes)]:
asm_passed += 1
else:
asm_failed += 1
asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': llvm={llvm_bytes.hex()} orig={orig[:len(llvm_bytes)].hex()}")
else:
asm_skipped += 1
# Disasm comparison test
if our_disasm is None:
disasm_skipped += 1
elif idx in disasm_llvm_map:
llvm_disasm_str = disasm_llvm_map[idx]
if our_disasm == llvm_disasm_str:
disasm_passed += 1
else:
disasm_failed += 1
disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm_str}'")
else:
disasm_skipped += 1
print(f"[{arch}] decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
print(f"[{arch}] asm via llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
print(f"[{arch}] disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
self.assertEqual(decode_failed, 0, "Decode failures:\n" + "\n".join(decode_failures[:20]))
self.assertEqual(asm_failed, 0, "Asm failures:\n" + "\n".join(asm_failures[:20]))
# Note: disasm string comparison is informational only - formatting differences between LLVM versions are expected
# Basic unary ops
def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
# Binary ops
def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
# Reductions
def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
# Matmul
def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
# Complex ops
def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
# Memory patterns
def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
# Cast ops
def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
# Comparison ops
def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
# Fused ops
def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
class TestTinygradKernelRoundtripRDNA4(TestTinygradKernelRoundtrip): arch = 'rdna4'
@unittest.skip("CDNA decode roundtrip not yet supported")
class TestTinygradKernelRoundtripCDNA(TestTinygradKernelRoundtrip): arch = 'cdna'
if __name__ == "__main__":
unittest.main()

108
tinygrad_repo/test/amd/test_sqtt_encoder.py Normal file
View File

@@ -0,0 +1,108 @@
#!/usr/bin/env python3
"""Tests for SQTT encoder: verifies the emulator produces correct SQTT traces for known kernels.
Run with: DEV=MOCK+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
"""
import ctypes, unittest
from tinygrad.helpers import Context
from tinygrad.renderer.amd.sqtt import decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST, InstOp
from tinygrad.runtime.autogen.amd.rdna3.ins import *
def _run_kernel(instructions: list, lx=1, ly=1, lz=1, gx=1, gy=1, gz=1, args_ptr=0) -> bytes:
"""Assemble instructions, run on emulator with PROFILE=1, return the SQTT blob."""
from test.mockgpu.amd.emu import run_asm, sqtt_traces
code = b''.join(inst.to_bytes() for inst in instructions)
buf = (ctypes.c_char * len(code))(*code)
lib = ctypes.addressof(buf)
sqtt_traces.clear()
with Context(PROFILE=1):
run_asm(lib, len(code), gx, gy, gz, lx, ly, lz, args_ptr)
assert len(sqtt_traces) == 1, f"expected 1 trace, got {len(sqtt_traces)}"
return sqtt_traces.pop()
class TestSQTTEncoder(unittest.TestCase):
def test_simple_salu(self):
"""A simple s_mov + s_endpgm kernel emits SALU INST packet."""
blob = _run_kernel([s_mov_b32(s[0], 42), s_endpgm()])
packets = list(decode(blob))
inst_pkts = [p for p in packets if isinstance(p, INST)]
self.assertEqual(len(inst_pkts), 1)
self.assertEqual(inst_pkts[0].op, InstOp.SALU)
def test_valu_emits_valuinst(self):
"""Regular VALU ops emit VALUINST packets."""
blob = _run_kernel([v_mov_b32_e32(v[0], 0), v_add_f32_e32(v[1], v[0], v[0]), s_endpgm()])
packets = list(decode(blob))
valu_pkts = [p for p in packets if isinstance(p, VALUINST)]
self.assertEqual(len(valu_pkts), 2)
# no INST packets for regular VALU
self.assertEqual(len([p for p in packets if isinstance(p, INST)]), 0)
def test_waitcnt_emits_immediate(self):
"""s_waitcnt and s_nop emit IMMEDIATE packets."""
blob = _run_kernel([s_nop(simm16=0), s_waitcnt(simm16=0), s_endpgm()])
imm_pkts = [p for p in decode(blob) if isinstance(p, IMMEDIATE)]
self.assertEqual(len(imm_pkts), 2) # s_nop + s_waitcnt
def test_endpgm_skipped(self):
"""s_endpgm does not emit any packet."""
blob = _run_kernel([s_endpgm()])
packets = list(decode(blob))
self.assertEqual(len([p for p in packets if isinstance(p, INST)]), 0)
self.assertEqual(len([p for p in packets if isinstance(p, IMMEDIATE)]), 0)
def test_wave_lifecycle(self):
"""Every WAVESTART has a matching WAVEEND."""
blob = _run_kernel([s_mov_b32(s[0], 0), s_endpgm()])
packets = list(decode(blob))
self.assertEqual(sum(1 for p in packets if isinstance(p, WAVESTART)), sum(1 for p in packets if isinstance(p, WAVEEND)))
def test_layout_header(self):
"""First packet is LAYOUT_HEADER with layout=3."""
blob = _run_kernel([s_endpgm()])
packets = list(decode(blob))
self.assertIsInstance(packets[0], LAYOUT_HEADER)
self.assertEqual(packets[0].layout, 3)
def test_blob_32byte_aligned(self):
"""SQTT blob is 32-byte aligned."""
blob = _run_kernel([s_mov_b32(s[0], 0), s_mov_b32(s[1], 1), s_endpgm()])
self.assertEqual(len(blob) % 32, 0)
def test_multiple_waves(self):
"""Multiple wavefronts each get their own WAVESTART/WAVEEND."""
blob = _run_kernel([s_mov_b32(s[0], 0), s_endpgm()], lx=64) # 64 threads = 2 waves (WAVE_SIZE=32)
packets = list(decode(blob))
self.assertEqual(sum(1 for p in packets if isinstance(p, WAVESTART)), 2)
self.assertEqual(sum(1 for p in packets if isinstance(p, WAVEEND)), 2)
def test_branch_taken_and_not_taken(self):
"""A loop with s_cbranch_scc1 emits JUMP when taken, JUMP_NO on final iteration."""
# s[0] = 2; loop: s[0] -= 1; cmp s[0] != 0 (SCC=1 if true); cbranch_scc1 loop; endpgm
# iteration 1: s[0]=2→1, SCC=1 (1!=0), branch taken (JUMP)
# iteration 2: s[0]=1→0, SCC=0 (0==0), branch not taken (JUMP_NO)
blob = _run_kernel([s_mov_b32(s[0], 2), s_sub_u32(s[0], s[0], 1), s_cmp_lg_u32(s[0], 0), s_cbranch_scc1(simm16=-3), s_endpgm()])
inst_pkts = [p for p in decode(blob) if isinstance(p, INST)]
ops = [p.op for p in inst_pkts]
self.assertIn(InstOp.JUMP, ops)
self.assertIn(InstOp.JUMP_NO, ops)
def test_timestamps_monotonic(self):
"""Timestamps are monotonically non-decreasing."""
blob = _run_kernel([s_mov_b32(s[0], 0), s_mov_b32(s[1], 1), s_mov_b32(s[2], 2), s_endpgm()])
times = [p._time for p in decode(blob)]
self.assertEqual(times, sorted(times))
def test_no_trace_without_profile(self):
"""No SQTT trace is emitted when PROFILE=0."""
from test.mockgpu.amd.emu import run_asm, sqtt_traces
code = s_endpgm().to_bytes()
buf = (ctypes.c_char * len(code))(*code)
sqtt_traces.clear()
with Context(PROFILE=0):
run_asm(ctypes.addressof(buf), len(code), 1, 1, 1, 1, 1, 1, 0)
self.assertEqual(len(sqtt_traces), 0)
if __name__ == "__main__":
unittest.main()

235
tinygrad_repo/test/amd/test_sqtt_examples.py Normal file
View File

@@ -0,0 +1,235 @@
#!/usr/bin/env python3
"""Tests for SQTT packet decoding using real captured examples."""
import pickle, unittest, ctypes, threading
from pathlib import Path
from tinygrad.helpers import DEBUG
from tinygrad.runtime.autogen import rocprof
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.renderer.amd import decode_inst
from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, WAVEEND_RDNA4, INST, INST_RDNA4, VALUINST,
IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
print_packets, CDNA_WAVEEND, CDNA_INST)
from test.amd.helpers import TARGET_TO_ARCH
from test.amd.test_sqttmap import needs_rocprof
import tinygrad
EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
# ═══════════════════════════════════════════════════════════════════════════════
# ROCPROF DECODER
# ═══════════════════════════════════════════════════════════════════════════════
def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
"""Run rocprof decoder on SQTT blobs, returning raw occupancy and instruction records."""
image, sections, _ = elf_loader(lib)
text = next((sh for sh in sections if sh.name == ".text"), None)
assert text is not None, "no .text section found"
text_off, text_size = text.header.sh_addr, text.header.sh_size
blob_iter, current_blob = iter(blobs), [None] # type: ignore[var-annotated]
occupancy_records: list[tuple[int, int, int, int, bool]] = [] # (wave_id, simd, cu, time, is_start)
wave_insts: list[list[tuple[int, int]]] = [] # per-wave list of (time, stall)
@rocprof.rocprof_trace_decoder_se_data_callback_t
def copy_cb(buf, buf_size, _): # type: ignore[no-untyped-def]
blob = next(blob_iter, None)
if blob is None: return 0
current_blob[0] = (ctypes.c_ubyte * len(blob)).from_buffer_copy(blob) # type: ignore[call-overload]
buf[0] = ctypes.cast(current_blob[0], ctypes.POINTER(ctypes.c_ubyte)) # type: ignore[arg-type]
buf_size[0] = len(current_blob[0]) # type: ignore[arg-type]
return len(current_blob[0]) # type: ignore[arg-type]
@rocprof.rocprof_trace_decoder_trace_callback_t
def trace_cb(record_type, events_ptr, n, _):
if record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr):
occupancy_records.append((ev.wave_id, ev.simd, ev.cu, ev.time, ev.start))
elif record_type == rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr):
if ev.instructions_size > 0:
sz = ev.instructions_size * ctypes.sizeof(rocprof.rocprofiler_thread_trace_decoder_inst_t)
insts_blob = bytearray(sz)
ctypes.memmove((ctypes.c_char * sz).from_buffer(insts_blob), ev.instructions_array, sz)
insts = list((rocprof.rocprofiler_thread_trace_decoder_inst_t * ev.instructions_size).from_buffer(insts_blob))
wave_insts.append([(inst.time, inst.stall) for inst in insts])
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
arch = TARGET_TO_ARCH[target]
@rocprof.rocprof_trace_decoder_isa_callback_t
def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
offset = pc.address - base
if offset < text_off or offset >= text_off + text_size:
mem_size_ptr[0] = 0
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
try:
inst = decode_inst(image[offset:], arch=arch)
mem_size_ptr[0] = inst._size()
# this could be an error in our decode_inst
except (ValueError, AssertionError):
mem_size_ptr[0] = 0
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
if isinstance(inst, SOPP) and inst.op == SOPPOp.S_ENDPGM: mem_size_ptr[0] = 0
# rocprof parses instruction string to determine type; v_nop works for all
if (max_sz := size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
ctypes.memmove(instr_ptr, b"v_nop", min(5, max_sz - 1))
size_ptr[0] = min(5, max_sz - 1)
return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
exc = None
def worker():
nonlocal exc
try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
except Exception as e: exc = e
(t:=threading.Thread(target=worker, daemon=True)).start()
t.join(timeout=5)
if exc is not None: raise exc
if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
return occupancy_records, wave_insts
class SQTTExamplesTestBase(unittest.TestCase):
target: str
examples: dict
@classmethod
def setUpClass(cls):
if cls is SQTTExamplesTestBase: raise unittest.SkipTest("base class")
cls.examples = {}
for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
with open(pkl_path, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
if sqtt_events and prg:
cls.examples[pkl_path.stem] = (sqtt_events, prg.lib, prg.base)
def test_examples_loaded(self):
self.assertGreater(len(self.examples), 0, "no example files found")
def test_decode_all_examples(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
packets = list(decode(event.blob))
if DEBUG >= 2:
print(f"\n=== {name} event {i} ===")
print_packets(packets)
self.assertGreater(len(packets), 0, f"no packets decoded from {name} event {i}")
self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")
def test_packet_types_valid(self):
all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values()) | set(PACKET_TYPES_CDNA.values())
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
for pkt in decode(event.blob):
# Use isinstance to handle layout-specific subclasses (e.g., WAVESTART_RDNA4)
self.assertTrue(any(isinstance(pkt, cls) for cls in all_classes), f"unknown packet type {type(pkt)} in {name}")
def test_wave_lifecycle(self):
for name, (events, *_) in self.examples.items():
if "empty" in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
def test_time_monotonic(self):
for name, (events, *_) in self.examples.items():
for i, event in enumerate(events):
with self.subTest(example=name, event=i):
times = [p._time for p in decode(event.blob)]
self.assertEqual(times, sorted(times), f"timestamps not monotonic in {name}")
def test_gemm_has_instructions(self):
for name, (events, *_) in self.examples.items():
if "gemm" not in name: continue
with self.subTest(example=name):
all_packets = [p for e in events for p in decode(e.blob)]
inst_packets = [p for p in all_packets if isinstance(p, (INST, INST_RDNA4, CDNA_INST))]
self.assertGreater(len(inst_packets), 0, f"no INST packets in {name}")
if isinstance(inst_packets[0], (INST, INST_RDNA4)):
self.assertGreater(len([p for p in inst_packets if p.op.name.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
expected: dict[str, list[int]] = {} # override in subclasses
def test_packet_counts(self):
if not self.expected: self.skipTest("no expected packet counts for this target")
for name, (events, *_) in self.examples.items():
with self.subTest(example=name):
if not self.expected.get(name): continue
counts = [len(list(decode(e.blob))) for e in events]
self.assertEqual(counts, self.expected[name], f"packet count mismatch in {name}")
@needs_rocprof
def test_rocprof_wave_times_match(self):
"""Wave start/end times must match rocprof exactly."""
for name, (events, lib, base) in self.examples.items():
with self.subTest(example=name):
occupancy, _ = run_rocprof_decoder([e.blob for e in events], lib, base, self.target)
# extract from rocprof occupancy records
roc_starts: dict[tuple[int, int, int], int] = {}
roc_waves: list[tuple[int, int]] = []
for wave_id, simd, cu, time, is_start in occupancy:
key = (wave_id, simd, cu)
if is_start: roc_starts[key] = time
elif key in roc_starts: roc_waves.append((roc_starts.pop(key), time))
# extract from our decoder
our_waves: list[tuple[int, int]] = []
for event in events:
wave_starts: dict[tuple[int, int, int], int] = {}
first_timestamp:int|None = None
for p in decode(event.blob):
if first_timestamp is None: first_timestamp = p._time
if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
elif isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
our_waves.append((wave_starts[key], p._time))
for st in wave_starts.values():
self.assertGreater(st, first_timestamp, "wave start must be after the first packet")
# rocprof fails non deterministically and gives inaccurate timestamps.
#self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
for st, et in our_waves:
self.assertGreater(et, st, "wave end must be after start")
@needs_rocprof
def test_rocprof_inst_times_match(self):
"""Instruction times must match rocprof exactly (excluding s_endpgm)."""
for name, (events, lib, base) in self.examples.items():
with self.subTest(example=name):
_, wave_insts = run_rocprof_decoder([e.blob for e in events], lib, base, self.target)
# skip last inst per wave (s_endpgm) - it needs special handling (time + duration instead of time + stall)
roc_insts = [time + stall for insts in wave_insts for time, stall in insts[:-1]]
# extract from our decoder
our_insts: list[int] = []
for event in events:
for p in decode(event.blob):
# INST ops for non-traced SIMDs (excluded from instruction count)
if isinstance(p, (INST, INST_RDNA4)) and not p.op.name.startswith("OTHER_"): our_insts.append(p._time)
elif isinstance(p, VALUINST): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
elif isinstance(p, IMMEDIATE_MASK):
for _ in range(bin(p.mask).count('1')): our_insts.append(p._time)
self.assertEqual(sorted(our_insts), sorted(roc_insts), f"instruction times mismatch in {name}")
class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
target = "gfx1100"
expected = {
"profile_empty_run_0": [1880, 1867, 1920, 1971, 1998, 1904],
"profile_empty_run_1": [1880, 1867, 1920, 1971, 1998, 1904],
"profile_gemm_run_0": [3275, 3278, 2426, 2475, 2511, 2431],
"profile_gemm_run_1": [3264, 3268, 2420, 2469, 2504, 2401],
"profile_ops_run_0": [1944, 4903, 1984, 2035, 2062, 1968],
"profile_ops_run_1": [1944, 4918, 1984, 2035, 2062, 1968],
"profile_plus_run_0": [1938, 1932, 1978, 2029, 2056, 1962],
"profile_plus_run_1": [1891, 1874, 1931, 1982, 2009, 1915],
}
class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
target = "gfx950"
def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
def test_rocprof_inst_times_match(self): self.skipTest("TODO: requires timestamp patching")
if __name__ == "__main__":
unittest.main()

99
tinygrad_repo/test/amd/test_sqtt_profiler.py Normal file
View File

@@ -0,0 +1,99 @@
import unittest, contextlib
from tinygrad import Device, Tensor, Context, TinyJit
from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.viz.serve import load_amd_counters, VizData
@contextlib.contextmanager
def save_sqtt():
data = VizData()
yield data.ctxs
Device[Device.DEFAULT].synchronize()
Device[Device.DEFAULT]._at_profile_finalize()
load_amd_counters(data, Compiled.profile_events)
data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]
@unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
class TestSQTTProfiler(unittest.TestCase):
# TODO: can we enable SQTT profiling in context?
@classmethod
def setUpClass(cls):
if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")
def setUp(self):
Device[Device.DEFAULT].synchronize()
Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
def test_simple(self):
t = Tensor.empty(1) + 1
with save_sqtt() as sqtt:
linear = t.schedule_linear()
run_linear(linear)
fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
self.assertEqual(len(sqtt), 1)
self.assertEqual(sqtt[0]["name"], f"SQTT {fn_name}")
def test_multiple_runs(self):
t = Tensor.empty(1) + 1
with save_sqtt() as sqtt:
linear = t.schedule_linear()
for _ in range(N:=3): run_linear(linear)
fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
self.assertEqual(len(sqtt), N)
for i in range(1, N):
self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name} n{i+1}")
def test_multiple_kernels(self):
t = ((Tensor.empty(1) + 1).contiguous() + 2)
linear = t.schedule_linear()
with save_sqtt() as sqtt:
run_linear(linear)
self.assertEqual(len(sqtt), len(linear.src))
for i,call in enumerate(linear.src):
fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}")
def test_multiple_kernels_lower(self):
t = ((Tensor.empty(1) + 1).contiguous() + 2)
linear = t.schedule_linear()
with save_sqtt() as sqtt:
run_linear(linear)
self.assertEqual(len(sqtt), len(linear.src))
for i,call in enumerate(linear.src):
fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}")
def test_jit(self):
@TinyJit
def f(a): return a + 1
t = Tensor.empty(1)
with save_sqtt() as sqtt:
for _ in range(N:=5):
f(t).realize()
self.assertEqual(len(sqtt), N)
kernel_name = sqtt[0]["name"]
for i,s in enumerate(sqtt[1:], start=1): self.assertEqual(s["name"], f"{kernel_name} n{i+1}")
# TODO: can we trace SQTT for graphed kernels?
def test_jit_graph(self, kernel_count=3*1):
@TinyJit
def f(a): return ((a + 1).contiguous() + 2).contiguous().sum()
t = Tensor.empty(32)
with save_sqtt() as sqtt:
for _ in range(5):
f(t).realize()
names = [s["name"] for s in sqtt]
k0, k1, k2 = names[:3]
for i in range(3, len(sqtt), 3):
n = (i // 3)+1
self.assertEqual(names[i], f"{k0} n{n}")
self.assertEqual(names[i+1], f"{k1} n{n}")
self.assertEqual(names[i+2], f"{k2} n{n}")
self.assertEqual(len(sqtt), kernel_count)
@Context(JIT=2)
def test_jit_multiple_kernels(self): self.test_jit_graph(kernel_count=3*5)
if __name__ == "__main__":
unittest.main()

174
tinygrad_repo/test/amd/test_sqttmap.py Normal file
View File

@@ -0,0 +1,174 @@
# test to compare every packet with the rocprof decoder
import unittest, pickle, functools, json
from typing import Iterator
from pathlib import Path
from tinygrad.helpers import DEBUG, getenv, temp, ansistrip, Context
from tinygrad.renderer.amd.sqtt import print_packets, map_insts
from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
from tinygrad.viz.serve import sqtt_timeline, amd_decode
from test.amd.disasm import disasm
from test.null.test_viz import run_cli
import tinygrad
EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
def needs_rocprof(fn):
@functools.wraps(fn)
def wrapper(self, *args, **kwargs):
# check if latest rocprof is available, if not, skip rocprof comparison tests
# rocprof doesn't have a version string, decode a known pickle to validate it's the latest
try:
from extra.sqtt.roc import decode as roc_decode
with open(EXAMPLES_DIR/"gfx1200"/"profile_plus_run_0.pkl", "rb") as f:
data = pickle.load(f)
sqtt = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"][1]
kern = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}[sqtt.kern]
rctx = roc_decode([sqtt], {kern.tag:{addr+kern.base:inst for addr,inst in amd_decode(kern.lib, "gfx1200").items()}})
insts = [e.time for e in list(rctx.inst_execs.values())[0][0].unpack_insts()]
self.assertListEqual(insts, [28178, 28179, 28180, 28181, 28182, 29882, 29883, 29884, 29885, 30966, 30983, 30985, 30992, 30993])
except Exception as e: self.skipTest(f"latest rocprof not available, install with extra/sqtt/install_rocprof_decoder.py: {e}")
return fn(self, *args, **kwargs)
return wrapper
def rocprof_inst_traces_match(sqtt, prg, target):
from extra.sqtt.roc import decode as roc_decode, InstExec
addr_table = amd_decode(prg.lib, target)
disasm_map = {addr+prg.base:inst for addr,inst in addr_table.items()}
rctx = roc_decode([sqtt], {prg.tag:disasm_map})
rwaves = rctx.inst_execs.get((sqtt.kern, sqtt.exec_tag), [])
rwaves_iter:dict[int, list[Iterator[InstExec]]] = {} # wave unit (0-15) -> list of inst trace iterators for all executions on that unit
for w in rwaves: rwaves_iter.setdefault(w.wave_id, []).append(w.unpack_insts())
if not rwaves: return 0, 0, 0
passed_insts = 0
for pkt, info in map_insts(sqtt.blob, prg.lib, target):
if DEBUG >= 2: print_packets([(pkt, info)])
if info is None: continue
if DEBUG >= 2: print(f"{' '*29}{disasm(info.inst)}")
rocprof_inst = next(rwaves_iter[info.wave][0])
ref_pc = rocprof_inst.pc-prg.base
# always check pc matches
assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
# special handling for s_endpgm, it marks the wave completion.
if info.inst == s_endpgm():
completed_wave = list(rwaves_iter[info.wave].pop(0))
assert len(completed_wave) == 0, f"incomplete instructions in wave {info.wave}"
# otherwise the packet timestamp is time + "stall"
else:
assert pkt._time == rocprof_inst.time+rocprof_inst.stall
passed_insts += 1
for k,v in rwaves_iter.items():
assert len(v) == 0, f"incomplete wave {k}"
return passed_insts, len(rwaves), len(rwaves_iter)
class TestSQTTMapBase(unittest.TestCase):
target: str
examples: dict
@classmethod
def setUpClass(cls):
if cls is TestSQTTMapBase: raise unittest.SkipTest("base class")
cls.examples = {}
for pkl_path in ([Path(temp("profile.pkl", append_user=True))] if getenv("LOAD_PROFILE") else sorted((EXAMPLES_DIR/cls.target).glob("*.pkl"))):
with open(pkl_path, "rb") as f:
data = pickle.load(f)
sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
kern_events = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}
if sqtt_events and kern_events:
cls.examples[pkl_path.stem] = (sqtt_events, kern_events, cls.target)
@needs_rocprof
def test_rocprof_inst_traces_match(self):
for name, (events, kern_events, target) in self.examples.items():
if "sync" in name and self.target.startswith("gfx12"):
self.skipTest("our timestamps are off by a few cycles because rocprof patches timestamps for rdna4 barriers")
for event in events:
if not event.itrace: continue
if event.kern not in kern_events: continue
with self.subTest(example=name, kern=event.kern):
passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")
def test_sqtt_timeline(self):
for name, (events, kern_events, target) in self.examples.items():
for event in events:
if (p:=kern_events.get(event.kern)) is None: continue
with self.subTest(example=name, kern=event.kern):
# skip if there's no SQTT frequency data
if not (timeline:=list(sqtt_timeline(event.blob, p.lib, target))): continue
if not (frequency:=[e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]): continue
mean = sum(frequency) / len(frequency)
variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
self.assertGreater(mean, 0)
if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
insts, execs = 0, 0
for e in events:
if "EXEC" in e.device:
if "ALT" not in e.name.display_name: execs += 1
elif "WAVE" in e.device:
# sopk/immediates don't get ALU/MEM EXEC
if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
"WAVEEND", "WAVEEND_RDNA4", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
self.assertEqual(execs, insts)
def test_wave_sync(self):
for name, (events, kern_events, target) in self.examples.items():
for event in events:
wave_barriers = {}
for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target):
if type(e).__name__ == "ProfileRangeEvent" and e.name.display_name == "BARRIER": wave_barriers.setdefault(e.device, []).append(e)
if not wave_barriers: continue
for row, events in wave_barriers.items():
for e in events:
assert e.en-e.st > 1, f"all barriers must have a duration greater than 1, got {e}"
def test_sqtt_cli(self):
for pkl_path in sorted((EXAMPLES_DIR/self.target).glob("*.pkl")):
out = run_cli("--profile-path", str(pkl_path), "--ls")
sqtt_traces = [l["value"].strip() for l in out if "SQTT" in l["value"]]
for name in sqtt_traces:
lines = run_cli("--profile-path", str(pkl_path), "-s", ansistrip(name))
self.assertIn("Clk", lines[0]["value"])
waves = [r["clk"] for r in lines[2:] if "WAVE" in r["unit"]]
self.assertEqual(waves, sorted(waves), f"wave timestamps not monotonic in {name}")
with Context(DEBUG=2):
kernels = run_cli("--profile-path", str(pkl_path), "-s", "AMD")
self.assertEqual(len(kernels), len(self.examples[pkl_path.stem][1]))
class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
class TestSQTTMapRDNA4(TestSQTTMapBase):
target = "gfx1200"
@unittest.expectedFailure
def test_pipes(self):
events, kernels, target = self.examples["profile_handwritten_run_0"]
lib = list(kernels.values())[0].lib
dispatch_st:dict[str, int] = {}
row_ends:dict[str, int] = {}
row_counts:dict[str, int] = {}
for e in sqtt_timeline(events[1].blob, lib, target):
if type(e).__name__ != "ProfileRangeEvent": continue
info = json.loads(e.name.ret) if e.name.ret else {}
if e.device.startswith("WAVE"):
idx = row_counts.get(e.device, 0)
dispatch_st[f"{e.device}-{idx}"] = int(e.st)
row_counts[e.device] = idx + 1
elif info.startswith("LINK:"):
delay = int(e.st) - dispatch_st[info[len("LINK:"):]]
self.assertGreaterEqual(delay, 1, f"EXEC {e.device} starts before DISPATCH: delay={delay}")
if (prev_en:=row_ends.get(e.device)) is not None:
self.assertGreaterEqual(e.st, prev_en, f"EXEC overlap in {e.device}: {e.st} < prev end {prev_en}")
row_ends[e.device] = int(e.en)
class TestSQTTMapCDNA(TestSQTTMapBase):
target = "gfx950"
def test_rocprof_inst_traces_match(self): self.skipTest("requires timestamp patching to match rocprof, currently it's off by a few cycles")
if __name__ == "__main__":
unittest.main()

0
tinygrad_repo/test/backend/__init__.py Normal file
View File

246
tinygrad_repo/test/backend/test_arange.py Normal file
View File

@@ -0,0 +1,246 @@
import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
from tinygrad.helpers import Context, getenv, DEV
from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import needs_second_gpu
class TestArange(unittest.TestCase):
def _get_flops(self, tensor, desired):
GlobalCounters.reset()
linear = compile_linear(tensor.schedule_linear())
self.assertEqual(len(linear.src), 1)
run_linear(linear)
np.testing.assert_equal(tensor.numpy(), desired)
return estimate_uop(linear.src[-1]).ops
def test_arange_complexity(self):
self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)
@unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
def test_arange_cumsum(self):
np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())
def test_arange_cat(self):
t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])
def test_eye_complexity(self):
with Context(NOOPT=1):
# NOTE: not every backend supports CMPEQ
self.assertLessEqual(self._get_flops(Tensor.eye(2560).contiguous(), np.eye(2560)), 2*2560*2560)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexing is weird")
def test_tri_complexity(self):
with Context(NOOPT=1):
t = Tensor.ones(256, 256).contiguous().realize()
linear = compile_linear(t.triu().schedule_linear())
self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
DSET, DDIM = 2048, 32
class TestIndexing(unittest.TestCase):
def test_arange_2_reduce(self):
needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
needle[1337] = 1
needle.realize()
with Context(NOOPT=1):
GlobalCounters.reset()
out = ((Tensor.arange(1,16385)-1)*needle).sum()
linear, var_vals = out.linear_with_vars()
self.assertEqual(len(linear.src), 1)
run_linear(linear, var_vals)
self.assertEqual(out.item(), 1337)
def test_manual_index(self):
dataset = Tensor.rand(DSET, DDIM).realize()
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=1):
GlobalCounters.reset()
rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1)
idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1)
reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False))
X = full.sum(axis=(2,3))
linear, var_vals = X.linear_with_vars()
self.assertEqual(len(linear.src), 1)
run_linear(linear, var_vals)
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
np.testing.assert_allclose(real_index, X.numpy())
def test_index_variable(self):
dataset = Tensor.rand(DSET, DDIM).realize()
v = Variable("v", 0, DDIM-1)
with Context(NOOPT=1):
GlobalCounters.reset()
vb = Tensor(v.bind(12))
comp = dataset[vb].numpy()
# no global ops because they are all indexing
self.assertEqual(GlobalCounters.global_ops, 0)
np.testing.assert_allclose(comp, dataset.numpy()[12])
def test_index(self):
dataset = Tensor.rand(DSET, DDIM).realize()
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=1):
GlobalCounters.reset()
X = dataset[idxs]
assert X.shape == (4,DDIM)
linear, var_vals = X.linear_with_vars()
self.assertEqual(len(linear.src), 1)
run_linear(linear, var_vals)
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
np.testing.assert_allclose(real_index, X.numpy())
def test_index_fused(self, noopt=1):
dataset = Tensor.rand(DSET, DDIM).realize()
idxs = Tensor([0,3,5,6]).realize()
real_index = dataset.numpy()[idxs.numpy()]
print("*** indexing ***")
with Context(NOOPT=noopt):
GlobalCounters.reset()
X = dataset[idxs]
assert X.shape == (4,DDIM)
linear, var_vals = X.linear_with_vars()
self.assertEqual(len(linear.src), 1)
run_linear(linear, var_vals)
assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
np.testing.assert_allclose(real_index, X.numpy())
@unittest.skip("not ready")
def test_index_fused_opt(self): self.test_index_fused(0)
def test_index_fused_out_of_bounds(self):
dataset = Tensor.rand(256, 256).realize()
idxs = Tensor([-19238, -257, 256, 495, 10982377]).realize()
with Context(NOOPT=1):
X = dataset[idxs]
np.testing.assert_equal(X.numpy(), 0)
def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
# WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
if Device.DEFAULT == "WEBGPU": op_limit *= 15
# from tinygrad.nn.datasets import mnist
X_train, Y_train = Tensor.randint(DSET, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(DSET, dtype='uchar').realize()
with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop):
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
GlobalCounters.reset()
x = X_train[samples].numpy()
y = Y_train[samples].numpy()
assert GlobalCounters.global_ops < op_limit, f"too many ops {GlobalCounters.global_ops} != {op_limit}"
np.testing.assert_allclose(X_train.numpy()[samples.numpy()], x)
np.testing.assert_allclose(Y_train.numpy()[samples.numpy()], y)
def test_index_mnist_opt(self): self.test_index_mnist(0)
def test_index_mnist_split(self): self.test_index_mnist(1, split_reduceop=1)
def test_index_mnist_opt_split(self): self.test_index_mnist(0, split_reduceop=1)
def test_llama_embedding(self, noopt=1, op_limit=65536):
# llama3 is 128256
vocab_size, embed_size = (10, 3)
emb = nn.Embedding(vocab_size, embed_size)
emb_w = emb.weight.numpy()
x = Tensor([1,2,3,4])
with Context(NOOPT=noopt):
GlobalCounters.reset()
z = emb(x).realize()
self.assertLessEqual(GlobalCounters.global_ops, op_limit)
self.assertEqual(GlobalCounters.kernel_count, 2)
if getenv("CHECK", 1):
import torch
with torch.no_grad():
torch_emb = torch.nn.Embedding(vocab_size, embed_size).eval()
torch_emb.weight[:] = torch.tensor(emb_w, dtype=torch.float32)
torch_z = torch_emb(torch.tensor(x.numpy()))
# TODO: reshape to match torch, should we do this in nn?
np.testing.assert_allclose(z.numpy().reshape(4, embed_size), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# at least the arange is being fused
def test_llama_embedding_opt(self): self.test_llama_embedding(0, 1_736_704_000)
# NOTE: call doesn't work with SPEC=2
@unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
@Context(USE_ATOMICS=1, SPEC=1)
def test_llama_8b_embedding_backward(self):
from tinygrad.renderer.cstyle import CStyleLanguage
if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
vocab_size, embed_size = 1000, 128
bs, seqlen = 4, 256
idx = Tensor.randint(bs, seqlen, high=vocab_size)
emb = nn.Embedding(vocab_size, embed_size)
emb.weight = Tensor.ones(vocab_size, embed_size)
gt = Tensor.zeros(bs, seqlen, embed_size)
Tensor.realize(idx, emb.weight, gt)
GlobalCounters.reset()
loss = (emb(idx)-gt).square().sum()
loss.backward()
emb.weight.grad.realize()
bwd_ops = GlobalCounters.global_ops
print(f"embedding bwd: {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
self.assertLess(bwd_ops, bs*seqlen*embed_size*20, f"backward ops {bwd_ops:,} should be less than 20 per with atomic scatter-add")
# correctness check
expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
for i in idx.flatten().numpy(): expected_grad[i] += 2
np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
@needs_second_gpu
@unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
@Context(USE_ATOMICS=1, SPEC=1)
def test_embedding_backward_vocab_sharded(self):
from tinygrad.renderer.cstyle import CStyleLanguage
if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
vocab_size, embed_size = 1000, 128
bs, seqlen = 4, 256
idx = Tensor.randint(bs, seqlen, high=vocab_size)
emb = nn.Embedding(vocab_size, embed_size)
emb.weight = Tensor.ones(vocab_size, embed_size)
gt = Tensor.zeros(bs, seqlen, embed_size)
Tensor.realize(idx, emb.weight, gt)
# compute expected grad on single device
expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
for i in idx.flatten().numpy(): expected_grad[i] += 2
# now shard the embedding weight on vocab axis and recompute
emb.weight = Tensor.ones(vocab_size, embed_size)
emb.weight.shard_(devices, axis=0)
idx = idx.shard(devices, axis=None)
gt = gt.shard(devices, axis=None)
Tensor.realize(idx, emb.weight, gt)
loss = (emb(idx)-gt).square().sum()
loss.backward()
np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
@unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and DEV.arch.startswith("gfx")), "tests AMD bf16 cast overhead")
def base_test_llama_8b_rope_backward(self, dtype, ops_scale=1):
from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
bs, seqlen, dim, n_heads = 1, 512, 256, 4
head_dim = dim // n_heads
x = Tensor.randn(bs, seqlen, dim, dtype=dtype)
wq = Tensor.randn(dim, dim, dtype=dtype)
freqs_cis = precompute_freqs_cis(head_dim, seqlen).cast(dtype)
Tensor.realize(x, wq, freqs_cis)
xq = (x @ wq.T)
# main llama does not fuse it
#xq = xq.contiguous_backward()
xq = xq.reshape(bs, seqlen, n_heads, head_dim)
xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
xq_rope.sum().backward()
linear = compile_linear(wq.grad.schedule_linear())
assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
bwd_ops = estimate_uop(linear.src[0]).ops
expected_ops = bs*seqlen*dim*dim*ops_scale
print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
self.assertLess(bwd_ops, expected_ops, f"rope bwd ops {bwd_ops:,} should be < {ops_scale} per (got {bwd_ops/(bs*seqlen*dim*dim):.1f})")
def test_llama_8b_rope_backward_f16(self):
self.base_test_llama_8b_rope_backward(dtypes.float16, ops_scale=2)
# bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
def test_llama_8b_rope_backward_bf16(self):
self.base_test_llama_8b_rope_backward(dtypes.bfloat16, ops_scale=2 if Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950") else 25)
if __name__ == "__main__":
unittest.main()

252
tinygrad_repo/test/backend/test_asm_gemm.py Normal file
View File

@@ -0,0 +1,252 @@
import unittest
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.helpers import getenv, system, DEV
from extra.gemm.cdna_asm_gemm import asm_gemm
from test.helpers import needs_second_gpu
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
# On non CDNA4 it will only validate the Tensor.custom_kernel integration
# Use DEV=NULL:HIP:gfx950 to also test the assembly
def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")
def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
Tensor.manual_seed(0)
input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
with Context(DEBUG=0):
Tensor.realize(a_rand, b_rand)
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None
if dtype == FP8_DTYPE:
a_rand, x_scale, _ = quantize_fp8(a_rand)
b_rand, w_scale, _ = quantize_fp8(b_rand)
grad_amax_state = Tensor.full((), FP8_MAX, dtype=dtypes.float32, device=devs).contiguous()
with Context(DEBUG=0):
Tensor.realize(a_rand, x_scale, b_rand, w_scale, grad_amax_state)
# clone all inputs before any backward: a clone copies the source's current .grad
a, b = a_rand.clone(), b_rand.clone()
if dtype == FP8_DTYPE:
a_ref, b_ref = a_rand.detach().cast(dtypes.bfloat16), b_rand.detach().cast(dtypes.bfloat16)
else:
a_ref, b_ref = a_rand.clone(), b_rand.clone()
if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
if dtype == FP8_DTYPE:
tst = asm_gemm(a, b, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state)
else:
tst = asm_gemm(a, b)
tst.sum().backward()
Tensor.realize(tst, a.grad, b.grad)
if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
if dtype == FP8_DTYPE:
ref = ((a_ref @ b_ref) * x_scale * w_scale).cast(dtypes.bfloat16)
else:
ref = a_ref @ b_ref
ref.sum().backward()
Tensor.realize(ref, a_ref.grad, b_ref.grad)
# no validation on the NULL device
if a_rand.device.startswith("NULL"): return None
atol, rtol = (2e-1, 1e-2) if dtype == dtypes.bfloat16 else (256, 1e-2) if dtype == FP8_DTYPE else (1e-2, 1e-3)
# allow more rtol for multi because of ALLREDUCE_CAST
grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, 2e-2 if multi else rtol)
with Context(DEBUG=0):
# enable for debugging, slow for larger gemms
if getenv("USE_NPY"):
import numpy as np
np.testing.assert_allclose(tst.numpy(), ref.numpy(), atol=atol, rtol=rtol)
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
assert tst.allclose(ref, atol=atol, rtol=rtol).item(), "forward mismatch"
assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"
def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)
def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
# 128x smaller than usual
# uses the UOp GEMM, runs on non CDNA4 and CI
@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
class TestGemm(unittest.TestCase):
def setUp(self):
if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.half)
def test_gemm(self): verify_asm_gemm(1, 64, 32, 112)
def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
@needs_second_gpu
def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
@needs_second_gpu
def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
@needs_second_gpu
def test_gemm_m_sharded(self): verify_asm_gemm_m_sharded(2*64, 64, 32, gpus=2)
@needs_second_gpu
def test_gemm_n_sharded(self): verify_asm_gemm_n_sharded(1, 64, 64, 32, gpus=2)
@needs_second_gpu
def test_gemm_n_sharded_2d(self): verify_asm_gemm_n_sharded_2d(64, 2*64, 32, gpus=2)
@needs_second_gpu
def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)
# uses the smallest size for the cdna assembly gemm
class TestAsmGEMM(unittest.TestCase):
def setUp(self):
if not is_cdna4():
self.skipTest("assembly gemm is only for cdna4")
def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
def test_verify_with_numpy(self):
import numpy as np
M, N, K = 256, 256, 64
rng = np.random.default_rng(0)
a_np = (rng.random((M, K), dtype=np.float32) - 0.5).astype(np.half)
b_np = (rng.random((K, N), dtype=np.float32) - 0.5).astype(np.half)
c_np = a_np @ b_np
a, b = Tensor(a_np), Tensor(b_np)
c = asm_gemm(a, b)
c.realize()
# no validation on the NULL device
if a.device.startswith("NULL"): return None
np.testing.assert_allclose(c.numpy(), c_np, atol=2e-3, rtol=5e-2)
def test_unsupported_batch(self):
with self.assertRaisesRegex(AssertionError, "batch size"):
verify_asm_gemm(3, 256, 256, 256)
def test_unsupported_k(self):
with self.assertRaisesRegex(AssertionError, "not a multiple"):
verify_asm_gemm(1, 1024, 1024, 100)
def test_unsupported_m(self):
with self.assertRaisesRegex(AssertionError, "not a multiple"):
verify_asm_gemm(1, 1000, 256, 256)
def test_unsupported_n(self):
with self.assertRaisesRegex(AssertionError, "not a multiple"):
verify_asm_gemm(1, 256, 1000, 256)
# test the Asm GEMM with Llama shapes, only run on the real machine for speed
class TestGemmLlama(unittest.TestCase):
dtype = dtypes.bfloat16
def setUp(self):
if not is_cdna4() or DEV.interface.startswith("MOCK"):
self.skipTest("very slow on non mi350x")
def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
def test_empty_bw(self):
x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype)
y = Tensor.empty((N, N), dtype=self.dtype)
if self.dtype == FP8_DTYPE:
x_scale = Tensor.empty((), dtype=dtypes.float32)
w_scale = Tensor.empty((), dtype=dtypes.float32)
grad_amax_state = Tensor.empty((), dtype=dtypes.float32).contiguous()
z = asm_gemm(x, y, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state)
else:
z = asm_gemm(x, y)
z.sum().backward()
Tensor.realize(z, x.grad, y.grad)
# FP8 GEMM stores bf16 output and its backward produces bf16 gradients.
grad_dtype = dtypes.bfloat16 if self.dtype == FP8_DTYPE else self.dtype
assert z.dtype == dtypes.bfloat16
assert x.grad.dtype == y.grad.dtype == grad_dtype
def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=self.dtype)
def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336, dtype=self.dtype)
def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096, dtype=self.dtype)
def test_gemm1(self): verify_asm_gemm(8, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
@unittest.skip("disabled, asm in this shape is slower than tinygrad")
def test_gemm2(self): verify_asm_gemm(8, 8192, 128256, 4096, dtype=self.dtype, gpus=8)
def test_gemm3(self): verify_asm_gemm(8, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, dtype=self.dtype, gpus=8)
def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, dtype=self.dtype, gpus=8)
@unittest.skip("disabled, asm in this shape is slower than tinygrad")
def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
def test_gemm8(self): verify_asm_gemm(1, 4096, 14336, 8192, dtype=self.dtype)
def test_gemm9(self): verify_asm_gemm(8, 4096, 14336, 8192, dtype=self.dtype, gpus=8)
def test_gemm10(self): verify_asm_gemm(1, 4096, 8192, 4096, dtype=self.dtype)
def test_gemm_previously_unsupported(self): verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)
def test_k_sharded_1(self): verify_asm_gemm_k_sharded(14336, 4096, 8*8192, dtype=self.dtype, gpus=8)
def test_k_sharded_2(self): verify_asm_gemm_k_sharded(4096, 14336, 8*8192, dtype=self.dtype, gpus=8)
def test_k_sharded_3(self): verify_asm_gemm_k_sharded(4096, 4096, 8*8192, dtype=self.dtype, gpus=8)
# M-sharded 2D
def test_m_sharded_1(self): verify_asm_gemm_m_sharded(8*8192, 4096, 4096, dtype=self.dtype, gpus=8)
def test_m_sharded_2(self): verify_asm_gemm_m_sharded(8*4096, 14336, 4096, dtype=self.dtype, gpus=8)
# N-sharded 2D
def test_n_sharded_2d_1(self): verify_asm_gemm_n_sharded_2d(8192, 8*4096, 4096, dtype=self.dtype, gpus=8)
def test_n_sharded_2d_2(self): verify_asm_gemm_n_sharded_2d(4096, 8*14336, 4096, dtype=self.dtype, gpus=8)
# tensor parallel shapes (Llama 8B, MP=8)
def test_tp_n_sharded_wq(self): verify_asm_gemm_n_sharded(1, 8192, 4096, 4096, dtype=self.dtype, gpus=8)
def test_tp_n_sharded_w1(self): verify_asm_gemm_n_sharded(1, 8192, 14336, 4096, dtype=self.dtype, gpus=8)
def test_tp_k_sharded_wo(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 4096, dtype=self.dtype, gpus=8)
def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=self.dtype, gpus=8)
# more shapes: vary M, N, K independently
def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256)
def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256)
def test_shape_small_rect_n(self): verify_asm_gemm(1, 256, 512, 256)
def test_shape_small_rect_k(self): verify_asm_gemm(1, 256, 256, 512)
def test_shape_tall(self): verify_asm_gemm(1, 2048, 256, 256)
def test_shape_wide(self): verify_asm_gemm(1, 256, 2048, 256)
def test_shape_deep(self): verify_asm_gemm(1, 256, 256, 4096)
def test_shape_non_square(self): verify_asm_gemm(1, 1024, 2048, 512)
def test_shape_batched_small(self): verify_asm_gemm(2, 256, 256, 256)
def test_shape_batched_rect(self): verify_asm_gemm(2, 512, 1024, 256)
# K edge cases: iters=1,2,3 exercise different loop paths
def test_shape_k64(self): verify_asm_gemm(1, 256, 256, 64)
def test_shape_k128(self): verify_asm_gemm(1, 256, 256, 128)
def test_shape_k192(self): verify_asm_gemm(1, 256, 256, 192)
def test_llama3_out1(self): verify_asm_gemm(1, 8192, 128256, 4096, dtype=self.dtype)
def test_llama3_out2(self): verify_asm_gemm(1, 8192, 4096, 128256, dtype=self.dtype)
def test_llama3_out3(self): verify_asm_gemm(1, 4096, 128256, 8192, dtype=self.dtype)
def has_hipcc():
try: system("hipcc --version")
except Exception: return False
return True
@unittest.skipUnless(has_hipcc(), "FP8 gemm requires hipcc to compile")
class TestGemmLlamaFP8(TestGemmLlama): dtype = FP8_DTYPE
class TestMagicGu(unittest.TestCase):
def test_magicgu_matches_old(self):
from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
old_gemm_shapes = [
(8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
(8192, 8192, 8192), (4096, 4096, 4096), (4096, 14336, 4096),
(4096, 14336, 8192), (4096, 4096, 14336), (14336, 4096, 8192),
(4096, 8192, 14336), (4096, 4096, 8192), (4096, 8192, 4096),
]
for M, N, K in old_gemm_shapes:
iters = K // TILE_K
total = (M // TILE_M) * (N // TILE_N) * iters
for batch in [1, 2]:
magic, shift = _magicgu_mulhi(iters, total * batch)
old_magic, old_shift = old_iters_args[iters]
self.assertEqual((magic, shift), (old_magic, old_shift), f"mismatch for ({M},{N},{K}) batch={batch} iters={iters}")
if __name__ == "__main__":
unittest.main()

209
tinygrad_repo/test/backend/test_const_folding.py Normal file
View File

@@ -0,0 +1,209 @@
import unittest, math
from tinygrad import Tensor, Device, dtypes
from tinygrad.dtype import DTYPES_DICT
from tinygrad.uop.ops import Ops, UOp
import numpy as np
from test.helpers import not_support_multi_device
def _check_ast_count(desired_count:int, t:Tensor):
# NOTE: this has side effect because everything can be scheduled only once
schedule = t.schedule_linear()
asts = [s for s in schedule.src if s.src[0].op is Ops.SINK]
len(asts)
# NOT SUPPORTED ANYMORE
#assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
class TestMovedConstFolding(unittest.TestCase):
def test_contiguous_deviceless_const(self):
t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous()
self.assertIs(t.uop.op, Ops.CONST)
self.assertIsNone(t.uop.device)
def test_add_shrunk_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
def test_add_padded_zero(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
def test_mul_shrunk_one(self):
_check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
def test_add_padded_one(self):
_check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
def test_copy_padded_const(self):
schedule = Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").schedule_linear()
assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded"
np.testing.assert_equal(Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
def test_cast_padded(self):
# NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
_check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
# folded
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
class TestReduceOpsConstFolding(unittest.TestCase):
def test_const_sum(self):
_check_ast_count(0, Tensor.ones(4, 5, 6).sum())
np.testing.assert_equal(Tensor.ones(4, 5, 6).sum().numpy(), 4 * 5 * 6)
_check_ast_count(0, Tensor.ones(4, 5, 6).sum(axis=0))
np.testing.assert_equal(Tensor.ones(4, 5, 6).sum(axis=0).numpy(), np.full((5, 6), 4))
_check_ast_count(0, Tensor(4).sum())
np.testing.assert_equal(Tensor(4).sum().numpy(), 4)
def test_padded_const_sum(self):
_check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).sum())
np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)
# NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
_check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).exp().sum())
np.testing.assert_allclose(Tensor.ones(4).pad(((1, 1),)).exp().sum().numpy(), 4 * math.e + 2)
def test_bool_zero_max(self):
_check_ast_count(0, Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)))
np.testing.assert_equal(Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)).numpy(), False)
def test_zero_size_ops(self):
for reduceop in [lambda x:x.prod(), lambda x:x.sum()]: # lambda x:x.max() NOTE: numpy gives "reduction operation maximum which has no identity"
_check_ast_count(0, reduceop(Tensor.empty(1, 0)))
np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty(shape)))
def test_zero_size_ops_view(self):
for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
_check_ast_count(0, reduceop(Tensor.empty(1, 0, 4).permute((1, 2, 0)).contiguous()))
np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty((shape))))
def test_zero_size_ops_realized(self):
for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
_check_ast_count(0, reduceop((Tensor.randn(0, 1)+1).realize()))
np.testing.assert_equal(reduceop((Tensor.randn(shape:=(0, 1))+1).realize()).numpy(), reduceop(np.empty(shape)))
def test_zero_size_realize_folded(self):
# non contiguous folded output doesn't realize
_check_ast_count(0, Tensor.empty(1, 0).sum())
# contiguous folded const can still schedule
a = Tensor.empty(1, 0).sum().contiguous()
_check_ast_count(2, a+2)
self.assertIs(a.uop.base.op, Ops.BUFFER)
np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)
# otherwise we just fuse it
_check_ast_count(1, (Tensor.empty(1, 0).sum()+2).contiguous())
np.testing.assert_equal((Tensor.empty(1, 0).sum()+2).numpy(), 2)
def test_const_prod(self):
_check_ast_count(0, Tensor.full((2, 3), fill_value=2).prod())
np.testing.assert_equal(Tensor.full((2, 3), fill_value=2).prod().numpy(), 2**(2*3))
_check_ast_count(0, Tensor.full((4, 5, 6), fill_value=2).prod(axis=0))
np.testing.assert_equal(Tensor.full((4, 5, 6), fill_value=2).prod(axis=0).numpy(), np.full((5, 6), 2**4))
_check_ast_count(0, Tensor(4).prod())
np.testing.assert_equal(Tensor(4).prod().numpy(), 4)
def test_const_max(self):
_check_ast_count(0, Tensor.ones(4, 5, 6).max())
np.testing.assert_equal(Tensor.ones(4, 5, 6).max().numpy(), 1)
_check_ast_count(0, Tensor(4).max())
np.testing.assert_equal(Tensor(4).max().numpy(), 4)
def test_sum_output_dtype(self):
# sum output dtype can be different from input
for dt in DTYPES_DICT.values():
if dt in Device[Device.DEFAULT].renderer.supported_dtypes():
t = Tensor.ones(16, dtype=dt).reshape(4, 4)
assert t.sum().dtype == t.contiguous().sum().dtype
@unittest.skipIf(not_support_multi_device() or True, "no multi, RANGEIFY doesn't support multi const folding")
class TestMultiConstFolding(unittest.TestCase):
def test_multi_const_folding_literal(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
t = Tensor.arange(16).float().clone().to(ds).realize()
# non const folding case creates one ast on each shard
_check_ast_count(4, t + 1)
_check_ast_count(4, 1 + t)
_check_ast_count(4, t * 2)
_check_ast_count(4, 2 * t)
# const folded
_check_ast_count(0, t + 0)
_check_ast_count(0, 0 + t)
_check_ast_count(0, t * 0)
_check_ast_count(0, 0 * t)
_check_ast_count(0, t * 1)
_check_ast_count(0, 1 * t)
np.testing.assert_equal((t + 0).numpy(), np.arange(16))
np.testing.assert_equal((t * 0).numpy(), [0] * 16)
np.testing.assert_equal((t * 1).numpy(), np.arange(16))
_check_ast_count(0, t ** 0)
_check_ast_count(0, t ** 1)
_check_ast_count(0, 1 ** t)
def test_multi_const_folding_tensor(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
t = Tensor.arange(16).float().clone().to(ds).realize()
zero = Tensor.zeros(16).to(ds).realize()
one = Tensor.ones(16).to(ds).realize()
# const folded
_check_ast_count(0, t + zero)
_check_ast_count(0, zero + t)
_check_ast_count(0, t * zero)
_check_ast_count(0, zero * t)
_check_ast_count(0, t * one)
_check_ast_count(0, one * t)
np.testing.assert_equal((t + zero).numpy(), np.arange(16))
np.testing.assert_equal((t * zero).numpy(), [0] * 16)
np.testing.assert_equal((t * one).numpy(), np.arange(16))
_check_ast_count(0, t ** zero)
_check_ast_count(0, t ** one)
_check_ast_count(0, one ** t)
np.testing.assert_equal((t ** zero).numpy(), [1] * 16)
np.testing.assert_equal((t ** one).numpy(), np.arange(16))
np.testing.assert_equal((one ** t).numpy(), [1] * 16)
class TestThreefryConstFolding(unittest.TestCase):
def test_threefry(self):
x = UOp.const(dtypes.uint64, 5, Device.DEFAULT, ()).threefry(UOp.const(dtypes.uint64, 10, Device.DEFAULT, ()))
self.assertIs(x.simplify().op, Ops.CONST)
class TestTautologicalCompare(unittest.TestCase):
# without const folding, these would have triggered -Wtautological-compare in clang
def test_lt_false(self):
# bool < False is always false
np.testing.assert_equal((Tensor([True, False]) < False).numpy(), [False, False])
def test_true_lt(self):
# True < bool is always false
np.testing.assert_equal((True < Tensor([True, False])).numpy(), [False, False])
def test_truth_table(self):
np.testing.assert_equal((Tensor(False) < Tensor(False)).numpy(), False)
np.testing.assert_equal((Tensor(False) < Tensor(True)).numpy(), True)
np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
def test_a_eq_a(self):
# self eq is always true for int or bool
a = Tensor([1, 2, 3])
np.testing.assert_equal((a == a).numpy(), [True, True, True])
# not true for nan
a = Tensor([math.nan, 1.0, 2.0])
np.testing.assert_equal((a == a).numpy(), [False, True, True])
def test_a_ne_a(self):
# self not eq is always false for int or bool
a = Tensor([1, 2, 3])
np.testing.assert_equal((a != a).numpy(), [False, False, False])
# not true for nan
a = Tensor([math.nan, 1.0, 2.0])
np.testing.assert_equal((a != a).numpy(), [True, False, False])
if __name__ == '__main__':
unittest.main()

446
tinygrad_repo/test/backend/test_custom_kernel.py Normal file
View File

@@ -0,0 +1,446 @@
import unittest
from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
from tinygrad.dtype import AddrSpace, dtypes
from tinygrad.uop.ops import KernelInfo, AxisType, Ops
# **** kernels ****
def custom_arange_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
def custom_eye_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
j = UOp.range(C.shape[1], 1)
return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
A,B = A.flatten(), B.flatten()
assert B.numel() == A.numel()
i = UOp.range(A.numel(), 0)
return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.numel()}"))
def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
C,A,B = C.flatten(), A.flatten(), B.flatten()
i = UOp.range(C.numel(), 0)
return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.numel()}")).simplify()
def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
assert C.numel() == D.numel()
i = UOp.range(C.numel(), 0)
store_c = C[i].store(A[i]+B[i])
store_d = D[i].store(A[i]*B[i])
return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.numel()}")).simplify()
def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
assert A.shape[1] == B.shape[0]
i, j, k = UOp.range(C.shape[0], 0), UOp.range(C.shape[1], 1), UOp.range(A.shape[1], 2, axis_type=AxisType.REDUCE)
C = C[i, j].set(0.0)
C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
prog = C.end(i, j)
return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=()))
def custom_sum(B:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE)
B = B[0].set(0.0)
B = B[0].set(B.after(i)[0] + A[i], end=i)
return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=()))
def flip_contract_kernel(dest:UOp, src:UOp):
i = UOp.range(dest.shape[0], 0)
j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
vec = src[i, j].contract(j)
store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
def slice_sum_kernel(dest:UOp, src:UOp):
G = UOp.range(src.shape[0], 0)
slice_src = src[G, :]
reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
reg = reg.after(G)[0].set(0)
R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
ast = dest[G].set(reg[0], end=G)
return ast.sink(arg=KernelInfo(name=f"slice_sum_{src.shape[0]}_{src.shape[1]}", opts_to_apply=()))
def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
# attention without softmax
N, d = Q.shape[0], Q.shape[1]
i = UOp.range(N, 0) # output row
d_out = UOp.range(d, 1) # output column
j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
qk_acc = qk_acc.after(i, j)[0].set(0.0)
qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
qk_score = qk_acc[0] / (d ** 0.5)
out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
out_acc = out_acc.after(i, d_out)[0].set(0.0)
out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
store = O[i, d_out].store(out_acc[0])
return store.end(d_out).end(i).sink(arg=KernelInfo(name=f"simple_qkv_{N}_{d}", opts_to_apply=()))
# **** backward callbacks ****
def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
out, a, b = kernel.src[1:]
grad_a = (Tensor(gradient) @ Tensor(b).T).uop
grad_b = (Tensor(a).T @ Tensor(gradient)).uop
return (None, grad_a, grad_b)
def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
out, a, b = kernel.src[1:]
grad_a = Tensor.empty_like(Tensor(a)).custom_kernel(Tensor(gradient), Tensor(b).T, fxn=custom_gemm)[0].uop
grad_b = Tensor.empty_like(Tensor(b)).custom_kernel(Tensor(a).T, Tensor(gradient), fxn=custom_gemm)[0].uop
return (None, grad_a, grad_b)
# **** tests ****
class TestCustomKernel(unittest.TestCase):
def test_empty(self):
a = Tensor.empty(1)
a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]
a.realize()
def test_simple(self):
a = Tensor.ones(16, 16).contiguous()
b = Tensor.ones(16, 16).contiguous()
c = Tensor.empty(16, 16)
c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
out = c.flatten().tolist()
assert all(x == 2 for x in out), "all 2"
def test_simple_sharded(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
# ugly construction to get a sharded empty tensor
c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
out = c.flatten().tolist()
assert all(x == 2 for x in out), "all 2"
def test_sharded_add_one(self):
# PYTHON backend explicitly checks for OOB access for wrong multi shape regression
devs = ("PYTHON:0", "PYTHON:1")
a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
assert (c == 2).all().item()
def test_multioutput(self):
a = Tensor.full((16, 16), 3.).contiguous()
b = Tensor.full((16, 16), 3.).contiguous()
c = Tensor.empty(16, 16)
d = Tensor.empty(16, 16)
c,d = Tensor.custom_kernel(c,d,a,b, fxn=custom_elementwise_addmul_kernel)[:2]
Tensor.realize(c,d)
assert all(x == 6 for x in c.flatten().tolist()), "all 6"
assert all(x == 9 for x in d.flatten().tolist()), "all 9"
def test_arange(self):
ref = Tensor.arange(100)
tst = Tensor.empty_like(ref)
tst = tst.custom_kernel(fxn=custom_arange_kernel)[0]
self.assertTrue((ref == tst).all().item())
def test_eye(self):
ref = Tensor.eye(1024).clone().realize()
tst = Tensor.empty_like(ref)
tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
self.assertTrue((ref == tst).all().item())
@unittest.skip("contract shouldn't be supported here")
def test_flip_contract(self):
a = Tensor.randn(10,4)
b = Tensor.empty_like(a)
b = b.custom_kernel(a, fxn=flip_contract_kernel)[0]
self.assertTrue((a.flip(1) == b).all().item())
def test_noncontig(self):
a = Tensor.ones(16, 16).contiguous()
tst = Tensor.empty_like(a)
b = a+1
b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0]
self.assertTrue((b_p1 == 3).all().item())
def test_sum(self):
a = Tensor([1.0, 2, 3, 4, 5])
tst = Tensor.empty(1)
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 15)
def test_sum_int(self):
a = Tensor([1, 2, 3, 4, 5])
tst = Tensor.empty(1, dtype=a.dtype)
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 15)
def test_slice_sum(self):
A = Tensor.randn(16, 16).contiguous()
B = Tensor.empty(16)
B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
self.assertTrue(B.allclose(A.sum(1)).item())
def test_gemm(self):
N = 16
a = Tensor.randn(N, N)
b = Tensor.randn(N, N)
c = Tensor.empty(N, N)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
err = (tst - (a@b)).square().max()
self.assertLess(err.item(), 1e-6)
def test_gemm_multi(self):
devs = ("CPU:0", "CPU:1")
N = 16
a = Tensor.randn(N, N).shard_(devs, axis=0)
b = Tensor.randn(N, N).to(devs)
c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
err = (tst - (a@b)).square().max()
self.assertLess(err.item(), 1e-6)
def test_gemm_backward_custom(self): self.test_gemm_backward(True)
# NOTE: grad_fxn doesn't work with pyrender
def test_gemm_backward(self, custom_backward_gemm=False):
N = 4
a_rand = Tensor.randn(N, 8)
b_rand = Tensor.randn(8, N)
Tensor.realize(a_rand, b_rand)
a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
c = Tensor.empty(N, N)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
tst.sum().backward()
grad_a, grad_b = a.grad, b.grad
Tensor.realize(tst, grad_a, grad_b)
a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
ref = (a@b)
ref.sum().backward()
real_grad_a, real_grad_b = a.grad, b.grad
Tensor.realize(ref, real_grad_a, real_grad_b)
err = (tst - ref).square().max()
self.assertLess(err.item(), 1e-6)
err = (grad_a - real_grad_a).square().max()
self.assertLess(err.item(), 1e-6)
err = (grad_b - real_grad_b).square().max()
self.assertLess(err.item(), 1e-6)
def test_simple_qkv(self):
N, d = 8, 4
Q = Tensor.randn(N, d)
K = Tensor.randn(N, d)
V = Tensor.randn(N, d)
O = Tensor.empty(N, d)
O_custom = Tensor.custom_kernel(O, Q, K, V, fxn=lambda o,q,k,v: simple_qkv_kernel(o,q,k,v))[0]
O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
Tensor.realize(O_custom, O_ref)
err = (O_custom - O_ref).square().max()
self.assertLess(err.item(), 1e-6)
def test_multi_after_schedule_order(self):
"""Test correct scheduling order when custom_kernel has multiple outputs.
custom_kernel with 4 arguments creates 4 AFTERs from the same kernel.
The custom_kernel depends on both A2 and B2, so it must be scheduled after both.
E only depends on A2, so E can run before custom_kernel finishes waiting for B2.
Expected schedule order: [A2, B2, E, custom_addmul, final_sum]
The custom_addmul kernel should be at index 3.
"""
A, B = Tensor.empty(4, 4), Tensor.empty(4, 4)
A2 = (A + 1).contiguous() # kernel 0: depends on A
B2 = (B * 2).contiguous() # kernel 1: depends on B
C, D = Tensor.empty(4, 4), Tensor.empty(4, 4)
C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel) # depends on A2 AND B2
E = (A2 * 3).contiguous() # kernel 2: depends only on A2
result = (C + D + E).sum() # kernel 3: custom_addmul, then kernel 4: sum
schedule = result.schedule_linear().src
# Find the custom_addmul kernel position
custom_idx = next((i for i, item in enumerate(schedule)
if hasattr(item.src[0], "arg") and hasattr(item.src[0].arg, "name")
and "custom_addmul" in item.src[0].arg.name), None)
self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
def test_invalids_into_custom_kernel_no_empty_kernel(self):
from tinygrad.engine.realize import compile_linear
a = Tensor.full((4, 4), 3.).contiguous()
b = Tensor.full((4, 4), 2.).contiguous()
Tensor.realize(a, b)
out = Tensor.invalids(*a.shape, dtype=a.dtype)
out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
compiled = compile_linear(out.schedule_linear())
for call in compiled.src:
prg = call.src[0]
if prg.op is not Ops.PROGRAM: continue
self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}")
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
def test_invalids_into_custom_kernel_with_beam(self):
a = Tensor.full((4, 4), 3.).contiguous()
b = Tensor.full((4, 4), 2.).contiguous()
Tensor.realize(a, b)
with Context(BEAM=1, IGNORE_BEAM_CACHE=1):
out = Tensor.invalids(*a.shape, dtype=a.dtype)
out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
result = out.flatten().tolist()
self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}")
@unittest.skip("what are anonymous buffers?")
def test_anonymous_buffers_in_function(self):
"""Test that custom kernels with anonymous output buffers work inside @function."""
a = Tensor.full((4, 4), 3.).contiguous()
b = Tensor.full((4, 4), 2.).contiguous()
Tensor.realize(a, b)
def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
i = UOp.range(o1.numel(), 0)
store_o1 = o1[i].store(A[i]+B[i])
store_o2 = o2[i].store(A[i]+B[i]+2)
return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.numel()}")).simplify()
from tinygrad import function
@function(precompile=True)
def run(x:Tensor, w:Tensor) -> Tensor:
out = Tensor.invalids(*x.shape, dtype=x.dtype)
tmp = Tensor.invalids(*x.shape, dtype=x.dtype)
out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
return out+tmp
result = run(a, b).flatten().tolist()
expected = (3+2)*2+2
assert all(x == expected for x in result), f"expected all {expected}, got {result}"
def test_custom_kernel_sched(self, use_custom=False):
x = Tensor.arange(32).reshape(8, 4).clone().realize()
y = Tensor.empty_like(x)
y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
if use_custom:
z = Tensor.empty_like(x)
z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
else: z = y.T.T+1
GlobalCounters.reset()
z.realize()
self.assertEqual(GlobalCounters.kernel_count, 2)
self.assertEqual(z.tolist(), x.add(2).tolist())
@unittest.expectedFailure
def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
@unittest.expectedFailure
def test_sliced_buffer_function(self):
x = Tensor.arange(32).reshape(8, 4).clone().realize()
from tinygrad import function
@function(precompile=True)
def run(x:Tensor) -> Tensor:
y = Tensor.invalids(*x.shape, dtype=x.dtype)
return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
y = run(x[0]).realize()
# it's copying the input and the output
self.assertEqual(GlobalCounters.kernel_count, 1)
self.assertEqual(y.tolist(), [1, 2, 3, 4])
@Context(DEV="CPU")
def test_simple_from_source(self):
a = Tensor([0., 1., 2.]).realize()
src = "void test_src(float* restrict a) { a[0] = 1.0; }"
# TODO: it currently requires a compiler for Ops.BINARY
from tinygrad.device import Device
binary = Device[a.device].renderer.compiler.compile(src)
def custom_src_kernel(A:UOp) -> UOp:
sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
self.assertEqual(a.tolist(), [1., 1., 2.])
class TestUOpReduce(unittest.TestCase):
def test_uop_sum(self):
a = Tensor([1.0, 2, 3, 4, 5])
self.assertAlmostEqual(Tensor(a.uop.sum(axis=0)).item(), 15.0)
def test_uop_sum_2d(self):
a = Tensor.arange(6).reshape(2, 3).float()
result = Tensor(a.uop.sum(axis=1)).numpy()
assert result[0] == 3 and result[1] == 12
def test_uop_sum_all(self):
a = Tensor.arange(6).reshape(2, 3).float()
self.assertAlmostEqual(Tensor(a.uop.sum()).item(), 15.0)
def test_uop_sum_keepdim(self):
a = Tensor.arange(6).reshape(2, 3).float()
result = Tensor(a.uop.sum(axis=1, keepdim=True))
assert result.shape == (2, 1)
def test_uop_sum_negative_axis(self):
a = Tensor.arange(6).reshape(2, 3).float()
result = Tensor(a.uop.sum(axis=-1)).numpy()
assert result[0] == 3 and result[1] == 12
def test_uop_sum_multi_axis(self):
a = Tensor.arange(24).reshape(2, 3, 4).float()
ref = a.sum(axis=(0, 2)).numpy()
result = Tensor(a.uop.sum(axis=(0, 2))).numpy()
for i in range(3): self.assertAlmostEqual(result[i], ref[i])
def test_uop_sum_dtype(self):
a = Tensor([1.0, 2, 3], dtype=dtypes.float16)
result = Tensor(a.uop.sum(axis=0, dtype=dtypes.float32))
self.assertEqual(result.dtype, dtypes.float)
self.assertAlmostEqual(result.item(), 6.0, places=2)
def test_uop_prod(self):
a = Tensor([1.0, 2, 3, 4, 5])
self.assertAlmostEqual(Tensor(a.uop.prod(axis=0)).item(), 120.0)
def test_uop_max(self):
a = Tensor([1.0, 5, 3, 2, 4])
self.assertAlmostEqual(Tensor(a.uop.max(axis=0)).item(), 5.0)
def test_uop_max_2d(self):
a = Tensor([[1, 5, 3], [4, 2, 6]]).float()
result = Tensor(a.uop.max(axis=0)).numpy()
assert result[0] == 4 and result[1] == 5 and result[2] == 6
def test_uop_std(self):
a = Tensor([2.0, 4, 4, 4, 5, 5, 7, 9])
self.assertAlmostEqual(Tensor(a.uop.std()).item(), a.std().item(), places=5)
class TestUOpWhere(unittest.TestCase):
def test_uop_where_both_const(self):
cond = Tensor([True, False, True])
result = Tensor(cond.uop.where(1, 0))
self.assertEqual(result.tolist(), [1, 0, 1])
result = Tensor(cond.uop.where(1.5, 0))
self.assertEqual(result.tolist(), [1.5, 0, 1.5])
if __name__ == '__main__':
unittest.main()

9
tinygrad_repo/test/backend/test_device.py Normal file
View File

@@ -0,0 +1,9 @@
import unittest
from tinygrad import Device
class TestDeviceCount(unittest.TestCase):
def test_count(self):
self.assertGreaterEqual(Device[Device.DEFAULT].count(), 1)
if __name__ == "__main__":
unittest.main()

419
tinygrad_repo/test/backend/test_dtype.py Normal file
View File

@@ -0,0 +1,419 @@
import contextlib, unittest, math
import numpy as np
import torch
from typing import Any, List
from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, DEV
from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
from tinygrad import Context, Device, Tensor, dtypes
from hypothesis import given, settings, strategies as strat
from test.helpers import rand_for_dtype
from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
import pytest
pytestmark = pytest.mark.filterwarnings("ignore")
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
def get_available_cast_dtypes(dtype: DType) -> List[DType]:
dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
return dts
def _to_torch_storage_type(dtype:DType):
if dtype == dtypes.bfloat16: return torch.float32
if dtype in dtypes.fp8s: return torch.float32
return _to_torch_dtype(dtype)
def _test_to_np(a:Tensor, np_dtype, target):
if DEBUG >= 2: print(a)
na = a.numpy()
if DEBUG >= 2: print(na, na.dtype, a.uop.base.realized)
try:
assert na.dtype == np_dtype
np.testing.assert_allclose(na, target)
except AssertionError as e:
raise AssertionError(f"\ntensor {a.numpy()} does not match target {target} with np_dtype {np_dtype}") from e
def _test_op(fxn, target_dtype:DType, target):
_assert_eq(fxn(), target_dtype, target)
def _test_cast(a:Tensor, target_dtype:DType):
if a.is_floating_point() and dtypes.is_unsigned(target_dtype):
# converting negative float to unsigned integer is undefined
a = a.abs()
expected = list(a.numpy().astype(_to_np_dtype(target_dtype)))
if target_dtype in dtypes.fp8s: expected = [truncate[target_dtype](x) for x in expected]
_test_op(lambda: a.cast(target_dtype), target_dtype, expected)
def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist()
if target_dtype in dtypes.fp8s: expected = [fp8_to_float(x, target_dtype) for x in expected]
_test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)
class TestDType(unittest.TestCase):
DTYPE: Any = None
DATA: Any = None
@classmethod
def setUpClass(cls):
if cls.DTYPE is None: raise unittest.SkipTest("base class")
cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes)
def test_to_np(self):
_test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
def test_casts_to(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
_test_cast(Tensor(self.DATA, dtype=dtype), self.DTYPE)
def test_casts_from(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
_test_cast(Tensor(self.DATA, dtype=self.DTYPE), dtype)
def test_same_size_ops(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype.itemsize == self.DTYPE.itemsize:
_test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
def test_upcast_ops(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype.itemsize > self.DTYPE.itemsize:
_test_ops(a_dtype=self.DTYPE, b_dtype=dtype)
def test_upcast_to_ops(self):
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype.itemsize < self.DTYPE.itemsize:
_test_ops(a_dtype=dtype, b_dtype=self.DTYPE)
def test_bitcast(self):
if self.DTYPE == dtypes.bool: raise unittest.SkipTest("no bools in bitcast")
for dtype in get_available_cast_dtypes(self.DTYPE):
if dtype != dtypes.bool:
_test_bitcast(Tensor(self.DATA[:8], dtype=self.DTYPE), dtype)
@unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
def test_uint_overflow(self):
if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
v = self.DTYPE.max
_test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
_test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)
def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
_assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
_assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
_assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16])
_assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]])
_assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4)))
_assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4)))
_assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4)))
class TestFp8sConversions(unittest.TestCase):
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
def test_float_to_fp8e4m3(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
def test_float_to_fp8e4m3_extreme_values(self):
for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX))
def test_float_to_fp8e5m2(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
def test_float_to_fp8e5m2_extreme_values(self):
for x in [FP8E5M2_MAX, FP8E5M2_MAX*1.01, -FP8E5M2_MAX, -FP8E5M2_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e4m3_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fn).float().item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e5m2_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
def test_float_to_fp8e4m3fnuz(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
def test_float_to_fp8e4m3fnuz_extreme_values(self):
for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
@given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
def test_float_to_fp8e5m2fnuz(self, x):
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
def test_float_to_fp8e5m2fnuz_extreme_values(self):
for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e4m3fnuz_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
@given(strat.integers(min_value=0, max_value=255))
def test_fp8e5m2fnuz_to_float(self, x):
np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
class TestBFloat16DType(unittest.TestCase):
def test_bf16_to_float(self):
_test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
def test_float_to_bf16(self):
_test_cast(Tensor([100000], dtype=dtypes.float32), dtypes.bfloat16)
def test_bf16(self):
t = Tensor([10000, -1, -1000, -10000, 20]).cast(dtypes.bfloat16)
t.realize()
back = t.cast(dtypes.float32)
assert tuple(back.numpy().tolist()) == (9984., -1, -1000, -9984, 20)
class TestBFloat16DTypeCast(unittest.TestCase):
def test_f16_to_bf16_conversion(self):
original_tensor = Tensor([1.0, 2.0, 3.0], dtype=dtypes.float16)
converted_tensor = original_tensor.cast(dtypes.bfloat16)
self.assertEqual(converted_tensor.dtype, dtypes.bfloat16)
back_to_float32 = converted_tensor.cast(dtypes.float32)
original_to_float32 = original_tensor.cast(dtypes.float32)
np.testing.assert_allclose(back_to_float32.numpy(), original_to_float32.numpy(), rtol=1e-2, atol=1e-3)
def test_f16_to_bf16_edge_cases(self):
edge_cases = Tensor([0.0, -0.0, float('inf'), float('-inf'), float('nan')], dtype=dtypes.float16)
converted = edge_cases.cast(dtypes.bfloat16).cast(dtypes.float32)
np.testing.assert_equal(converted.numpy(), edge_cases.cast(dtypes.float32).numpy())
def test_f16_to_bf16_range_precision(self):
large_value = Tensor([65504.0], dtype=dtypes.float16) # Max representable in float16
small_value = Tensor([6.1035e-5], dtype=dtypes.float16) # Smallest positive normal float16
large_converted = large_value.cast(dtypes.bfloat16).cast(dtypes.float32)
small_converted = small_value.cast(dtypes.bfloat16).cast(dtypes.float32)
np.testing.assert_allclose(large_converted.numpy(), large_value.cast(dtypes.float32).numpy(), rtol=1e-2, atol=1e-3)
np.testing.assert_equal(small_converted.numpy(), small_value.cast(dtypes.float32).numpy())
def test_f16_to_bf16_randomized(self):
np.random.seed(42) # For reproducibility
random_values = Tensor(np.random.uniform(-65504, 65504, 1000), dtype=dtypes.float16)
converted = random_values.cast(dtypes.bfloat16).cast(dtypes.float32)
np.testing.assert_allclose(converted.numpy(), random_values.cast(dtypes.float32).numpy(), rtol=1e-2, atol=1e-3)
class TestHalfDType(TestDType): DTYPE = dtypes.half
class TestEmulatedHalf(TestHalfDType):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="half"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestFloatDType(TestDType):
DTYPE = dtypes.float
def test_float_to_uint(self):
_test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
[0, 0, 1])
@unittest.skipUnless(dtypes.double in supported_dtypes, f"no double on {Device.DEFAULT}")
class TestDoubleDType(TestDType):
DTYPE = dtypes.double
@unittest.skipIf((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"CUDA", "NV"}) or \
isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR") # TODO: why not?
def test_float64_increased_precision(self):
for func in [
lambda t: t.exp(),
lambda t: t.exp2(),
lambda t: t.log(),
lambda t: t.log2(),
lambda t: t.sqrt(),
lambda t: t.rsqrt(),
lambda t: t.sin(),
lambda t: t.cos(),
lambda t: t.tan(),
lambda t: t.sigmoid(),
]:
a = [2, 3, 4]
np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)
def test_float64_to_float32_cast_inf(self):
_test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
dtypes.float32, [float('inf'), 3.4e38, 1, 0])
class TestInt8DType(TestDType):
DTYPE = dtypes.int8
@unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
def test_int8_to_uint8_negative(self):
_test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint8), dtypes.uint8, [255, 254, 253, 252])
def test_int8_to_uint16_negative(self):
_test_op(lambda: Tensor([-1, -2, -3, -4], dtype=dtypes.int8).cast(dtypes.uint16), dtypes.uint16, [2**16-1, 2**16-2, 2**16-3, 2**16-4])
def test_bitcast_alt(self):
a = Tensor([72, -90, 27, 40, -53, 70, 96, 51], dtype=dtypes.int8).bitcast(dtypes.short)
self.assertListEqual(a.tolist(), [-22968, 10267, 18123, 13152])
class TestUint8DType(TestDType):
DTYPE = dtypes.uint8
@unittest.skipIf(Device.DEFAULT == "CUDA" or isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "cuda saturation works differently")
def test_uint8_to_int8_overflow(self):
_test_op(lambda: Tensor([255, 254, 253, 252], dtype=dtypes.uint8).cast(dtypes.int8), dtypes.int8, [-1, -2, -3, -4])
class TestBitCast(unittest.TestCase):
@given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
def test_shape_change_bitcast(self, dt1, dt2):
data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
if dt2 in dtypes.fp8s:
expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
_test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())
def test_shape_change_bitcast_exceptions(self):
with self.assertRaises(RuntimeError):
# should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16)
def test_bitcast_float_to_int32(self):
a = Tensor([1.,2,3])
b = a.bitcast(dtypes.int32)
assert b.numpy()[0] == 0x3f800000
def test_bitcast_upcasted(self):
a = Tensor.zeros(100, 4, dtype=dtypes.int32).contiguous() + 0x3f800000
b = a.bitcast(dtypes.float32)
assert b.numpy()[0,0] == 1.
class TestInt16DType(TestDType): DTYPE = dtypes.int16
class TestUint16DType(TestDType):
DTYPE = dtypes.uint16
def test_uint16_to_int8_overflow(self):
_test_op(lambda: Tensor([2**16-1, 2**16-2, 1, 0], dtype=dtypes.uint16).cast(dtypes.int8), dtypes.int8, [-1, -2, 1, 0])
class TestInt32DType(TestDType): DTYPE = dtypes.int32
class TestUint32DType(TestDType): DTYPE = dtypes.uint32
class TestInt64DType(TestDType): DTYPE = dtypes.int64
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
class TestEmulatedInt64DType(TestInt64DType):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestUint64DType(TestDType):
DTYPE = dtypes.uint64
def test_uint64_load(self):
assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
class TestEmulatedUInt64DType(TestUint64DType):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestBoolDType(TestDType): DTYPE = dtypes.bool
class TestBFloat16Type(TestDType): DTYPE = dtypes.bfloat16
class TestEmulatedBFloat16Type(TestBFloat16Type):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="bfloat16"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestFp8e4m3(TestDType): DTYPE = dtypes.fp8e4m3
class TestEmulatedFp8e4m3(TestFp8e4m3):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e4m3"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestFp8e5m2(TestDType): DTYPE = dtypes.fp8e5m2
class TestEmulatedFp8e5m2(TestFp8e5m2):
@classmethod
def setUpClass(cls):
cls.stack = contextlib.ExitStack()
cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e5m2"))
cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
@classmethod
def tearDownClass(cls): cls.stack.close()
class TestImplicitFunctionTypeChange(unittest.TestCase):
def test_functions(self):
result = []
for func in [
lambda t: t.exp(),
lambda t: t.exp2(),
lambda t: t.log(),
lambda t: t.log2(),
lambda t: t.sqrt(),
lambda t: t.sin(),
]:
t = func(Tensor([4.0, 3.0])).max() == func(Tensor([4.0, 3.0]))
result.append(t.numpy().sum())
assert all(result)
class TestTensorMethod(unittest.TestCase):
@given(strat.sampled_from(core_dtypes))
def test_abs_diff(self, dt):
if dt == dtypes.bool or dt not in supported_dtypes: return
a, b = Tensor([2], dtype=dt), Tensor([1], dtype=dt)
ret = (a - b).abs()
np.testing.assert_allclose(ret.numpy(), np.abs(a.numpy()-b.numpy()))
class TestDtypeUsage(unittest.TestCase):
def test_max_w_alu(self):
for d in dtypes.ints:
if d in supported_dtypes:
t = Tensor([[1, 2], [3, 4]], dtype=d)
(t*t).max().item()
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
class TestOpsBFloat16(unittest.TestCase):
def test_cast(self):
# TODO: helper_test_op breaks in unrelated part
data = [60000.0, 70000.0, 80000.0]
np.testing.assert_allclose(Tensor(data).cast("bfloat16").numpy(), torch.tensor(data).type(torch.bfloat16).float().numpy())
# some CPUs there is no native bfloat16 sqrt
@unittest.skipIf(Device.DEFAULT == "CPU", "no approximation")
def test_no_approximation(self):
data = [326.0, 339.0, 10603200512.0]
expected = torch.tensor(data, dtype=torch.bfloat16).sqrt().float().numpy()
np.testing.assert_allclose(Tensor(data, dtype=dtypes.bfloat16).sqrt().numpy(), expected)
if __name__ == '__main__':
unittest.main()

383
tinygrad_repo/test/backend/test_dtype_alu.py Normal file
View File

@@ -0,0 +1,383 @@
import unittest, operator, math
from tinygrad import Context, Tensor, dtypes, Device
from tinygrad.dtype import DType, truncate, fp8_to_float
from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv
from tinygrad.tensor import _to_np_dtype
from tinygrad.runtime.ops_python import from_storage_scalar
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
from tinygrad.uop import Ops
import numpy as np
import pytest
from hypothesis import assume, given, strategies as strat, settings
pytestmark = pytest.mark.filterwarnings("ignore")
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
print(settings.default)
dtypes_float = (dtypes.float16, dtypes.float32, dtypes.float64)
dtypes_int = (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
dtypes_bool = (dtypes.bool,)
binary_operations = [operator.add, operator.sub, operator.mul, operator.lt, operator.eq]
integer_binary_operations = binary_operations + [(Tensor.bitwise_xor, np.bitwise_xor), (Tensor.bitwise_and, np.bitwise_and),
(Tensor.bitwise_or, np.bitwise_or), (Tensor.maximum, np.maximum), operator.mod]
integer_unary_operations = [operator.neg]
unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin),
(Tensor.sqrt, np.sqrt), (Tensor.reciprocal, np.reciprocal), (Tensor.cos, np.cos)]
# TODO: enable this (this is a dtype issue)
#binary_operations.append(operator.truediv)
# TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
if ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"})
or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)):
unary_operations.remove((Tensor.sin, np.sin))
unary_operations.remove((Tensor.cos, np.cos))
# transcendental isn't accurate enough
if Ops.SQRT not in Device[Device.DEFAULT].renderer.code_for_op: unary_operations.remove((Tensor.sqrt, np.sqrt))
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
class ht:
float64 = strat.floats(width=64, allow_subnormal=False)
float32 = strat.floats(width=32, allow_subnormal=False)
float16 = strat.floats(width=16, allow_subnormal=False)
uint8 = strat.integers(0, 255)
uint16 = strat.integers(0, 65535)
uint32 = strat.integers(0, 2**32-1)
uint64 = strat.integers(0, 2**64-1)
int8 = strat.integers(-128, 127)
int16 = strat.integers(-32768, 32767)
int32 = strat.integers(-2147483648, 2147483647)
int64 = strat.integers(-9223372036854775808, 9223372036854775807)
bool = strat.booleans()
ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0) # filter subnormal bfloat16
ht.fp8e4m3 = ht.uint8
ht.fp8e5m2 = ht.uint8
ht.fp8e4m3fnuz = ht.uint8
ht.fp8e5m2fnuz = ht.uint8
def universal_test(a, b, dtype, op):
if not isinstance(op, tuple): op = (op, op)
if op[0] == operator.mod and b == 0: return
# lt and max with nan is undefined in tinygrad
if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
if dtype in dtypes.fp8s and op[0] not in (operator.lt, operator.eq):
tensor_value = fp8_to_float((op[0](ta.realize(), tb.realize())).bitcast(dtypes.uint8).item(), dtype)
numpy_value = truncate[dtype](op[1](ta.numpy(), tb.numpy()).item())
else: tensor_value, numpy_value = (op[0](ta, tb)).numpy(), op[1](ta.numpy(), tb.numpy())
if dtype in dtypes.floats:
if dtype not in supported_dtypes or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
fe, fm = dtypes.finfo(dtype)
atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
else: np.testing.assert_equal(tensor_value, numpy_value)
def universal_test_unary(a, dtype, op):
if not isinstance(op, tuple): op = (op, op)
ta = Tensor([a], dtype=dtype)
# TODO: cos does not match for large input
if op[0] == Tensor.cos and abs(a) > 30: return
if op[0] == Tensor.log and a <= 0: return
if dtype in dtypes.fp8s:
# denormals are zero
if (dtype in EMULATED_DTYPES.tolist(dtypes) or dtype not in supported_dtypes
and abs(ta.numpy().item()) < 0.015625): return
tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
# cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
if math.isinf(v): return
else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
if dtype in dtypes.floats:
atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
else: np.testing.assert_equal(tensor_value, numpy_value)
def universal_test_cast(a, in_dtype, dtype):
tensor_value = Tensor([a], dtype=in_dtype).cast(dtype)
numpy_value = np.array([a], dtype=_to_np_dtype(in_dtype)).astype(_to_np_dtype(dtype))
np.testing.assert_equal(tensor_value.numpy(), numpy_value)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Inf and nan cases are wrong on WebGPU")
def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
if not isinstance(op1, tuple): op1 = (op1, op1)
if not isinstance(op2, tuple): op2 = (op2, op2)
if op1[0] == operator.mod and b == 0: return
# lt and max with nan is undefined in tinygrad
if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return
at, bt, ct = Tensor([a], dtype=d1), Tensor([b], dtype=d1), Tensor([c], dtype=d2)
an, bn, cn = np.array([a]).astype(_to_np_dtype(d1)), np.array([b]).astype(_to_np_dtype(d1)), np.array([c]).astype(_to_np_dtype(d2))
tensor_value = op2[0](op1[0](at, bt).cast(d2), ct).numpy()
numpy_value = op2[1](op1[1](an, bn).astype(_to_np_dtype(d2)), cn)
np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) else 1e-7)
class TestDTypeALU(unittest.TestCase):
@unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
@given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)
@given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
@given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="half")
def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
@given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
def test_bfloat16(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
@given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="bfloat16")
def test_emulated_bfloat16(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
@unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
@given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
def test_fp8e4m3(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e4m3")
def test_emulated_fp8e4m3(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
@given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
def test_fp8e5m2(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e5m2")
def test_emulated_fp8e5m2(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
@given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
def test_fp8e4m3fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
@given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
def test_fp8e5m2fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e4m3fnuz")
def test_emulated_fp8e4m3fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
@Context(EMULATED_DTYPES="fp8e5m2fnuz")
def test_emulated_fp8e5m2fnuz(self, a, b, op):
universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.float32, strat.sampled_from(unary_operations))
def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@given(ht.float16, strat.sampled_from(unary_operations))
def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
@given(ht.float16, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="half")
def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
@given(ht.bfloat16, strat.sampled_from(unary_operations))
def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
@given(ht.bfloat16, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="bfloat16")
def test_emulated_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)
@unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
@given(ht.fp8e4m3, strat.sampled_from(unary_operations))
def test_fp8e4m3_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@given(ht.fp8e4m3, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e4m3")
def test_emulated_fp8e4m3_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
@given(ht.fp8e5m2, strat.sampled_from(unary_operations))
def test_fp8e5m2_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@given(ht.fp8e5m2, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e5m2")
def test_emulated_fp8e5m2_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
@given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
def test_fp8e4m3fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
@given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
def test_fp8e5m2fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e4m3fnuz")
def test_emulated_fp8e4m3fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
@given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
@Context(EMULATED_DTYPES="fp8e5m2fnuz")
def test_emulated_fp8e5m2fnuz_unary(self, a, op):
if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)
@unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
@given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)
@unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
@given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)
@unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
@given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)
@given(ht.int8, ht.int8, strat.sampled_from(integer_binary_operations))
def test_int8(self, a, b, op): universal_test(a, b, dtypes.int8, op)
@given(ht.int16, ht.int16, strat.sampled_from(integer_binary_operations))
def test_int16(self, a, b, op): universal_test(a, b, dtypes.int16, op)
@given(ht.int32, ht.int32, strat.sampled_from(integer_binary_operations))
def test_int32(self, a, b, op): universal_test(a, b, dtypes.int32, op)
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
def test_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.int64, ht.int64, strat.sampled_from(integer_binary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
@given(ht.uint8, strat.sampled_from(integer_unary_operations))
def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
@unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
@given(ht.uint16, strat.sampled_from(integer_unary_operations))
def test_uint16_unary(self, a, op): universal_test_unary(a, dtypes.uint16, op)
@unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
@given(ht.uint32, strat.sampled_from(integer_unary_operations))
def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op)
@unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.uint64, strat.sampled_from(integer_unary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)
@given(ht.int8, strat.sampled_from(integer_unary_operations))
def test_int8_unary(self, a, op): universal_test_unary(a, dtypes.int8, op)
@given(ht.int16, strat.sampled_from(integer_unary_operations))
def test_int16_unary(self, a, op): universal_test_unary(a, dtypes.int16, op)
@given(ht.int32, strat.sampled_from(integer_unary_operations))
def test_int32_unary(self, a, op): universal_test_unary(a, dtypes.int32, op)
@given(ht.int64, strat.sampled_from(integer_unary_operations))
def test_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
@given(ht.int64, strat.sampled_from(integer_unary_operations))
@Context(EMULATED_DTYPES="long")
def test_emulated_int64_unary(self, a, op): universal_test_unary(a, dtypes.int64, op)
@given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)
@given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)
# Metal and (MOCK)CUDA and HIP and NIR behave differently than numpy for overflows
skip_overflow = ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"AMD", "NV", "CUDA"})
or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer))
@given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
@unittest.skipIf(Device.DEFAULT == "PYTHON", "TODO: fix cast inf to int32 in PYTHON")
@unittest.skip("broken on Mac")
def test_float_midcast_int32(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.float32, dtypes.int32)
@unittest.skip("broken. TODO: fix it")
@given(ht.float32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
def test_float_cast(self, a, dtype): universal_test_cast(a, dtypes.float32, dtype)
@unittest.skip("broken. TODO: fix it")
@given(ht.int32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype)
@given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
universal_test_cast(a, float_dtype, unsigned_dtype)
@unittest.skip("relied on hacks")
@given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
universal_test_cast(a, float_dtype, unsigned_dtype)
@unittest.skip("relied on hacks")
@given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
universal_test_cast(a, float_dtype, unsigned_dtype)
@unittest.expectedFailure
def test_unsafe_cast_float_to_int_failure(self):
val = float(dtypes.int32.max - 1)
t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
np.testing.assert_equal(t1.item(), t2.item())
if __name__ == '__main__':
unittest.main()

246
tinygrad_repo/test/backend/test_edgecases.py Normal file
View File

@@ -0,0 +1,246 @@
# end to end tests of tinygrad that you think might be edge cases.
# using the documentation, write code you think should work.
# you can compare the outputs to torch or numpy, or just tinygrad assert/raise while doing things that should be valid
# i'm not interested in tests that currently pass, i'm only interested in tests that you think should pass but don't.
# mark them with @unittest.expectedFailure
# all the tests in here didn't pass until bugs were fixed
# get creative! think about things that failed in pytorch or tensorflow for a long time until they were fixed.
# every test should surface a unique bug. if tinygrad throws an error saying something is not supported, this is probably not a bug.
# the tests don't have to test the same parts of the code that these current ones test, more diversity is better
# focus on making tinygrad throw runtime errors or assertions for valid things, or find clear numerical mismatches from pytorch
# confirm any bugs found are valid by doing the same thing in pytorch in the test.
# for any failing tests, explain in a comment why tinygrad is wrong and what the desired behavior should be.
# don't worry about running mypy or linters. focus on writing more of these tests and running them to confirm broken behavior.
# surface level bugs, like issues with empty tensors, input validation, or nans, are not that interesting.
# focus on bugs that would frustrate real users.
# these are not bugs, these are desired behavior. don't add failing tests for them:
# tinygrad only accepts tinygrad dtypes or strings of the tinygrad dtype.
# boolean indexing, or anything with unknown output shape of tensor at compile time isn't supported.
# invalid indexing in things like gather and one_hot is not an error in tinygrad. nothing that depends on the value is
# repeat_interleave doesn't support a tensor as the dim. check tinygrad type signature before claiming something is a bug
import unittest
import numpy as np
import torch
from tinygrad import Tensor, dtypes, nn
from tinygrad.device import Device
from tinygrad.helpers import DEV
from tinygrad.renderer.nir import NIRRenderer
MOCKGPU = DEV.interface.startswith("MOCK")
class TestNaNEdgeCases(unittest.TestCase):
# we don't need more of these. it's unclear if torch's behavior is desired here
@unittest.expectedFailure
def test_max_nan(self):
# Reductions with NaN should propagate NaN like PyTorch.
arr = [1.0, float('nan'), 3.0]
torch_out = torch.tensor(arr).max().item()
out = Tensor(arr).max().numpy()
if np.isnan(torch_out):
self.assertTrue(np.isnan(out))
else:
np.testing.assert_equal(out, torch_out)
@unittest.skip("passes on webgpu")
@unittest.expectedFailure
def test_argmax_nan(self):
# PyTorch returns the index of the NaN, tinygrad returns the index of the maximum value.
arr = [1.0, float('nan'), 3.0]
torch_idx = torch.tensor(arr).argmax().item()
idx = Tensor(arr).argmax().item()
self.assertEqual(idx, torch_idx)
@unittest.expectedFailure
def test_sort_with_nan(self):
# Sorting a tensor containing NaN should keep NaN at the end like PyTorch.
arr = [1.0, float('nan'), 3.0]
torch_vals, torch_idxs = torch.tensor(arr).sort()
vals, idxs = Tensor(arr).sort()
np.testing.assert_equal(vals.numpy(), torch_vals.numpy())
np.testing.assert_equal(idxs.numpy(), torch_idxs.numpy().astype(np.int32))
class TestEmptyTensorEdgeCases(unittest.TestCase):
# we don't need more of these
def test_sort_empty(self):
# Sorting an empty tensor works in PyTorch and should return empty
# values and indices. tinygrad raises an error instead.
torch_vals, torch_idxs = torch.tensor([]).sort()
values, indices = Tensor([]).sort()
np.testing.assert_equal(values.numpy(), torch_vals.numpy())
np.testing.assert_equal(indices.numpy(), torch_idxs.numpy().astype(np.int32))
@unittest.expectedFailure
def test_max_empty(self):
# Max on an empty tensor should also raise an error.
with self.assertRaises(RuntimeError):
torch.tensor([]).max()
with self.assertRaises(RuntimeError):
Tensor([]).max()
@unittest.expectedFailure
def test_argmax_empty(self):
# Argmax on an empty tensor should raise an error like torch does.
with self.assertRaises(RuntimeError):
torch.tensor([]).argmax()
with self.assertRaises(RuntimeError):
Tensor([]).argmax()
def test_masked_select_empty(self):
# Masked select on empty tensors should return an empty tensor.
torch_out = torch.tensor([], dtype=torch.float32).masked_select(torch.tensor([], dtype=torch.bool))
out = Tensor([], dtype=dtypes.float32).masked_select(Tensor([], dtype=dtypes.bool))
np.testing.assert_equal(out.numpy(), torch_out.numpy())
class TestDropoutProbabilityEdgeCases(unittest.TestCase):
# we don't need more of these
def test_dropout_rate_one(self):
with Tensor.train():
out = Tensor.ones(100).dropout(1.0)
np.testing.assert_allclose(out.numpy(), np.zeros(100))
def test_dropout_invalid_prob(self):
with self.assertRaises(ValueError):
torch.nn.functional.dropout(torch.ones(10), -0.1, True)
with self.assertRaises(ValueError):
with Tensor.train():
Tensor.ones(10).dropout(-0.1)
class TestInputValidation(unittest.TestCase):
# we don't need more of these, input validation bugs are not very interesting, many are WONTFIX
@unittest.expectedFailure
def test_repeat_negative(self):
# repeating with a negative value should error like PyTorch
with self.assertRaises(RuntimeError):
torch.tensor([1, 2, 3]).repeat(-1, 2)
with self.assertRaises(RuntimeError):
Tensor([1, 2, 3]).repeat(-1, 2)
def test_negative_weight_decay(self):
with self.assertRaises(ValueError):
torch.optim.AdamW([torch.tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1)
with self.assertRaises(ValueError):
nn.optim.AdamW([Tensor([1.])], lr=0.1, weight_decay=-0.1)
def test_negative_lr(self):
with self.assertRaises(ValueError):
torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=-0.1)
with self.assertRaises(ValueError):
nn.optim.SGD([Tensor([1.])], lr=-0.1)
def test_negative_momentum(self):
with self.assertRaises(ValueError):
torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1)
with self.assertRaises(ValueError):
nn.optim.SGD([Tensor([1.])], lr=0.1, momentum=-0.1)
class TestZeroFolding(unittest.TestCase):
# we don't need more of these
# folding rules treat x/x, x//x and x%x as constants even when x can be zero
@unittest.expectedFailure
def test_divide_by_self_with_zero(self):
x = Tensor([0.0, 1.0])
torch_out = torch.tensor([0.0, 1.0]) / torch.tensor([0.0, 1.0])
out = (x / x).numpy()
np.testing.assert_allclose(out, torch_out.numpy(), equal_nan=True)
@unittest.expectedFailure
def test_floordiv_by_self_with_zero(self):
x = Tensor([0])
with self.assertRaises(RuntimeError):
torch.tensor([0]) // torch.tensor([0])
with self.assertRaises(RuntimeError):
(x // x).numpy()
@unittest.expectedFailure
def test_mod_by_self_with_zero(self):
x = Tensor([0])
with self.assertRaises(RuntimeError):
torch.tensor([0]) % torch.tensor([0])
with self.assertRaises(RuntimeError):
(x % x).numpy()
class TestAssignIssues(unittest.TestCase):
# these are good failures. i'm not sure we need more, but we need to fix these.
def test_assign_permuted_view_constant(self):
# assigning to a permuted view should modify the underlying tensor
arr = np.arange(6).reshape(2, 3).astype(np.float32)
torch_tensor = torch.tensor(arr)
torch_tensor.t().copy_(torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
t = Tensor(arr).contiguous().realize()
t.permute(1, 0).assign(Tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]))
np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
def test_assign_shrink_view_constant(self):
# assigning to a shrunk view should update the base tensor
arr = np.arange(9).reshape(3, 3).astype(np.float32)
torch_tensor = torch.tensor(arr)
torch_tensor[1:3, 1:3] = torch.ones(2, 2)
t = Tensor(arr).contiguous().realize()
t.shrink(((1, 3), (1, 3))).assign(Tensor.ones(2, 2))
np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
@unittest.expectedFailure
def test_assign_broadcast(self):
# broadcasting during assign should behave like PyTorch
# NOTE: we don't want implicit dtype casting (int64 -> float32 loses precision), so this fails
torch_tensor = torch.zeros(3, 5)
torch_tensor[:] = torch.arange(5)
t = Tensor.zeros(3, 5)
t.assign(Tensor.arange(5))
np.testing.assert_allclose(t.numpy(), torch_tensor.numpy())
class TestUOpValidationIssue(unittest.TestCase):
# these fail with UOp verification error.
# we want more of these with diverse errors!
@unittest.skipIf(MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "hangs gpuocelot, NIR cannot render")
def test_tensor_index_overflow(self):
val = Tensor([1])
big = val.expand(2**31 + 3)
idx = Tensor([0, 2**31 + 2])
np.testing.assert_equal(big[idx].numpy(), np.array([1, 1]))
def test_float_floordiv_scalar(self):
(Tensor.arange(4, dtype=dtypes.float32) // 2).realize()
def test_float_floordiv_tensor(self):
(Tensor.arange(4, dtype=dtypes.float32) // Tensor.ones(4, dtype=dtypes.float32)).realize()
class TestEdgeCases(unittest.TestCase):
# add tests exposing new and diverse kinds of bugs that might impact real users here
def test_circular_pad_negative(self):
# negative pads with circular mode should wrap like PyTorch
arr = np.arange(9).reshape(1, 1, 3, 3).astype(np.float32)
torch_out = torch.nn.functional.pad(torch.tensor(arr), (1, -1, 1, -1), mode='circular')
out = Tensor(arr).pad((1, -1, 1, -1), mode='circular')
np.testing.assert_equal(out.numpy(), torch_out.numpy())
def test_arange_float_step(self):
# float steps should match PyTorch exactly
torch_out = torch.arange(0, 2, 0.3).numpy()
out = Tensor.arange(0, 2, 0.3).numpy()
np.testing.assert_allclose(out, torch_out, atol=1e-7)
@unittest.skip("this is flaky")
@unittest.expectedFailure
def test_topk_ties_indices(self):
# topk should match PyTorch tie-breaking behavior when values are equal
arr = [1.0, 1.0, 1.0, 1.0]
_, ti = torch.tensor(arr).topk(2)
_, i = Tensor(arr).topk(2)
np.testing.assert_equal(i.numpy(), ti.numpy().astype(np.int32))
if __name__ == "__main__":
unittest.main()

149
tinygrad_repo/test/backend/test_encodings.py Normal file
View File

@@ -0,0 +1,149 @@
import unittest
from tinygrad import Device
from tinygrad.uop.ops import UOp, Ops
from tinygrad.dtype import dtypes
from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
class TestEncodingsX86(unittest.TestCase):
# NOTE: x86 supports a single displacement as memory address and index without base memory address
# these have no use cases so they aren't supported
def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u])
# displacement of 0 isn't emitted
def test_base_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RDI)
# mov edi, dword ptr [rdi]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
# rsp/r12 require a sib byte when used as base memory address
def test_rsp_base_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RSP)
# mov esp, dword ptr [rsp]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
# rbp/r13 require a displacement when used as base memory address
def test_rbp_base_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RBP)
# mov ebp, dword ptr [rbp + 0]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
# test [base + index*scale]
def test_base_index_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0)), RAX)
# mov eax, dword ptr [rax + rdx*4]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
# rsp as index means no index
def test_rsp_index_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0)), RAX)
# mov eax, dword ptr [rax]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
# however r12 is a valid index
def test_r12_index_address(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0)), RAX)
# mov eax, dword ptr [rax + r12*4]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
# test [base + index*scale + 8bit disp]
def test_complex_address_8bit_disp(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)), RDI)
# mov edi, dword ptr [rdi + rsi*4 + 0xa]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
# test [base + index*scale + 32bit disp]
def test_complex_address_32bit_disp(self):
load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000)), RDI)
# mov edi, dword ptr [rdi + rsi*4 + 0x2710]
self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
# 8bit variants of legacy instructions subtract 1 from opcode
def test_8bit_legacy_encoding(self):
cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDX),), RAX)
# movsx eax, dl
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2"))
# accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix
def test_lower_8bits_reg(self):
cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDI),), RAX)
# movsx eax, dil
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7"))
# test 16 bit variant of legacy instruction
def test_16bit_legacy_encoding(self):
cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(dtypes.int8, RDX),), RAX)
# movsx ax, dl
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2"))
# test 64 bit variant of legacy instruction
def test_64bit_legacy_encoding(self):
cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(dtypes.int8, RDX),), RAX)
# movsx rax, dl
self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2"))
# test compact vex encoding
def test_compact_vex_encoding(self):
xmm0, xmm1 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1])
add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0])
# vaddss xmm0, xmm0, xmm1
self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1"))
# test long vex encoding
def test_long_vex_encoding(self):
xmm0, xmm8 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[8])
add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0])
# vaddss xmm0, xmm0, xmm8
self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
# test ymm encoding
def test_ymm_encoding(self):
xmm0, xmm1 = def_reg(dtypes.float32.vec(8), XMM[0]), def_reg(dtypes.float32.vec(8), XMM[1])
add = ins(X86Ops.VADDPS, dtypes.float32.vec(8), (xmm0, xmm1), XMM[0])
# vaddps ymm0, ymm0, ymm1
self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
# test encoding where register is in the immediate field
def test_reg_in_imm_field(self):
xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0])
# vblendvps xmm0, xmm0, xmm1, xmm2
self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20"))
# when writting to mem the uop takes the store form where dtype is void and there's no definition
def test_write_mem(self):
base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
xmm0 = def_reg(dtypes.float32, XMM[0])
extr = ins(X86Ops.VPEXTRD, dtypes.void, (base, index, disp, xmm0, imm(dtypes.uint8, 0)))
# vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
# test two address instruction with fused load works
def test_two_address_load(self):
base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
cmove = ins(X86Ops.CMOVE, dtypes.int32, (base, index, disp), RAX)
# cmove eax, dword ptr [rdi + rsi*4 + 0xa]
self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
# test instruction where displacement and imm have the same value
def test_disp_imm_same_value(self):
base, index, disp = def_reg(dtypes.int8.ptr(), RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10)
mov = ins(X86Ops.MOVi, dtypes.void, (base, index, disp, disp))
# mov byte ptr [rdi + rsi + 0xa], 0xa
self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10)
imul = ins(X86Ops.IMULi, dtypes.int32, (base, index, disp) + (imm(dtypes.int32, 10),), RDI)
# imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
# cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
def test_cmove_ignore_cmp(self):
cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
# cmove edx, eax
self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
if __name__ == "__main__":
unittest.main()

319
tinygrad_repo/test/backend/test_graph.py Normal file
View File

@@ -0,0 +1,319 @@
import numpy as np
import functools, unittest, ctypes
from tinygrad.device import Device, Buffer
from tinygrad.tensor import Tensor
from tinygrad.helpers import Context, from_mv
from tinygrad.dtype import dtypes
from tinygrad.engine.jit import MultiGraphRunner
from tinygrad.engine.realize import run_linear, compile_linear
from tinygrad.uop.ops import UOp, Ops, buffers
from test.helpers import needs_second_gpu
np.random.seed(1337)
Tensor.manual_seed(1337)
BUF_SIZE = 4096
RUN_CNT = 5
# cache AST by (device, num_inputs)
cached_asts: dict[tuple[str, int], UOp] = {}
def get_ast(device:str, num_inputs:int) -> UOp:
if (device, num_inputs) not in cached_asts:
with Context(DEBUG=0):
fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
s = fst[0]
for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
cached_asts[(device, num_inputs)] = s.schedule_linear().src[-1].src[0]
return cached_asts[(device, num_inputs)]
def make_buffer(device, size=BUF_SIZE, fill=False):
buf = Buffer(device, size, dtypes.int).ensure_allocated()
if fill:
with Context(DEBUG=0):
buf.copyin(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized.as_memoryview())
return buf
def make_view(base, offset_elems, size_elems):
return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
if buf not in cache:
cache[buf] = u = UOp.new_buffer(buf.device, buf.size, buf.dtype)
buffers[u] = buf
return cache[buf]
def make_graph(graph_cls, calls:list[UOp]):
linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)))
cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph")
return graph_cls(cf, [])
def run_schedule(calls:list[UOp]):
run_linear(UOp(Ops.LINEAR, src=tuple(calls)))
def zero_bufs(bufs):
for b in bufs:
mv = memoryview(bytearray(b.nbytes))
ctypes.memset(from_mv(mv), 0, len(mv))
b.copyin(mv)
@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
class TestGraph(unittest.TestCase):
def skip_if_no_offset(self):
if not hasattr(Device[Device.DEFAULT].allocator, "_offset"): self.skipTest("device does not support _offset")
def skip_if_not_multigraph(self):
graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
def test_order_2_writes_to_same_buf(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(5)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
]
zero_bufs([b[0]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_read_write_same_buf(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(5)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
]
zero_bufs([b[0], b[1]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0], b[1]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_write_read_same_buf(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(5)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c), metadata=()),
]
zero_bufs([b[0], b[1]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0], b[1]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_copy_writed(self):
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(4)]
c: dict[Buffer,UOp] = {}
calls = [
get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
UOp(Ops.COPY).call(get_buf_uop(b[3],c), get_buf_uop(b[0],c), metadata=()),
]
zero_bufs([b[0], b[3]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[0], b[3]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_order_copy_then_read(self):
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(4)]
c: dict[Buffer,UOp] = {}
calls = [
UOp(Ops.COPY).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
]
zero_bufs([b[1], b[3]])
run_schedule(calls)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs([b[1], b[3]])
make_graph(Device[d0].graph, calls)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
def test_read_write_several_graphs(self):
d0 = Device.DEFAULT
b = [make_buffer(d0, fill=True) for _ in range(8)]
c: dict[Buffer,UOp] = {}
calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=())]
calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c), metadata=())]
calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c), metadata=())]
out = [b[3], b[4], b[5]]
zero_bufs(out)
run_schedule(calls1 + calls2 + calls3)
expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
for _ in range(RUN_CNT):
zero_bufs(out)
make_graph(Device[d0].graph, calls1)([], {})
make_graph(Device[d0].graph, calls2)([], {})
make_graph(Device[d0].graph, calls3)([], {})
for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
@needs_second_gpu
def test_copies_2_devs(self):
self.skip_if_not_multigraph()
d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
b0 = [make_buffer(d0, fill=True) for _ in range(3)]
b1 = [make_buffer(d1, fill=True)]
c: dict[Buffer,UOp] = {}
calls = [
UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c), metadata=()),
]
out = [b1[0], b0[2]]
zero_bufs(out)
run_schedule(calls)
expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in b0 + b1}
for _ in range(RUN_CNT):
zero_bufs(out)
make_graph(Device[d0].graph, calls)([], {})
for buf in b0 + b1: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
def test_graph_offset_bufs(self):
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")
b0 = make_buffer(d0, fill=True)
b1 = make_view(b0, 0, b0.size)
b2 = make_view(b0, 0, b0.size)
c: dict[Buffer,UOp] = {}
calls = [
UOp(Ops.COPY).call(get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
]
zero_bufs([b0])
run_schedule(calls)
expected = np.frombuffer(b0.as_memoryview(), np.int32).copy()
for _ in range(RUN_CNT):
zero_bufs([b0])
make_graph(Device[d0].graph, calls)([], {})
np.testing.assert_equal(expected, np.frombuffer(b0.as_memoryview(), np.int32))
def test_partial_write_preserves_write_dep(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
base = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_src_full = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_src_lo = make_buffer(d0, fill=True)
v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
a, out = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
c: dict[Buffer,UOp] = {}
calls = [
UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
]
zero_bufs([base, out])
run_schedule(calls)
expected = {base: np.frombuffer(base.as_memoryview(), np.int32).copy(), out: np.frombuffer(out.as_memoryview(), np.int32).copy()}
for _ in range(RUN_CNT):
zero_bufs([base, out])
make_graph(Device[d0].graph, calls)([], {})
for buf in [base, out]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
def test_partial_write_preserves_read_dep(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
base = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_dst = make_buffer(d0, BUF_SIZE * 2, fill=True)
copy_src_lo = make_buffer(d0, fill=True)
v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
a, b = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
c: dict[Buffer,UOp] = {}
calls = [
UOp(Ops.COPY).call(get_buf_uop(copy_dst,c), get_buf_uop(base,c), metadata=()),
UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c), metadata=()),
]
zero_bufs([copy_dst, base])
run_schedule(calls)
expected = {copy_dst: np.frombuffer(copy_dst.as_memoryview(), np.int32).copy(), base: np.frombuffer(base.as_memoryview(), np.int32).copy()}
for _ in range(RUN_CNT):
zero_bufs([copy_dst, base])
make_graph(Device[d0].graph, calls)([], {})
for buf in [copy_dst, base]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
def test_middle_write_splits_write_dep(self):
self.skip_if_not_multigraph()
self.skip_if_no_offset()
d0 = Device.DEFAULT
base = make_buffer(d0, BUF_SIZE * 3, fill=True)
copy_src_full = make_buffer(d0, BUF_SIZE * 3, fill=True)
copy_src_mid = make_buffer(d0, fill=True)
v_lo, v_mid, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE), make_view(base, BUF_SIZE * 2, BUF_SIZE)
a, out1, out2 = make_buffer(d0, fill=True), make_buffer(d0, fill=True), make_buffer(d0, fill=True)
c: dict[Buffer,UOp] = {}
calls = [
UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
UOp(Ops.COPY).call(get_buf_uop(v_mid,c), get_buf_uop(copy_src_mid,c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c), metadata=()),
get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
]
outs = [base, out1, out2]
zero_bufs(outs)
run_schedule(calls)
expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in outs}
for _ in range(RUN_CNT):
zero_bufs(outs)
make_graph(Device[d0].graph, calls)([], {})
for buf in outs: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
if __name__ == '__main__':
unittest.main()

52
tinygrad_repo/test/backend/test_interop.py Normal file
View File

@@ -0,0 +1,52 @@
#!/usr/bin/env python
import unittest, os
import torch
import numpy as np
from tinygrad.helpers import DEV
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
MOCKGPU = DEV.interface.startswith("MOCK")
@unittest.skipIf(Device.DEFAULT not in ["METAL", "CUDA"] or MOCKGPU, f"no support on {Device.DEFAULT}")
class TestInterop(unittest.TestCase):
def setUp(self):
if Device.DEFAULT == "CUDA": self.torch_device = "cuda"
elif Device.DEFAULT == "METAL": self.torch_device = "mps"
def test_torch_interop(self):
inp = torch.rand(2, 2, 3, device=torch.device(self.torch_device))
if self.torch_device == "mps": torch.mps.synchronize()
else: torch.cuda.synchronize()
tg_data = Tensor.from_blob(inp.data_ptr(), inp.shape, dtype=_from_torch_dtype(inp.dtype))
tg_out = tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140
tg_res = tg_out.numpy()
if self.torch_device == "mps" and os.getenv("CI", "") != "":
# MPS backend out of memory: https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773
# Calculate expected value on cpu.
inp = inp.cpu()
torch_out = inp[:, :, 0] * 0.2989 + inp[:, :, 1] * 0.5870 + inp[:, :, 2] * 0.1140
np.testing.assert_allclose(tg_res, torch_out.cpu().numpy(), atol=1e-5, rtol=1e-5)
def test_torch_interop_write(self):
tg_data = Tensor.randn((4, 4), device=Device.DEFAULT)
out = torch.empty(4, 4, device=torch.device(self.torch_device), dtype=_to_torch_dtype(tg_data.dtype))
tg_out = Tensor.from_blob(out.data_ptr(), out.shape, dtype=_from_torch_dtype(out.dtype))
tg_out.assign(tg_data).realize()
Device[Device.DEFAULT].synchronize()
torch_out_np = out.cpu().numpy()
np.testing.assert_allclose(tg_data.numpy(), torch_out_np, atol=1e-5, rtol=1e-5)
if __name__ == '__main__':
unittest.main()

147
tinygrad_repo/test/backend/test_isel.py Normal file
View File

@@ -0,0 +1,147 @@
import unittest
from typing import cast
from tinygrad import Device
from tinygrad.uop import Ops
from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
from tinygrad.renderer.isa import IselContext
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
class TestIselX86(unittest.TestCase):
def isel_rewrite(self, x:UOp):
return graph_rewrite(x, cast(X86Renderer, Device[Device.DEFAULT].renderer).isel_matcher, IselContext(x), bottom_up=True)
def _check_op(self, dt_op, expr):
nargs = expr.__code__.co_argcount
for dt,op in dt_op:
with self.subTest(dtype=dt):
v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
n = self.isel_rewrite(expr(*v))
self.assertIs(n.arg, op)
def test_cmove(self):
a = UOp.variable("a", 0, 0, dtypes.int32)
b = UOp.variable("b", 0, 0, dtypes.int32)
c = (a < b).where(a, b)
d = (a != b).where(a, b)
f = c + d
n = self.isel_rewrite(f)
self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
# both comparisons become the same instruction
self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
def test_vmax(self):
dt_op = [(dtypes.float32, X86Ops.VMAXSS), (dtypes.float64, X86Ops.VMAXSD),
(dtypes.float32.vec(4), X86Ops.VMAXPS), (dtypes.float64.vec(4), X86Ops.VMAXPD)]
self._check_op(dt_op, lambda a,b: (a < b).where(b, a))
def test_vmin(self):
dt_op = [(dtypes.float32, X86Ops.VMINSS), (dtypes.float64, X86Ops.VMINSD),
(dtypes.float32.vec(4), X86Ops.VMINPS), (dtypes.float64.vec(4), X86Ops.VMINPD)]
self._check_op(dt_op, lambda a,b: (a < b).where(a, b))
def test_vfmadd(self):
dt_op = [(dtypes.float32, X86Ops.VFMADD213SS), (dtypes.float64, X86Ops.VFMADD213SD),
(dtypes.float32.vec(4), X86Ops.VFMADD213PS), (dtypes.float64.vec(4), X86Ops.VFMADD213PD)]
self._check_op(dt_op, lambda a,b,c: a * b + c)
# don't use fmadd if op being fused (mul) is used multiple times
def test_no_vfmadd(self):
dt_op = [(dtypes.float32, X86Ops.VADDSS), (dtypes.float64, X86Ops.VADDSD),
(dtypes.float32.vec(4), X86Ops.VADDPS), (dtypes.float64.vec(4), X86Ops.VADDPD)]
self._check_op(dt_op, lambda a,b: a * b + a * b)
def test_vpbroadcast(self):
a = UOp.variable("a", 0, 0, dtypes.int32)
n = self.isel_rewrite(a.broadcast(4))
# need to move src from gpr to xmm before broadcasting
self.assertTrue(n.arg is X86Ops.VPBROADCASTD and n.src[0].arg is X86Ops.VMOVD)
# if we can fuse a load we can skip the move and access memory directly
load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
n = self.isel_rewrite(load.broadcast(4))
self.assertTrue(n.arg is X86Ops.VPBROADCASTD and len(n.src) == 3)
def test_vbroadcastss(self):
a = UOp.variable("a", 0, 0, dtypes.float32)
valid = [UOp.vectorize(a, a, a, a), UOp.vectorize(a, a, a, a, a, a, a, a)]
for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VBROADCASTSS)
def test_vshufps(self):
a = UOp.variable("a", 0, 0, dtypes.float32.vec(8))
b = UOp.variable("b", 0, 0, dtypes.float32.vec(8))
c = UOp.variable("c", 0, 0, dtypes.float32)
d = UOp.variable("d", 0, 0, dtypes.float32)
valid = [UOp.vectorize(c, c, d, d),
UOp.vectorize(a.gep(0), a.gep(1), c, c),
UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
UOp.vectorize(a.gep(1), a.gep(2), a.gep(3), a.gep(0)),
UOp.vectorize(a.gep(3), a.gep(2), a.gep(1), a.gep(0), a.gep(7), a.gep(6), a.gep(5), a.gep(4)),
UOp.vectorize(a.gep(0), a.gep(0), b.gep(1), b.gep(1), a.gep(4), a.gep(4), b.gep(5), b.gep(5))]
for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
invalid = [UOp.vectorize(a.gep(0), a.gep(1), b.gep(4), b.gep(5)),
UOp.vectorize(a.gep(0), a.gep(5), b.gep(2), b.gep(3)),
UOp.vectorize(a.gep(0), a.gep(0), a.gep(0), a.gep(0), a.gep(4), a.gep(4), a.gep(4), a.gep(5)),
UOp.vectorize(a.gep(0), a.gep(0), b.gep(0), b.gep(0), a.gep(4), a.gep(4), b.gep(4), a.gep(4))]
for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
def test_vshufpd(self):
a = UOp.variable("a", 0, 0, dtypes.float64.vec(4))
b = UOp.variable("b", 0, 0, dtypes.float64.vec(4))
c = UOp.variable("c", 0, 0, dtypes.float64)
d = UOp.variable("d", 0, 0, dtypes.float64)
valid = [UOp.vectorize(c, d),
UOp.vectorize(a.gep(0), c),
UOp.vectorize(a.gep(1), b.gep(1)),
UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
UOp.vectorize(a.gep(1), a.gep(1), a.gep(3), a.gep(3))]
for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
invalid = [UOp.vectorize(c, c, c, c),
UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
UOp.vectorize(a.gep(2), b.gep(3), a.gep(2), b.gep(3)),
UOp.vectorize(a.gep(0), b.gep(1), a.gep(0), b.gep(1))]
for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
def test_vinsertps(self):
a = UOp.variable("a", 0, 0, dtypes.float32.vec(4))
b = UOp.variable("b", 0, 0, dtypes.float32.vec(4))
c = UOp.variable("c", 0, 0, dtypes.float32.vec(4))
d = UOp.variable("e", 0, 0, dtypes.float32)
# moving 0th element to position 0 does nothing so only 1 vinsertps is generated
n = self.isel_rewrite(UOp.vectorize(a.gep(0), d))
self.assertIs(n.arg, X86Ops.VINSERTPS)
self.assertIsNot(n.src[0].arg, X86Ops.VINSERTPS)
valid = [UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
UOp.vectorize(a.gep(3), b.gep(2), c.gep(1), d)]
for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
# complex address is [base + index*scale + displacement]
def test_complex_address(self):
a = UOp.variable("a", 0, 0, dtypes.int32)
load = UOp.param(0, dtypes.int32.ptr()).index(a + 1, ptr=True).load()
n = self.isel_rewrite(load)
# displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)
def test_fold_load(self):
load1 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
load2 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 1), ptr=True).load()
n = self.isel_rewrite(load1 + load2)
self.assertTrue(len(n.src) == 4)
# don't fold when used multiple times
def test_dont_fold_load(self):
load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
# used by multiple users
n = self.isel_rewrite(load + 1 + load)
self.assertTrue(len(n.src) == 2)
# used mutiple times by same user
n = self.isel_rewrite(load * load)
self.assertTrue(len(n.src) == 2)
if __name__ == "__main__":
unittest.main()

892
tinygrad_repo/test/backend/test_jit.py Normal file
View File

@@ -0,0 +1,892 @@
#!/usr/bin/env python
import unittest
import numpy as np
from hypothesis import given, settings, strategies as strat
from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
from tinygrad import Variable
from tinygrad.tensor import Tensor
from tinygrad.engine.jit import TinyJit, JitError, graph_class
from tinygrad.device import Device
from tinygrad.helpers import Context, JIT, DEV, GlobalCounters
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import Ops, UOp
from extra.models.unet import ResBlock
from tinygrad.renderer.isa.x86 import X86Renderer
def _simple_test(add, extract=lambda x: x, N=10):
for _ in range(5):
a = Tensor.randn(N, N)
b = Tensor.randn(N, N)
c = add(a, b)
np.testing.assert_allclose(extract(c).numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add, 1)
class TestJit(unittest.TestCase):
@settings(deadline=2e4)
@unittest.skipUnless(Device.DEFAULT in ["CPU"], f"no support on {Device.DEFAULT}")
@given(strat.sampled_from([Tensor.exp2, Tensor.log2, Tensor.sin]))
def test_approx_jit_timeout(self, op):
with Context(TRANSCENDENTAL=2):
model = [ResBlock(16, 24, 16) for _ in range(4)]
@TinyJit
def fw_approx(t, t2):
for l in model: t = l(t, t2)
return op(t).realize()
fw_approx(Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24))
def test_simple_jit(self):
@TinyJit
def add(a, b): return (a+b).realize()
_simple_test(add)
@unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar")
def test_hcq_core_id_runtimevar_merge(self):
N = 262144
@TinyJit
def f(x, st):
y = (x + 1).contiguous().realize()
z = x.shrink(((st, st + N),)).contiguous().realize()
return y, z
x = Tensor.arange(2*N).clone().realize()
for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
self.assertEqual(y.shape, (2*N,))
self.assertEqual(z.shape, (N,))
def test_jitbeam_triggers_beam(self):
from unittest.mock import patch
from tinygrad.helpers import getenv as _getenv
@TinyJit
def add(a, b): return (a+b).realize()
a, b = Tensor.ones(10, 10).contiguous().realize(), Tensor.ones(10, 10).contiguous().realize()
with patch("tinygrad.codegen.opt.search.beam_search", wraps=lambda k,*a,**kw: k) as mock_beam:
add(a, b)
assert mock_beam.call_count == 0
with patch("tinygrad.engine.jit.getenv", side_effect=lambda k, d=0: 1 if k == "JITBEAM" else _getenv(k, d)): add(a, b)
assert mock_beam.call_count == 1
def test_simple_jit_reset(self):
@TinyJit
def add(a, b): return (a+b).realize()
_simple_test(add)
add.reset()
_simple_test(add, N=20)
def test_simple_jit_norealize(self):
@TinyJit
def add(a, b): return (a+b)
_simple_test(add)
def test_simple_jit_norealize_list(self):
@TinyJit
def add(a, b): return [a+b]
_simple_test(add, extract=lambda x: x[0])
def test_simple_jit_norealize_dict(self):
@TinyJit
def add(a, b): return {"billy": a+b}
_simple_test(add, extract=lambda x: x["billy"])
def test_jit_input_view(self):
@TinyJit
def f(x): return (x[2:5].contiguous() + 1).realize()
for i in range(5):
x = (Tensor.arange(10).float() + i * 10).clone().realize()
np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)
def test_jit_multiple_outputs(self):
@TinyJit
def f(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize()
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
c, d, e = f(a, b)
np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(d.numpy(), a.numpy()-b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(f, 3)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "estimates are wrong for x86")
def test_global_counters_jit(self):
@TinyJit
def f(a, b):
c = (a + b).realize()
d = (c * 2).realize()
return (d - a).realize()
a, b = Tensor.randn(64, 64).realize(), Tensor.randn(64, 64).realize()
for _ in range(4):
GlobalCounters.reset()
f(a, b)
Device[a.device].synchronize()
self.assertGreater(GlobalCounters.global_mem, 0)
self.assertGreater(GlobalCounters.global_ops, 0)
def test_nothing_jitted(self):
@TinyJit
def add(a, b): return None
with self.assertRaises(JitError):
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
add(a, b)
def test_jit_zero_does_not_jit(self):
@TinyJit
def add(a, b): return (a+b).realize()
with Context(JIT=0):
for i in range(5):
a = Tensor([i])
b = Tensor([i])
c = add(a, b)
np.testing.assert_allclose(c.numpy(), 2*i)
assert_jit_cache_len(add, 0)
def test_jit_not_capturing(self):
@TinyJit
def add(a, b):
Tensor.zeros(4, 4).contiguous().realize() # no-op kernel is captured
return (a+b).realize()
for i in range(5):
a = Tensor([i])
b = Tensor([i])
c = add(a, b)
np.testing.assert_allclose(c.numpy(), 2*i)
assert_jit_cache_len(add, 2)
@TinyJit
def add2(a, b):
with Context(CAPTURING=0): # not captured
Tensor.zeros(4, 4).contiguous().realize()
return (a+b).realize()
for i in range(5):
a = Tensor([i])
b = Tensor([i])
c = add2(a, b)
np.testing.assert_allclose(c.numpy(), 2*i)
assert_jit_cache_len(add2, 1)
def test_jit_shape_mismatch(self):
@TinyJit
def add(a, b): return (a+b).realize()
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
add(a, b)
bad = Tensor.randn(20, 20)
with self.assertRaises(JitError):
add(a, bad)
def test_jit_shape_views_mismatch(self):
@TinyJit
def add(a): return (a+1).realize()
with self.assertRaises(JitError):
for i in range(1,5):
# a has an offset that the kernel doesn't know about
a = Tensor.randn(10, 10).realize()[:, i:i+2]
add(a)
def test_jit_duplicate_fail(self):
# the jit doesn't support duplicate arguments
@TinyJit
def add(a, b): return (a+b).realize()
a = Tensor.randn(10, 10)
with self.assertRaises(JitError):
add(a, a)
def test_jit_assign(self, dtype=dtypes.float32):
@TinyJit
def add(a):
a += 1
a.realize()
a = Tensor.zeros(1, dtype=dtype).contiguous().realize()
for _ in range(5): add(a)
self.assertEqual(a.item(), 5)
def test_jit_assign_int8(self): self.test_jit_assign(dtypes.int8)
def test_kwargs_jit(self):
@TinyJit
def add_kwargs(first, second): return (first+second).realize()
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
c = add_kwargs(first=a, second=b)
np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add_kwargs, 1)
def test_reorder_kwargs_jit(self):
@TinyJit
def add_kwargs(first, second): return (first/second).realize()
for _ in range(2):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
c = add_kwargs(second=b, first=a)
np.testing.assert_allclose(c.numpy(), a.numpy()/b.numpy(), atol=1e-4, rtol=1e-5)
for _ in range(2):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
c = add_kwargs(first=a, second=b)
np.testing.assert_allclose(c.numpy(), a.numpy()/b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add_kwargs, 1)
def test_array_jit(self):
@TinyJit
def add_array(a, arr): return (a+arr[0]).realize()
for _ in range(5):
a, b = Tensor.randn(10, 10).realize(), Tensor.randn(10, 10).realize()
np.testing.assert_allclose(add_array(a, [b]).numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add_array, 1)
def test_jit_copyin(self):
@TinyJit
def f(a):
return a + Tensor([1,2,3])
for _ in range(5):
b = Tensor.randn(3)
c = f(b)
np.testing.assert_allclose(c.numpy(), b.numpy()+[1,2,3], atol=1e-4, rtol=1e-5)
def test_method_jit(self):
class Fun:
def __init__(self):
self.a = Tensor.randn(10, 10)
@TinyJit
def __call__(self, b:Tensor) -> Tensor:
return (self.a+b).realize()
fun = Fun()
for _ in range(5):
b = Tensor.randn(10, 10)
c = fun(b)
np.testing.assert_allclose(c.numpy(), fun.a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(fun.__call__.func.__self__, 1)
def test_jit_size1_input(self):
@TinyJit
def f(a, b): return (a+b).realize()
a = Tensor([1, 2, 3])
for i in range(5):
np.testing.assert_allclose(f(a, Tensor([i])).numpy(), (a+i).numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(f, 1)
def test_jit_output_non_tensor_fail(self):
@TinyJit
def f(a, b, i): return (a+b).realize(), i
with self.assertRaises(JitError):
for i in range(3):
f(Tensor.randn(10, 10), Tensor.randn(10, 10), i)
def test_jit_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
return ((a+b)*rn).realize()
a = Tensor.randn(10, 10).realize() # realize these before resetting the random seed
b = Tensor.randn(10, 10).realize()
Tensor.manual_seed(1234)
jf = TinyJit(f)
res = set()
for _ in range(5):
o1 = jf(a, b)
res.add(o1.numpy()[0][0])
assert len(res) == 5, "All values should be different, rand works in jit."
Tensor.manual_seed(1234)
jf2 = TinyJit(f)
res2 = set()
for _ in range(5):
o1 = jf2(a, b)
res2.add(o1.numpy()[0][0])
assert len(res2) == 5, "All values should be different, rand works in jit."
assert res == res2, "Jit rand is not reproducible with the same seed"
Tensor.manual_seed(3421)
jf3 = TinyJit(f)
res3 = set()
for _ in range(5):
o1 = jf3(a, b)
res3.add(o1.numpy()[0][0])
assert len(res3) == 5, "All values should be different, rand works in jit."
assert res3 != res2, "Jit rand is diff with diff seeds"
def test_jit_v_nojit_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
rn = rn * a
rn2 = Tensor.randn(*a.shape)
rn2 = rn2 * b
rn = rn + rn2
rn2 = rn2 + Tensor.randn(*a.shape)
return ((a+b)*rn).realize(), ((a+b)*rn2).realize()
Tensor.manual_seed(0)
a = Tensor.randn(10, 10).realize() # realize these before resetting the random seed
b = Tensor.randn(10, 10).realize()
Tensor.manual_seed(1234)
without_jit = set()
for _ in range(5):
o1, o2 = f(a, b)
without_jit.add(o1.numpy()[0][0])
without_jit.add(o2.numpy()[0][0])
assert len(without_jit) == 10, "All values should be different."
Tensor.manual_seed(1234)
jf = TinyJit(f)
with_jit = set()
for _ in range(5):
o1, o2 = jf(a, b)
with_jit.add(o1.numpy()[0][0])
with_jit.add(o2.numpy()[0][0])
assert len(with_jit) == 10, "All values should be different."
assert with_jit == without_jit, "jit and non-jit should produce the same random values with the same seed"
def test_jit_multiple_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
rn = rn * a
rn2 = Tensor.randn(*a.shape)
rn2 = rn2 * b
rn = rn + rn2
rn2 = rn2 + Tensor.randn(*a.shape)
return ((a+b)*rn).realize(), ((a+b)*rn2).realize()
a = Tensor.randn(10, 10).realize() # realize these before resetting the random seed
b = Tensor.randn(10, 10).realize()
Tensor.manual_seed(1234)
jf = TinyJit(f)
res = set()
for _ in range(5):
o1, o2 = jf(a, b)
res.add(o1.numpy()[0][0])
res.add(o2.numpy()[0][0])
assert len(res) == 10, "All values should be different, rand works in jit."
Tensor.manual_seed(1234)
jf2 = TinyJit(f)
res2 = set()
for _ in range(5):
o1, o2 = jf2(a, b)
res2.add(o1.numpy()[0][0])
res2.add(o2.numpy()[0][0])
assert len(res2) == 10, "All values should be different, rand works in jit."
assert res == res2, "Jit rand is not reproducible with the same seed"
Tensor.manual_seed(3421)
jf3 = TinyJit(f)
res3 = set()
for _ in range(5):
o1, o2 = jf3(a, b)
res3.add(o1.numpy()[0][0])
res3.add(o2.numpy()[0][0])
assert len(res3) == 10, "All values should be different, rand works in jit."
assert res3 != res2, "Jit rand is diff with diff seeds"
def test_jit_random_after_unrealized_random(self):
@TinyJit
def f(): return Tensor.rand()
Tensor.manual_seed(1234)
Tensor.rand()
res = [f().numpy() for _ in range(3)]
assert res[1] != res[2]
def test_jit_realization_and_sampling(self):
w = Tensor.eye(5)
@TinyJit
def foo (x): return w.dot(x).realize()
arg = [
Tensor([1,2,3,4,5]),
Tensor([1,3,3,4,6]),
Tensor([1,2,5,4,7]),
Tensor([0,2,3,1,0]),
]
Y = [foo(e).numpy() for e in arg]
foo(Tensor([7,7,7,7,7]))
want = [[1., 2., 3., 4., 5.],
[1., 3., 3., 4., 6.],
[1., 2., 5., 4., 7.],
[0., 2., 3., 1., 0.]]
np.testing.assert_allclose(want, Y)
def test_jit_buffer_behavior(self):
@TinyJit
def foo(x) -> Tensor: return x.sum().realize()
result_1 = foo(Tensor([1] * 2))
result_2 = foo(Tensor([2] * 2))
result_3 = foo(Tensor([3] * 2))
# expect the buffer to share underlying buffer
np.testing.assert_allclose(result_1.numpy(), [2], atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(result_2.numpy(), [6], atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(result_3.numpy(), [6], atol=1e-4, rtol=1e-5)
def test_jit_batch_split(self):
if Device[Device.DEFAULT].graph is None or JIT >= 2: raise unittest.SkipTest("only test graphs")
# Create long jit with 83 kernels.
def f(a, b, c, d, e):
for _ in range(80):
a = (a+b).realize()
y = (a*c).realize()
z = (y*d).realize()
w = (z*e)
return w.realize()
a = Tensor.randn(10, 10).realize()
b = Tensor.randn(10, 10).realize()
c = Tensor.randn(10, 10).realize()
d = Tensor.randn(10, 10).realize()
e = Tensor.randn(10, 10).realize()
jf = TinyJit(f)
prev = None
for _ in range(5):
o = jf(a, b, c, d, e).numpy()
if prev is not None: np.testing.assert_allclose(o, prev, atol=1e-4, rtol=1e-5)
prev = o
# Checking that 2 graphs are inited.
assert len(jf.captured.linear.src) == 2
for si in jf.captured.linear.src:
assert call_is_graph(si)
def test_jitted_clone(self):
def f(a): return a.clone().realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
ja = jf(a)
np.testing.assert_allclose(a.numpy(), ja.numpy(), atol=1e-4, rtol=1e-5)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_jitted_transfers(self):
d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
def f(a, b):
x = a.to(d1)
y = b.to(d1)
return x.realize(), y.realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 10, device=d0).realize()
b = Tensor.randn(10, 10, device=d0).realize()
xc, yc = jf(a, b)
np.testing.assert_allclose(a.numpy(), xc.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(b.numpy(), yc.numpy(), atol=1e-4, rtol=1e-5)
def test_jit_several_devs(self):
d0, d1 = f"{Device.DEFAULT}:0", "CPU"
def f(a, b):
x = a.to(d0).realize()
y = b.to(d0).realize()
return x+y.realize(), x*y.realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 10, device=d1).realize()
b = Tensor.randn(10, 10, device=d1).realize()
zc, wc = jf(a, b)
np.testing.assert_allclose((a.numpy()+b.numpy()), zc.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose((a.numpy()*b.numpy()), wc.numpy(), atol=1e-4, rtol=1e-5)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_jitted_view(self):
d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
def f(a):
x1 = a.sum(axis=(1,))
x = (x1 + 5).bitcast(dtypes.int32)
y = x.to(d1)
return y.realize()
jf = TinyJit(f)
for _ in range(5):
a = Tensor.randn(10, 1000, device=d0).realize()
xc = jf(a)
np.testing.assert_allclose((a.numpy().sum(axis=(1,)) + 5).view(np.int32), xc.numpy(), atol=1e-4, rtol=5e-5)
def test_jit_output_clone(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
f(Tensor([0.0]))
f(Tensor([0.0]))
a = f(Tensor([1.0])).clone().realize()
b = f(Tensor([2.0]))
assert abs((a - b).item()) > 0.5
def test_jit_init_empty(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
f(Tensor.empty(1))
f(Tensor.empty(1))
# scalar const input is not allowed
with self.assertRaises(JitError):
f(Tensor(2.0)).item()
# self.assertEqual(f(Tensor([2.0])).item(), 1.0) # TODO: wrong output, should be 3.0. currently depends on empty value
def test_jit_const_input(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
with self.assertRaises(JitError):
f(Tensor(UOp.const(dtypes.float, 2.0))).item()
def test_jit_deviceless_compute_input(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
with self.assertRaises(JitError):
f(Tensor(UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item()
def test_jit_init_empty_alt(self):
@TinyJit
def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1)
for i in range(4):
a = Tensor([i])
b = Tensor.empty_like(a)
c = f(a, b)
self.assertEqual(c.item(), i+1)
@unittest.skip("Pending multioutput implementation #3607")
class TestMultioutputJit(unittest.TestCase):
def _test(self, f):
for _ in range(5):
a, b = Tensor.randn(10, 10), Tensor.randn(10, 10)
out0, out1, out2 = f(a, b)
np.testing.assert_allclose(out0.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(out1.numpy(), a.numpy()-b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(out2.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
def test_jit_multioutput_realize(self):
@TinyJit
def fxn(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize()
self._test(fxn)
assert_jit_cache_len(fxn, 3)
def test_jit_multioutput_norealize(self):
@TinyJit
def fxn(a, b): return a+b, a-b, a*b
self._test(fxn)
assert_jit_cache_len(fxn, 1)
def test_jit_multioutput_mix(self):
@TinyJit
def fxn(a, b): return a+b, a-b, (a*b).realize()
self._test(fxn)
assert_jit_cache_len(fxn, 2)
class TestJitInsideJit(unittest.TestCase):
def test_jit_jit_error(self):
@TinyJit
def f(t): return t + 1
@TinyJit
def g(t): return f(t) * 3
# NOTE: first does not raise
g(Tensor([1])).realize()
with self.assertRaisesRegex(RuntimeError, "having TinyJit inside another TinyJit is not supported"):
g(Tensor([1])).realize()
class TestCopyInsideJit(unittest.TestCase):
def test_copy_inside_jit(self):
@TinyJit
def add(x,y) -> Tensor: return x.to(Device.DEFAULT)+y
for _ in range(5):
# create a Tensor on CPU
a = Tensor.rand(16,16,device="CPU").realize()
b = Tensor.rand(16,16).realize()
out = add(a,b)
np.testing.assert_allclose(out.flatten().tolist(), [x+y for x,y in zip(a.flatten().tolist(), b.flatten().tolist())])
class TestJitPrune(unittest.TestCase):
def test_simple_prune(self):
weights = Tensor.rand(16).realize()
def w2(x) -> Tensor: return (weights*2).contiguous() + x
w2_noprune = TinyJit(w2)
w2_prune = TinyJit(w2, prune=True)
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_noprune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
assert_jit_cache_len(w2_noprune, 2)
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_prune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
assert_jit_cache_len(w2_prune, 1)
def test_prune_w_copy_correct(self):
weights = Tensor.rand(16).realize()
def w2(x) -> Tensor: return (weights*2).contiguous() + x.to(Device.DEFAULT)
w2_noprune = TinyJit(w2)
w2_prune = TinyJit(w2, prune=True)
for _ in range(3):
a = Tensor.rand(16, device="CPU").realize()
out = w2_noprune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
for _ in range(3):
a = Tensor.rand(16, device="CPU").realize()
out = w2_prune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
def test_prune_w_independent_copy_correct(self):
weights = Tensor.rand(16, device="CPU").realize()
def w2(x) -> Tensor: return (weights*2).contiguous().to(Device.DEFAULT) + x
w2_noprune = TinyJit(w2)
w2_prune = TinyJit(w2, prune=True)
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_noprune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_prune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
assert_jit_cache_len(w2_prune, 1)
class TestJitFree(unittest.TestCase):
def test_free_intermediates(self):
ext_tensor = Tensor([1,24,23,45,1])
@TinyJit
def fxn(x:Tensor):
t1 = (x * 2).contiguous().realize()
t2 = (t1 + ext_tensor).contiguous().realize()
out = (t2.sum()).contiguous().realize()
return out
for i in range(5):
out = fxn(inp:=Tensor([i,1,2,3,4]))
self.assertEqual(out.item(), 114+2*i)
pre_free = GlobalCounters.mem_used
fxn.captured.free_intermediates()
savings_after_free = pre_free - GlobalCounters.mem_used
expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out
self.assertGreaterEqual(savings_after_free, expected_savings)
out = fxn(Tensor([11,1,2,3,4]))
self.assertEqual(out.item(), 136)
# Try one more time...
pre_free = GlobalCounters.mem_used
fxn.captured.free_intermediates()
fxn.captured.free_intermediates() # 2nd time to validate
savings_after_free = pre_free - GlobalCounters.mem_used
self.assertGreaterEqual(savings_after_free, expected_savings)
out = fxn(Tensor([11,1,2,3,4]))
self.assertEqual(out.item(), 136)
def test_updated_not_freed(self):
x = Tensor([1]).realize()
@TinyJit
def fxn(y):
nonlocal x
x += y
return x
for _ in range(5): fxn(Tensor([1]))
self.assertEqual(x.item(), 6)
pre_free = GlobalCounters.mem_used
fxn.captured.free_intermediates()
savings_after_free = pre_free - GlobalCounters.mem_used
self.assertEqual(savings_after_free, 0)
fxn(Tensor([2]))
self.assertEqual(x.item(), 8)
class TestJitGraphSplit(unittest.TestCase):
def compute(self, device, inp):
assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
return (inp + 1.0).contiguous().realize()
def copy(self, device, to_device, inp):
assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
return inp.to(to_device).realize()
def expect(self, f, *args, graph=None, multigraph=None, hcqgraph=None):
def _numpies(tpl): return tpl.numpy() if tpl.__class__ is Tensor else tuple([t.numpy() for t in tpl])
expected = _numpies(f(*args))
for i in range(4):
res = _numpies(f(*args))
np.testing.assert_allclose(res, expected, atol=1e-4, rtol=1e-5)
dev = Device[Device.DEFAULT]
graph_t = graph_class(dev)
if graph_t is None: return
got = f.captured.linear.src
from tinygrad.runtime.graph.hcq import HCQGraph
from tinygrad.engine.jit import MultiGraphRunner
if graph_t is HCQGraph:
validate = hcqgraph
elif issubclass(graph_t, MultiGraphRunner):
validate = multigraph
else:
validate = graph
assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
for expected, si in zip(validate, got):
ast = si.src[0]
if expected["type"] == "graph":
assert call_is_graph(si), f"Expected graph, got {ast.op}"
inner_cnt = len(ast.src[0].src)
assert inner_cnt == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {inner_cnt}"
elif expected["type"] == "comp":
assert ast.op in (Ops.SINK, Ops.PROGRAM), f"Expected kernel, got {ast.op}"
elif expected["type"] in ("copy", "xfer"):
assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"
def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
def ji_comp(self): return {"type": "comp"}
def ji_copy(self): return {"type": "copy"}
def ji_xfer(self): return {"type": "xfer"}
def test_jit_split_simple(self):
@TinyJit
def f(inp):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute(Device.DEFAULT, op1)
return op2
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
self.expect(f, inp,
graph=[self.ji_graph(3)],
multigraph=[self.ji_graph(3)],
hcqgraph=[self.ji_graph(3)])
def test_jit_cpu_simple(self):
if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@TinyJit
def f(inp, inp_cpu):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute("CPU", inp_cpu)
op3 = self.compute(Device.DEFAULT, op1)
return op2, op3
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
self.expect(f, inp, inp_cpu,
graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
hcqgraph=[self.ji_graph(4)])
def test_jit_cpu_several(self):
if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@TinyJit
def f(inp, inp_cpu):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute("CPU", inp_cpu)
op3 = self.compute("CPU", op2)
op4 = self.compute(Device.DEFAULT, op1)
return op3, op4
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
self.expect(f, inp, inp_cpu,
graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
hcqgraph=[self.ji_graph(5)])
def test_jit_multidev(self):
if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
try: Device[f"{Device.DEFAULT}:1"]
except Exception: raise unittest.SkipTest("no multidevice")
@TinyJit
def f(inp, inp_d1):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
op3 = self.compute(f"{Device.DEFAULT}:1", op2)
op4 = self.compute(Device.DEFAULT, op1)
return op3, op4
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
self.expect(f, inp, inp_d1,
graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
multigraph=[self.ji_graph(5)],
hcqgraph=[self.ji_graph(5)])
def test_jit_multidev_xfer(self):
if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU is not a valid default device for this test (zero-copies)")
if Device.DEFAULT == "METAL": raise unittest.SkipTest("Metal is flaky, with multidevice (same as metal llama 4gpu?)")
try: Device[f"{Device.DEFAULT}:1"]
except Exception: raise unittest.SkipTest("no multidevice")
@TinyJit
def f(inp, inp_d1):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
op3 = self.copy(f"{Device.DEFAULT}:1", Device.DEFAULT, op2)
op4 = self.compute(f"{Device.DEFAULT}:1", op2)
op5 = self.compute(Device.DEFAULT, op3)
return op1, op4, op5
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
self.expect(f, inp, inp_d1,
graph=[self.ji_graph(2), self.ji_comp(), self.ji_xfer(), self.ji_comp(), self.ji_comp()],
multigraph=[self.ji_graph(6)],
hcqgraph=[self.ji_graph(6)])
@unittest.skip("this fails if you don't have SDMA or are using AMD_DISABLE_SDMA=1")
@unittest.skipIf(DEV.interface.startswith("MOCK"), "MockGPU does not support parallel copies")
def test_jit_multidev_copy(self):
if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
@TinyJit
def f(inp):
op0 = self.compute(Device.DEFAULT, inp)
op1 = self.compute(Device.DEFAULT, op0)
op2 = self.copy(Device.DEFAULT, "CPU", op1)
op3 = self.compute("CPU", op2)
return op3
inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
self.expect(f, inp,
graph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
hcqgraph=[self.ji_graph(4)])
class TestJitRandom(unittest.TestCase):
def test_jit_rangeify(self):
tst = {0:[], 1:[]}
for r in [0,1]:
Tensor.manual_seed(1337)
with Context(JIT=r):
_ = Tensor.randint(4, high=3)
# this second one makes the behavior different
_ = Tensor.randint(4, high=3)
@TinyJit
def f(): return Tensor.randint(20, high=5)
for _ in range(5): tst[r].append(f().tolist())
for i, (t0, t1) in enumerate(zip(tst[0], tst[1])):
self.assertListEqual(t0, t1, msg=f"mismatch at list {i}")
if __name__ == '__main__':
unittest.main()

78
tinygrad_repo/test/backend/test_jit_cases.py Normal file
View File

@@ -0,0 +1,78 @@
import unittest
from tinygrad import TinyJit, Tensor
# The JIT functions as a "capturing" JIT.
# Whatever kernels ran in the JIT the second run through the function will be the kernels that will run from then on.
# Explicit inputs to the function are updated in the JIT graph to the new inputs.
# JITs have four tensor types
# 1. Tensors that are explicit in the input, aka what's passed in. TODO: support lists/dicts/classes, anything get_state works on
# 2. Tensors that are explicit in the output, aka what's returned. TODO: same as above
# 3. Tensors that are implicit in the input as a closure.
# 4. Tensors that are implicit in the output because they were assigned to and realized.
# explicit inputs and outputs are realized on their way in and out of the JIT
# there's a whole bunch of edge cases and weirdness here that needs to be tested and clarified.
class TestJitCases(unittest.TestCase):
def test_explicit(self):
# this function has an explicit input and an explicit output
@TinyJit
def f(x:Tensor):
ret:Tensor = x*2
return ret
for i in range(5):
out = f(Tensor([i]))
self.assertEqual(out.item(), i*2)
def test_implicit_input(self):
# x is the implicit input (like a weight)
x = Tensor([0])
# this function has an implicit input and an explicit output
@TinyJit
def f():
ret:Tensor = x*2
return ret
for i in range(5):
# NOTE: this must be realized here, otherwise the update doesn't happen
# if we were explicitly tracking the implicit input Tensors, we might not need this realize
x.assign(Tensor([i])).realize()
out = f()
self.assertEqual(out.item(), i*2)
def test_implicit_output(self):
# out is the implicit output (it's assigned to)
out = Tensor([0])
# this function has an explicit input and an implicit output
@TinyJit
def f(x:Tensor):
# NOTE: this must be realized here
# if we were explicitly tracking the implicit output Tensors, we might not need this realize
out.assign(x*2).realize()
for i in range(5):
f(Tensor([i]))
self.assertEqual(out.item(), i*2)
def test_implicit_io(self):
# x is the implicit input (like a weight)
# out is the implicit output (it's assigned to)
x = Tensor([0])
out = Tensor([0])
# this function has an implicit input and an implicit output
@TinyJit
def f():
out.assign(x*2).realize() # NOTE: this must be realized here
for i in range(5):
x.assign(Tensor([i])).realize()
f()
self.assertEqual(out.item(), i*2)
if __name__ == '__main__':
unittest.main()

414
tinygrad_repo/test/backend/test_jit_footguns.py Normal file
View File

@@ -0,0 +1,414 @@
#!/usr/bin/env python
"""
JIT Footguns: Documenting unexpected behavior changes when using @TinyJit
Each test shows behavior that works without JIT but changes with JIT.
Comments marked "should be X!" indicate the intuitively expected value.
SILENT MISMATCHES (highest priority - wrong results, no error):
class_method_shared_across_instances EASY could check if first arg is self and warn
slice_assign_requires_realize MED assign graph not connected to read during JIT replay
output_buffer_reuse MED performance tradeoff, could add option or better docs
symbolic_pad_view_frozen MED pad view BIND values baked in at capture time
python_constants_frozen HARD inherent to tracing JITs
conditional_branches_frozen HARD inherent to tracing JITs
ERRORS RAISED (lower priority - at least users know):
item_bakes_in_values EASY raises JitError if .item()/.data() accessed during capture
unrealized_const_input_error EASY raises JitError for unrealized const inputs
non_tensor_outputs_error EASY raises JitError if return contains non-Tensor values
positional_kwargs_cannot_mix EASY normalize positional args to kwargs using function signature
duplicate_inputs_fail MED would need to handle aliasing in input_replace
nested_jit_fails_on_second_call MED could fail on first call instead of second
"""
import unittest
import numpy as np
from tinygrad import Tensor, TinyJit, Device
from tinygrad.engine.jit import JitError
from tinygrad.helpers import JIT
class TestJitFootguns(unittest.TestCase):
def test_output_buffer_reuse(self):
"""Output tensors share buffer after capture - old references get overwritten."""
@TinyJit
def f(x): return x.sum().realize()
r1 = f(Tensor([1, 1])) # warmup
r2 = f(Tensor([2, 2])) # capture
r3 = f(Tensor([3, 3])) # jit exec
self.assertEqual(r1.item(), 2) # warmup result independent
self.assertEqual(r3.item(), 6) # latest is correct
self.assertEqual(r2.item(), 6) # should be 4! (overwritten by r3)
def test_output_buffer_workaround(self):
"""Use .clone().realize() to get independent copies."""
@TinyJit
def f(x): return x.sum().realize()
r1 = f(Tensor([1, 1])).clone().realize()
r2 = f(Tensor([2, 2])).clone().realize()
r3 = f(Tensor([3, 3])).clone().realize()
self.assertEqual([r1.item(), r2.item(), r3.item()], [2, 4, 6])
def test_graph_input_output_aliasing(self):
"""Test that JIT handles input=output aliasing correctly, simulating LLM generate pattern.
The LLM generate pattern:
1. First "session": multiple iterations where output becomes next input
2. Second "session": starts with a NEW input tensor (not the previous output)
The bug: GraphRunner computes input_replace during _first_run. If at that time input buffer == output buffer
(aliasing), it incorrectly includes the output position in input_replace. Later, when a DIFFERENT input
is passed, the output position gets overwritten with the input, corrupting the computation.
This requires multiple kernels to trigger because single-kernel JITs don't get graphed ("only one kernel doesn't graph").
"""
if Device[Device.DEFAULT].graph is None or JIT != 1:
self.skipTest("test requires JIT graph support")
# Multiple operations to create multiple kernels that get batched into a GraphRunner
@TinyJit
def step(x):
y = (x + 1).realize() # kernel 1
z = (y * 2).realize() # kernel 2
return z
# Phase 1: warmup and capture
a = Tensor([10]).contiguous().realize()
step(a) # warmup (cnt=0)
b = Tensor([20]).contiguous().realize()
x = step(b) # capture (cnt=1), x = (20+1)*2 = 42
# Phase 2: first "session" - iterations where output becomes input (triggers _first_run with aliasing)
for _ in range(3):
x = step(x) # (42+1)*2=86, (86+1)*2=174, (174+1)*2=350
self.assertEqual(x.item(), 350)
# Phase 3: second "session" - NEW input tensor (simulates new generate() call)
# The bug: GraphRunner's input_replace incorrectly includes the output position
# When new input y is passed, it overwrites the output buffer, using old value (350) instead of new (100)
y = Tensor([100]).contiguous().realize()
for _ in range(3):
y = step(y) # should be (100+1)*2=202, (202+1)*2=406, (406+1)*2=814
self.assertEqual(y.item(), 814) # fails with 1406 if bug exists (uses 350 instead of 100)
def test_multiple_outputs_same_intermediate(self):
"""Multiple outputs derived from the same intermediate - JIT copies aliased inputs to prevent hazard."""
@TinyJit
def f(buf, frame):
new_buf = buf[1:].cat(frame, dim=0)
return new_buf.contiguous(), new_buf[:1].contiguous()
buf = Tensor([[0], [1], [2]]).contiguous().realize()
for i in range(4):
frame = Tensor([[10+i]]).contiguous().realize()
expected_first = buf[1:2].numpy().item()
new_buf, first = f(buf, frame)
self.assertEqual(first.numpy().item(), expected_first)
buf = new_buf
def test_intra_kernel_output_input_aliasing(self):
"""JIT must copy aliased input when output buffer is fed back as input (read-write race in same kernel)."""
N = 1 << 20
f = TinyJit(lambda buf, new: buf[N//2:].cat(new), prune=True)
buf = Tensor.zeros(N, dtype='int32').contiguous().realize()
for i in range(10):
buf = f(buf, Tensor(np.ones(N//2, dtype=np.int32)*(i+1)))
np.testing.assert_array_equal(buf[:N//2].numpy(), np.full(N//2, i, dtype=np.int32))
def test_slice_assign_works_without_realize(self):
"""Slice assign then read from same buffer - pending assigns are side-realized."""
from tinygrad import Variable
v_pos = Variable("pos", 0, 3)
cache = Tensor.zeros(4, 4).contiguous().realize()
@TinyJit
def f(pos):
cache[pos:pos+1, :].assign(Tensor.ones(1, 4))
return cache.sum().realize()
for i in range(4):
cache.assign(Tensor.zeros(4, 4)).realize()
self.assertEqual(f(v_pos.bind(i)).item(), 4.0)
def test_symbolic_pad_view_frozen(self):
"""Symbolic pad view has BIND values baked in at capture time. TODO: pad should be captured in jit."""
from tinygrad import Variable
a = Tensor.rand(3, 10).realize()
# broken: pad is a view, BIND values frozen at capture (i=2)
@TinyJit
def f_broken(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
for i in range(1, 5): f_broken(a[:, :Variable("i", 1, 10).bind(i)])
self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2) # should be 4!
# workaround: contiguous fuses pad into kernel
@TinyJit
def f_fixed(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
for i in range(1, 5): f_fixed(a[:, :Variable("i", 1, 10).bind(i)])
self.assertEqual(int((f_fixed(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 4)
def test_non_tensor_outputs_error(self):
@TinyJit
def f(x, mult): return (x * 2).realize(), mult * 10
with self.assertRaises(JitError):
for i in range(3): f(Tensor([i]), i)
def test_duplicate_inputs_fail(self):
"""JIT cannot handle the same tensor passed as multiple arguments."""
@TinyJit
def f(a, b): return (a + b).realize()
x = Tensor([1, 2, 3])
with self.assertRaises(JitError):
f(x, x)
def test_tensors_in_containers(self):
@TinyJit
def f(a, arr): return (a + arr[0]).realize()
for i in range(4):
a, b = Tensor([1, 1, 1]).realize(), Tensor([i, i, i]).realize()
np.testing.assert_array_equal(f(a, [b]).numpy(), [1+i, 1+i, 1+i])
def test_nested_jit_fails_on_second_call(self):
"""Nested JIT works on first call but fails on second."""
@TinyJit
def inner(t): return t + 1
@TinyJit
def outer(t): return inner(t) * 3
self.assertEqual(outer(Tensor([1])).realize().item(), 6) # works!
with self.assertRaises(RuntimeError):
outer(Tensor([2])).realize() # fails
def test_implicit_inputs_need_realize(self):
"""Closure tensors must be realized before JIT call."""
x = Tensor([0])
@TinyJit
def f(): return (x * 2).realize()
for i in range(5):
x.assign(Tensor([i])).realize() # must realize!
self.assertEqual(f().item(), i * 2)
def test_views_with_different_offsets_fail(self):
"""JIT requires consistent tensor views across calls."""
@TinyJit
def f(a): return (a + 1).realize()
base = Tensor.randn(10, 10).realize()
with self.assertRaises(JitError):
for i in range(1, 5):
f(base[:, i:i+2]) # different offset each time
def test_shape_change_after_capture_fails(self):
"""Shapes are locked at capture time."""
@TinyJit
def f(a, b): return (a + b).realize()
f(Tensor.randn(10, 10), Tensor.randn(10, 10)) # warmup
f(Tensor.randn(10, 10), Tensor.randn(10, 10)) # capture
with self.assertRaises(JitError):
f(Tensor.randn(20, 20), Tensor.randn(20, 20))
def test_python_constants_frozen(self):
"""Python variables inside JIT use capture-time values."""
mult = 1
@TinyJit
def f(x): return (x * mult).realize()
results = []
for i in range(5):
mult = i + 1
results.append(f(Tensor([10])).item())
self.assertEqual(results[0], 10) # warmup, mult=1
self.assertEqual(results[1], 20) # capture, mult=2
self.assertEqual(results[2], 20) # should be 30!
self.assertEqual(results[3], 20) # should be 40!
def test_unrealized_const_input_error(self):
"""Const tensors have no buffer to replace, so JIT raises an error. Even explicit .realize() doesn't help."""
@TinyJit
def f(a, b): return (a * b).realize()
# unrealized const fails
with self.assertRaises(JitError):
f(Tensor([1, 2, 3]).realize(), Tensor(2))
# explicit .realize() on const still fails - const cannot be realized to have a buffer
@TinyJit
def g(a, b): return (a * b).realize()
with self.assertRaises(JitError):
g(Tensor([1, 2, 3]).realize(), Tensor(2).realize())
def test_conditional_branches_frozen(self):
"""Only the branch taken during capture runs thereafter."""
@TinyJit
def f(x, use_square):
if use_square:
return (x * x).realize()
return (x * 2).realize()
f(Tensor([3]), True) # warmup
f(Tensor([3]), False) # capture (False branch)
result = f(Tensor([3]), True) # passing True but False branch runs
self.assertEqual(result.item(), 6) # should be 9!
def test_positional_kwargs_cannot_mix(self):
"""Must use same calling convention after capture."""
@TinyJit
def f(a, b): return (a + b).realize()
f(Tensor([1]), Tensor([2])) # warmup with positional
f(Tensor([1]), Tensor([2])) # capture with positional
with self.assertRaises(JitError):
f(a=Tensor([3]), b=Tensor([4])) # kwargs fail
def test_class_method_shared_across_instances(self):
"""JIT on instance methods is shared at class level."""
class Model:
def __init__(self, scale):
self.scale = Tensor([scale])
@TinyJit
def forward(self, x):
return (x * self.scale).realize()
m1, m2 = Model(2), Model(3)
m1.forward(Tensor([5])) # warmup
m1.forward(Tensor([5])) # capture with m1.scale=2
self.assertEqual(m1.forward(Tensor([5])).item(), 10)
self.assertEqual(m2.forward(Tensor([5])).item(), 10) # should be 15!
def test_side_effects_only_during_capture(self):
"""Function body not executed during JIT replay."""
call_count = [0]
@TinyJit
def f(x):
call_count[0] += 1
return (x * 2).realize()
f(Tensor([1])) # warmup
f(Tensor([2])) # capture
self.assertEqual(call_count[0], 2)
f(Tensor([3]))
f(Tensor([4]))
f(Tensor([5]))
self.assertEqual(call_count[0], 2) # still 2, not 5!
def test_nothing_realized_fails(self):
"""Must JIT at least one kernel."""
@TinyJit
def f(a, b): return None
with self.assertRaises(JitError):
for _ in range(3):
f(Tensor([1]), Tensor([2]))
def test_item_creates_unrealized_return(self):
""".item() in shape computation raises error during JIT capture."""
@TinyJit
def f(x): return Tensor.zeros(x.sum().item())
f(Tensor([1, 1, 1])) # warmup
with self.assertRaises(JitError):
f(Tensor([1, 1, 1])) # capture - .item() raises
def test_item_bakes_in_values(self):
""".item() during JIT capture raises error (would bake in value)."""
@TinyJit
def f(x, mask): return x.masked_select(mask)
f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])) # warmup
with self.assertRaises(JitError):
f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])) # capture - .item() raises
def test_masked_select_static_size_jittable(self):
@TinyJit
def f(x, mask): return x.masked_select(mask, size=4, fill_value=-1).realize()
for _ in range(3):
np.testing.assert_equal(f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])).numpy(), [1, 3, -1, -1])
np.testing.assert_equal(f(Tensor([5, 6, 7, 8]), Tensor([False, True, True, True])).numpy(), [6, 7, 8, -1])
np.testing.assert_equal(f(Tensor([9, 8, 7, 6]), Tensor([True, True, True, True])).numpy(), [9, 8, 7, 6])
np.testing.assert_equal(f(Tensor([1, 1, 1, 1]), Tensor([False, False, False, False])).numpy(), [-1, -1, -1, -1])
def test_nonzero_static_size_jittable(self):
@TinyJit
def f(x): return x.nonzero(size=3, fill_value=-1).realize()
for _ in range(3):
np.testing.assert_equal(f(Tensor([1, 0, 2, 0, 3])).numpy(), [[0], [2], [4]])
np.testing.assert_equal(f(Tensor([0, 0, 5, 0, 0])).numpy(), [[2], [-1], [-1]])
np.testing.assert_equal(f(Tensor([0, 0, 0, 0, 0])).numpy(), [[-1], [-1], [-1]])
def test_tolist_bakes_in_values(self):
""".tolist() raises error during JIT capture (would bake in values)."""
@TinyJit
def f(x): return Tensor(x.tolist())
f(Tensor([1, 2, 3])) # warmup
with self.assertRaises(JitError):
f(Tensor([1, 2, 3])) # capture - .tolist() raises
class TestJitCorrectBehavior(unittest.TestCase):
"""Behaviors that work correctly - documented for clarity."""
def test_random_regenerates(self):
"""Random tensors regenerate each call."""
@TinyJit
def f(x):
return (x + Tensor.rand(3)).realize()
f(Tensor([0, 0, 0])) # warmup
f(Tensor([0, 0, 0])) # capture
results = {tuple(f(Tensor([0, 0, 0])).numpy().tolist()) for _ in range(5)}
self.assertEqual(len(results), 5)
def test_unrealized_return_auto_realized(self):
"""Unrealized return tensors are auto-realized."""
@TinyJit
def f(a, b): return a + b # no explicit realize
for _ in range(5):
a, b = Tensor.randn(10), Tensor.randn(10)
np.testing.assert_allclose(f(a, b).numpy(), a.numpy() + b.numpy(), atol=1e-5)
def test_kwargs_order_doesnt_matter(self):
"""Kwargs are sorted by name, so order doesn't matter."""
@TinyJit
def f(first, second): return (first / second).realize()
for _ in range(3):
a, b = Tensor.randn(10), Tensor.randn(10) + 1
np.testing.assert_allclose(f(second=b, first=a).numpy(), a.numpy() / b.numpy(), atol=1e-4)
np.testing.assert_allclose(f(first=a, second=b).numpy(), a.numpy() / b.numpy(), atol=1e-4)
def test_input_mutation_consistent(self):
"""Input mutation via assign works consistently."""
@TinyJit
def f(x):
x += 1
x.realize()
return x
a = Tensor([0]).contiguous().realize()
for _ in range(5):
f(a)
self.assertEqual(a.item(), 5)
if __name__ == '__main__':
unittest.main()

29
tinygrad_repo/test/backend/test_kernel_cache.py Normal file
View File

@@ -0,0 +1,29 @@
#!/usr/bin/env python
import unittest
from tinygrad.tensor import Tensor
from tinygrad import Device
class TestKernelCache(unittest.TestCase):
def test_kernel_cache_in_action(self):
if Device.DEFAULT not in ["CPU"]:
self.skipTest("No custom kernel cache is implemented")
const_value = 0.6765677269
a = Tensor.rand(4,4).realize()
b = Tensor.rand(4,4).realize()
x = a + b + const_value
x.realize()
a1 = Tensor.rand(4,4).realize()
b1 = Tensor.rand(4,4).realize()
orig_compile_func = Device['CPU'].compiler.compile_cached
Device['CPU'].compiler.compile_cached = None # making it not callable
try:
x1 = a1 + b1 + const_value
x1.realize() # Same kernel should be from cache.
finally:
Device['CPU'].compiler.compile_cached = orig_compile_func
if __name__ == "__main__":
unittest.main()

459
tinygrad_repo/test/backend/test_linearizer.py Normal file
View File

@@ -0,0 +1,459 @@
import numpy as np
import unittest
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers
from tinygrad.device import Device, Buffer
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import CUDARenderer
from tinygrad.renderer.isa import ISARenderer
from test.helpers import replace_opts
MOCKGPU = DEV.interface.startswith("MOCK")
from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
class TestLinearizer(unittest.TestCase):
def test_arg_dedup(self):
# NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
# without contiguous folding, rand.to("CPU") and rand.contiguous().to("CPU") are different UOps.
# this test asserts they are the identical Buffer
# having different buffers is fine for correctness, because the outputs match.
a, b = Tensor.randn(4).realize(), Tensor.randn(4).realize()
np_a, np_b = a.numpy(), b.numpy()
c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
linear = c.schedule_linear()
run_linear(linear)
rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
def test_load_removed(self):
a = Tensor.rand(1).realize()
b = Tensor.rand(1).realize()
ta = Tensor.where(Tensor(True), a, b).numpy()
tb = Tensor.where(Tensor(False), a, b).numpy()
np.testing.assert_equal(a.numpy(), ta)
np.testing.assert_equal(b.numpy(), tb)
@unittest.skip("TODO: some backends insert more casts")
def test_cast_there_and_back(self):
tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
ast = helper_linearizer_opt(out)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
@unittest.expectedFailure
def test_cast_back_and_there(self):
tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
ast = helper_linearizer_opt(out)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
def test_late_bias_load(self):
img = Tensor.empty(1, 3, 16, 16)
w = Tensor.empty(16, 3, 3, 3)
b = Tensor.empty(16)
out = img.conv2d(w, b)
ast = helper_linearizer_opt(out)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
# slice at the last loop end
uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
# only valid test if outermost range is the reduce
if uops[uslice].src[-1].arg[-1] == AxisType.REDUCE:
load_types = [u.src[0].dtype for u in uops[uslice+1:] if u.op == Ops.LOAD]
# assert that there is a global load after the reduce ends
assert any(dt.addrspace == AddrSpace.GLOBAL for dt in load_types)
def _test_no_nested_ranges(self, lins, skip=None):
for l in lins:
range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG])
ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.END and u.src[0] in range_in_acc)]
for i,u in enumerate(ranges):
if skip and i in skip: continue
assert ranges[i-1] != u, f"multireduce nested the ranges! {ranges[i-1], {u}}"
def test_two_nested_range(self):
a = Tensor.randn(2, ).realize()
out = a.reshape(2, 1).expand(2, 3).sum()
ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
def test_three_nested_range(self):
a = Tensor.randn(2, ).realize()
out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
def test_two_nested_range_alt_indexing(self):
a = Tensor([2, 2]).realize()
out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
ast = helper_linearizer_opt(out, wanna_output=[24])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
# RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
# the index of the load doesnt depend on the second range
assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])
def test_range_outer_op_before_phi(self):
a = Tensor.randn(4, 1).realize()
b = Tensor.randn(1, 1).realize()
out = (a + b[0]).sum() + b[0]
ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
# LOAD -> RANGE -> LOAD -> STORE
assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
def test_range_outer_op_before_phi_nested_range(self):
a = Tensor.randn(2, ).realize()
b = Tensor.randn(1, 1).realize()
out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
def test_load_dedup(self):
# for different leaves in the AST, the same loads may occur.
a = Tensor.randn(4).realize()
# these are of size 3 to avoid float4 coalesce
r = a[:-1] + a[1:]
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).src[2].src)
num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
assert num_loads <= 4, "more load uops than needed"
assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@unittest.skip("this is handled at higher level now")
def test_upcast_cse(self):
# when upcasting, within a subtree, there may be common expressions.
a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
r = a.expand([2]) + b.expand([2])
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).src[2].src)
num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
assert num_ops <= 1, "more alu uops than needed"
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_reduce_upcast(self):
x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
r = Tensor.conv2d(x,w,padding=1).relu()
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
accs = [u for u in uops if u.op is Ops.DEFINE_REG]
stores = [u for u in uops if u.op is Ops.STORE]
assert len(accs) == 0 # it's removed now
assert len(stores) == 1
assert stores[0].src[1].dtype == dtypes.float.vec(4)
# NOTE: can reenable, it does work. it just makes BEAM slow
@unittest.expectedFailure
@unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
def test_upcast_with_locals_cpu(self):
out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
renderer=Device[Device.DEFAULT].renderer)
self.assertEqual(len(prg.src[3].arg.split("for")), 5)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
def test_upcast_with_locals(self):
x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
r = (x@y).relu()
opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
# the first store is to lds and can be upcasted
assert stores[0].src[1].dtype == dtypes.float.vec(4)
assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
# the second store is to gds with no upcasts
assert stores[1].src[1].dtype == dtypes.float
assert any(x.op is Ops.PARAM for x in stores[1].toposort())
def test_zero_fold(self):
a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
r = Tensor.stack(a, b)
uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).src[2].src)
num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
assert num_ops == 0, "more alu uops than needed"
def test_sum_acc_dtype(self):
for tensor_dtype, acc_dtype in (
(dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):
if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts:
a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
realized_ast = a.schedule_linear().src[-1].src[0]
program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
assert local[0].dtype.base == acc_dtype
def test_arg_acc_dtype(self):
def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
realized_ast = c.schedule_linear().src[-1].src[0]
program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
self.assertEqual(local[0].dtype.base, expected_dtype)
tests = (
(dtypes.float16, None, dtypes.float),
(dtypes.bfloat16, None, dtypes.float),
(dtypes.float, None, dtypes.float),
(dtypes.float16, dtypes.float16, dtypes.float16),
(dtypes.bfloat16, dtypes.bfloat16, dtypes.bfloat16),
(dtypes.float, dtypes.float16, dtypes.float16),
)
for tensor_dtype, acc_dtype, expected_dtype in tests:
if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts and expected_dtype in dts:
a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype)
helper_arg_acc_dtype(a.sum(dtype=acc_dtype), expected_dtype)
helper_arg_acc_dtype(a.matmul(b, dtype=acc_dtype), expected_dtype)
helper_arg_acc_dtype(Tensor.einsum("ki,ij->kj", a, b, dtype=acc_dtype), expected_dtype)
d, w = Tensor.rand(4, 8, 8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, 2, 2, dtype=tensor_dtype)
helper_arg_acc_dtype(d.conv2d(w, dtype=acc_dtype), expected_dtype)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_simple_unroll_no_between_phi_dependencies(self):
x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
r = (x@y).relu()
opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
ast = helper_linearizer_opt(r, [opt])
# the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
for u in uops:
if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace is AddrSpace.REG:
if uops.index(u) < begin_range:
assert u.src[1].op is Ops.CONST
else:
assert u.src[1].op in GroupOp.ALU
assert begin_range < uops.index(u) < end_range
# children of END are placed after ENDRANGE
if any(x.op is Ops.END and x.src[1].op in GroupOp.ALU for x in u.src):
assert end_range < uops.index(u)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
def test_default_global_reversed(self):
# shrink so that the dims do not collapse
t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
ast = helper_linearizer_opt(t+1)
uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
idxs = sorted(idxs, key=lambda uop: uop.arg)
assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
assert (idxs[1].arg, idxs[1].src[0].arg) == ('gidx1', 5), idxs[1].arg
assert (idxs[2].arg, idxs[2].src[0].arg) == ('gidx2', 4), idxs[2].arg
def test_sum_collapse(self):
t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK]
# sum_collapse is a full collapse now
assert len(sched) == 1
assert not any(u.op is Ops.REDUCE and len(u.arg[1]) > 0 for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
#lin = Kernel(sched[0].ast)
#assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse"
def test_assign_fold(self):
a = Tensor.ones(4, 4).contiguous().realize()
m = Tensor.ones(4, 4).shrink(((1, 2), None)).pad(((1, 2), None))
a.assign(a+m)
a.realize()
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
@unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
def test_where_fold(self):
a = Tensor.ones(4, 4).contiguous().realize()
b = a.shrink(((1, 2), None)).pad(((1, 2), None))
a.assign(b.where(2, a))
linear, var_vals = a.linear_with_vars()
assert len(linear.src) == 1
run_linear(linear, var_vals)
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded"
def test_phi_simplification(self):
def helper(t, max_ops=0):
ast = helper_linearizer_opt(t)
uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
# ignore kernel optimized IF statements for now
if if_op:=next((u for u in uops if u.op is Ops.IF), None):
uops = uops[:uops.index(if_op)]
assert len(set([u.op for u in uops if u.op in {Ops.RANGE, Ops.SPECIAL}])) == 1, "has either specials or ranges, not both"
reg_stores = [u for u in uops if u.op is Ops.STORE and isinstance(dt:=u.src[0].dtype, PtrDType) and dt.addrspace == AddrSpace.REG]
assert len(reg_stores) == 0, "STORE to reg should have been simplified"
assert len([u for u in uops if u.op is Ops.MAX]) <= max_ops, "no unnecessary MAX ops"
helper(Tensor.arange(5.5, (3.5*300), 3.5), max_ops=2)
helper(Tensor.arange(-1, -100, -5), max_ops=2)
# NOTE: both of these split the reduce (this just wasn't tracked before)
#helper(Tensor.arange(-3.2, 6.7, 0.64), max_ops=2)
#helper(Tensor.arange(256), max_ops=2)
helper(Tensor.arange(255), max_ops=2)
@unittest.skip("test implicitly depends on certain optimizations")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
def test_grouped_store_phis(self):
"""
float4 acc0 = float4(0.0,0.0,0.0,0.0);
{
acc0 = // ...
}
*((device float4*)(data0+alu2)) = float4(acc0.x,acc0.y,acc0.z,acc0.w);
simplifies to:
*((device float4*)(data0+alu2)) = acc0;
"""
x, y = Tensor.empty(64,64), Tensor.empty(64,64)
out = x.matmul(y)
with Context(TC=0):
ast = helper_linearizer_opt(out)
uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
# check that the float4 cast collapses
store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
for val in store_vals:
assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE
@unittest.skip("test implicitly depends on certain optimizations")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_grouped_store_values(self):
x = Tensor.randn((4,3,6,6)).realize()
out = x.flip((0,1)).contiguous()
ast = helper_linearizer_opt(out)
store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0]
assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_grouped_store_locals_and_globals(self):
x, y = Tensor.empty(64, 64), Tensor.empty(64, 64)
out = x@y
opt = [Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.GROUPTOP, 0, 8),
Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
ast = helper_linearizer_opt(out, opts=[opt])
def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
barrier = [u for u in uops if u.op is Ops.BARRIER]
assert len(barrier) == 1
# check that the float4 cast collapses for all stores
for store in local_stores+global_stores:
assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
# # check the children's vins
# TODO: src ALU are not the same, should it?
# assert barrier.src == tuple(local_stores)
assert len([u for u in uops if u.op is Ops.IF])
@unittest.skip("test implicitly depends on certain optimizations")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
def test_grouped_store_local_only(self):
x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
r = (x@y).relu()
ast = helper_linearizer_opt(r)
uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
# the float4 value stores directly in lds and we skip upcast
self.assertEqual(stores[0].src[1].dtype, dtypes.float.vec(4))
#assert stores[0].src[-1].op is not Ops.VECTORIZE
# the global store doesn't change
assert stores[1].src[1].dtype == dtypes.float
# *** helpers ***
def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
if isinstance(r, Tensor): r = [r]
linear, var_vals = Tensor.linear_with_vars(*r)
run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals) # run all kernels except the last one
last_call = linear.src[-1]
ast = last_call.src[0]
assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
# now all input buffers in last_call should be realized
# create fresh buffers for the outputs
bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
# ensure buffers are allocated
for b in bufs: b.ensure_allocated()
return ast, bufs
def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
assert isinstance(ast, UOp), "ast must be UOp"
inbufs = [x.uop.base.buffer for x in inputs]
outbufs = [Buffer(inbufs[-1].device if inbufs else Device.DEFAULT, out.size, out.src[1].dtype).allocate() for out in ast.src]
_helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs)
def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs):
realized_ast, real_bufs = helper_realized_ast(r)
_helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs)
return realized_ast
def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
def reset_bufs(bufs:list[Buffer]):
for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
outbufs = real_bufs[:len(realized_ast.src)]
wanna_output = [np.array(x).flatten() for x in wanna_output]
buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
for u,b in zip(buf_uops, real_bufs): buffers[u] = b
def run_prg(opts):
ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
run_linear(UOp(Ops.LINEAR, src=(ast.call(*buf_uops),)))
def check_opt(opts):
reset_bufs(outbufs)
run_prg(opts)
for x,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(x, want, atol=atol, rtol=rtol)
# Get baseline if it is not provided, which is not optimized at all.
run_prg(opts=())
if len(wanna_output) == 0: wanna_output = copyout_outputs(outbufs)
else:
for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
# Check correctness of handcoded optimiztions.
reset_bufs(outbufs)
run_prg(opts=None)
for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
for x in opts: # Check custom transformations if any.
check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)
if __name__ == '__main__':
unittest.main()

31
tinygrad_repo/test/backend/test_linearizer_dumb.py Normal file
View File

@@ -0,0 +1,31 @@
# ruff: noqa: E501
# tests where the Linearizer is doing something dumb
# like test_linearizer_failures, but they don't have to fail
import unittest
from tinygrad import Device, dtypes
from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad.codegen import to_program
class TestLinearizerFailure(unittest.TestCase):
@unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
def test_failure_beam_mnist(self):
c0 = UOp.param(0, dtypes.uchar.ptr(4014080))
c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
c4 = UOp.param(1, dtypes.int.ptr(512))
c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
c9 = UOp.param(2, dtypes.uchar.ptr(47040000))
c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
_ = to_program(ast, Device["METAL"].renderer)
if __name__ == '__main__':
unittest.main()

96
tinygrad_repo/test/backend/test_llama_kernels.py Normal file
View File

@@ -0,0 +1,96 @@
import unittest
from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
from tinygrad.helpers import getenv
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
from extra.llama_kernels.fused_ce import fused_ce_loss
from extra.llama_kernels import local_abs_max
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
from test.helpers import needs_second_gpu
def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
Tensor.manual_seed(0)
logits_rand = Tensor.randn(bs, seqlen, vocab).cast(dtypes.bfloat16)
targets = Tensor.randint(bs, seqlen, high=vocab, dtype=dtypes.int32)
logits, logits_ref = logits_rand.clone(), logits_rand.detach().float().contiguous()
with Context(DEBUG=0):
Tensor.realize(logits, logits_ref, targets)
loss = fused_ce_loss(logits, targets, label_smoothing=label_smoothing)
loss.backward()
Tensor.realize(loss, logits.grad)
ref = logits_ref.sparse_categorical_crossentropy(targets, label_smoothing=label_smoothing)
ref.backward()
Tensor.realize(ref, logits_ref.grad)
assert logits.grad.shape == (bs, seqlen, vocab)
with Context(DEBUG=0):
assert loss.allclose(ref, atol=2e-3, rtol=2e-3).item(), "forward mismatch"
assert logits.grad.allclose(logits_ref.grad, atol=2e-3, rtol=2e-3).item(), "grad mismatch"
class TestFusedCE(unittest.TestCase):
def setUp(self):
if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
def test_fused_ce_1_2_16(self): run_fused_ce(1, 2, 16, label_smoothing=0.2)
def test_fused_ce_2_16_128(self): run_fused_ce(2, 16, 128)
def test_fused_ce_4_128_1024(self): run_fused_ce(4, 128, 1024, label_smoothing=0.2)
# note: this is the shape used in llama 8b
#def test_fused_ce_smoothing_16_1024_128256(self): run_fused_ce(16, 1024, 128256, label_smoothing=0.2)
def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
Tensor.manual_seed(0)
x = Tensor.randn(*shape).cast(dtypes.bfloat16).contiguous()
amax_state = Tensor.full((), 2.0, dtype=dtypes.float32).contiguous()
with Context(DEBUG=0): Tensor.realize(x, amax_state)
if delayed:
fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
Tensor.realize(fp8, inv_scale, new_amax)
Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
else:
fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
ref_fp8, _, _ = quantize_fp8(x, amax_state=amax_state)
Tensor.realize(fp8)
Tensor.realize(ref_fp8)
with Context(DEBUG=0):
assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
if delayed:
assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
class TestQuantizeFP8(unittest.TestCase):
def setUp(self):
ren = Device[Device.DEFAULT].renderer
if dtypes.bfloat16 not in ren.supported_dtypes(): self.skipTest("need bfloat16")
if not ren.has_local or not ren.has_shared: self.skipTest("need local/shared")
def test_scalar(self): run_quantize_fp8((getenv("N", 1024), 32), delayed=False)
def test_delayed(self): run_quantize_fp8((getenv("N", 2048), 1024))
@needs_second_gpu
def test_multi(self):
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
x = Tensor(x, device=devs)
amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
Tensor.realize(fp8, new_amax)
assert fp8.uop.shape == x.uop.shape
assert new_amax.shape == ()
class TestLocalAmax(unittest.TestCase):
def test_multi_tensor_local_shard_amax(self):
devices = ("CPU:0", "CPU:1")
x = Tensor.arange(16, device=devices[0]).reshape(4, 4).cast(dtypes.float).contiguous().realize().shard(devices, axis=0).realize()
GlobalCounters.reset()
out = (x * local_abs_max(x)).contiguous().realize()
self.assertEqual(GlobalCounters.kernel_count, 4)
self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
if __name__ == '__main__':
unittest.main()

1400
tinygrad_repo/test/backend/test_multitensor.py Normal file
View File

File diff suppressed because it is too large Load Diff

624
tinygrad_repo/test/backend/test_nn.py Normal file
View File

@@ -0,0 +1,624 @@
#!/usr/bin/env python
import unittest
import numpy as np
import torch
from tinygrad import Tensor, Device, TinyJit, dtypes
from tinygrad.uop.ops import Ops
from tinygrad.helpers import GlobalCounters, Context
from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
from tinygrad.nn.state import load_state_dict
from tinygrad.engine.realize import run_linear
from test.helpers import not_support_multi_device, needs_second_gpu, slow
@slow
class TestNN(unittest.TestCase):
def test_batchnorm2d(self, training=False, threed=False, track_running_stats=True):
with Tensor.train(training):
szs = [4, 8, 16, 32]
for sz in szs:
# create in tinygrad
bn = BatchNorm(sz, eps=1e-5, track_running_stats=track_running_stats)
bn.weight = Tensor.randn(sz)
bn.bias = Tensor.randn(sz)
if track_running_stats:
bn.running_mean = Tensor.randn(sz)
bn.running_var = Tensor.randn(sz)
bn.running_var.numpy()[bn.running_var.numpy() < 0] = 0
# create in torch
with torch.no_grad():
if threed:
tbn = torch.nn.BatchNorm3d(sz, track_running_stats=track_running_stats).eval()
else:
tbn = torch.nn.BatchNorm2d(sz, track_running_stats=track_running_stats).eval()
tbn.training = training
tbn.weight[:] = torch.tensor(bn.weight.numpy())
tbn.bias[:] = torch.tensor(bn.bias.numpy())
if track_running_stats:
tbn.running_mean[:] = torch.tensor(bn.running_mean.numpy())
tbn.running_var[:] = torch.tensor(bn.running_var.numpy())
if track_running_stats:
np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6)
np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5, atol=1e-6)
# trial
if threed:
inn = Tensor.randn(2, sz, 3, 3, 3)
else:
inn = Tensor.randn(2, sz, 3, 3)
# in tinygrad
outt = bn(inn)
# in torch
toutt = tbn(torch.tensor(inn.numpy()))
# close
np.testing.assert_allclose(outt.numpy(), toutt.detach().numpy(), rtol=5e-4, atol=1e-6)
if track_running_stats:
np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6)
np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5, atol=1e-6)
def test_batchnorm2d_training(self): self.test_batchnorm2d(True, False, True)
def test_batchnorm2d_no_running_stats(self): self.test_batchnorm2d(False, False, False)
def test_batchnorm2d_training_no_running_stats(self): self.test_batchnorm2d(True, False, False)
def test_batchnorm3d(self): self.test_batchnorm2d(False, True, True)
def test_batchnorm3d_training(self): self.test_batchnorm2d(True, True, True)
def test_batchnorm3d_no_running_stats(self): self.test_batchnorm2d(False, True, False)
def test_batchnorm3d_training_no_running_stats(self): self.test_batchnorm2d(True, True, False)
def test_batchnorm_axis(self):
sz = (2, 4, 3, 2, 2)
x = Tensor.randn(sz)
weight = Tensor.randn(2, 3)
bias = Tensor.randn(2, 3)
mean = Tensor.randn(2, 3)
invstd = Tensor.randn(2, 3)
a = (x.batchnorm(weight, bias, mean, invstd, axis=(0, 2))
.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2))
b = (x.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2)
.batchnorm(weight.flatten(), bias.flatten(), mean.flatten(), invstd.flatten()))
t_x = torch.tensor(x.permute(1, 0, 2, 3, 4).reshape(4, 6, 2, 2).numpy())
t_weight, t_bias = torch.tensor(weight.flatten().numpy()), torch.tensor(bias.flatten().numpy())
t_mean, t_invstd = torch.tensor(mean.flatten().numpy()), torch.tensor(invstd.flatten().numpy())
torch.nn.functional.batch_norm(t_x, t_mean, 1.0 / t_invstd**2, t_weight, t_bias)
np.testing.assert_allclose(a.numpy(), b.numpy())
def test_linear(self):
def _test_linear(x, in_dim, out_dim):
# create in tinygrad
model = Linear(in_dim, out_dim)
z = model(x)
# create in torch
with torch.no_grad():
torch_layer = torch.nn.Linear(in_dim, out_dim).eval()
torch_layer.weight[:] = torch.tensor(model.weight.numpy(), dtype=torch.float32)
torch_layer.bias[:] = torch.tensor(model.bias.numpy(), dtype=torch.float32)
torch_x = torch.tensor(x.numpy(), dtype=torch.float32)
torch_z = torch_layer(torch_x)
# test
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
BS, T, in_dim, out_dim = 4, 2, 8, 16
_test_linear(Tensor.randn(BS, in_dim), in_dim, out_dim)
_test_linear(Tensor.randn(BS, T, in_dim), in_dim, out_dim) # test with more dims
def _test_conv(self, tiny_conv, torch_conv, BS, C1, DIMS, C2, K, S, P, D=1):
# create in tinygrad
layer = tiny_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D)
# create in torch
with torch.no_grad():
torch_layer = torch_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D).eval()
torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
# test
x = Tensor.uniform(BS, C1, *DIMS)
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
def test_conv1d(self): self._test_conv(Conv1d, torch.nn.Conv1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
def test_conv2d(self): self._test_conv(Conv2d, torch.nn.Conv2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
def test_conv1d_same_padding(self):
self._test_conv(Conv1d, torch.nn.Conv1d, BS=8, C1=3, DIMS=[32], C2=16, K=3, S=1, P='same')
def test_conv2d_same_padding_odd_input(self):
self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[29, 31], C2=32, K=5, S=1, P='same')
def test_conv2d_same_padding_large_kernel(self):
self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
def test_conv2d_same_padding_with_dilation(self):
self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)
def test_conv2d_same_padding_invalid_stride(self):
self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
def test_conv2d_same_padding_invalid_padding_str(self):
self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=1, padding='not_same')
@unittest.skip("Takes too long to compile for Compiled backends")
def test_conv2d_winograd(self):
BS, C1, H, W = 2, 8, 16, 16
C2, K, S, P = 8, 3, 1, 1
# create in tinygrad
layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
# create in torch
torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
torch_layer.weight = torch.nn.Parameter(torch.tensor(layer.weight.numpy(), dtype=torch.float32))
torch_layer.bias = torch.nn.Parameter(torch.tensor(layer.bias.numpy(), dtype=torch.float32))
# test
x = Tensor.uniform(BS, C1, H, W)
with Context(WINO=1):
z = layer(x)
m = z.mean()
m.backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
gw = layer.weight.grad.realize()
gb = layer.bias.grad.realize()
gx = x.grad.realize()
torch_z.mean().backward()
np.testing.assert_allclose(gw.numpy(), torch_layer.weight.grad.numpy(), atol=5e-4, rtol=1e-5)
np.testing.assert_allclose(gb.numpy(), torch_layer.bias.grad.numpy(), atol=5e-4, rtol=1e-5)
np.testing.assert_allclose(gx.numpy(), torch_x.grad.numpy(), atol=5e-4, rtol=1e-5)
def test_conv_transpose1d(self):
self._test_conv(ConvTranspose1d, torch.nn.ConvTranspose1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
def test_conv_transpose2d(self):
self._test_conv(ConvTranspose2d, torch.nn.ConvTranspose2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
def test_groupnorm(self):
BS, H, W, C, G = 20, 10, 10, 6, 3
# create in torch
torch_layer = torch.nn.GroupNorm(G, C).eval()
# create in tinygrad
layer = GroupNorm(G, C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(BS, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
def test_layernorm_forward(self):
N, C, H, W = 20, 5, 10, 10
# create in torch
torch_layer = torch.nn.LayerNorm([H, W]).eval()
# create in tinygrad
layer = LayerNorm([H, W])
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
x = Tensor.empty(N, C, H, W)
z = layer(x)
z.realize()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
# TODO: why is torch numbers all 0?
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=5e-6)
def test_layernorm(self):
N, C, H, W = 20, 5, 10, 10
# create in torch
torch_layer = torch.nn.LayerNorm([H, W]).eval()
# create in tinygrad
layer = LayerNorm([H, W])
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
def test_layernorm_2d(self):
N, C, H, W = 20, 5, 10, 10
# create in torch
torch_layer = torch.nn.LayerNorm([C]).eval()
# create in tinygrad
layer = LayerNorm2d(C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x.permute(0,2,3,1)).permute(0,3,1,2)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=5e-4, rtol=5e-4)
def test_instancenorm_2d(self):
N, C, H, W = 20, 10, 10, 10
# create in torch
torch_layer = torch.nn.InstanceNorm2d(C, affine=True).eval()
# create in tinygrad
layer = InstanceNorm(C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
def test_instancenorm_3d(self):
N, C, D, H, W = 20, 10, 10, 10, 10
# create in torch
torch_layer = torch.nn.InstanceNorm3d(C, affine=True).eval()
# create in tinygrad
layer = InstanceNorm(C)
layer.weight = Tensor(torch_layer.weight.detach().numpy())
layer.bias = Tensor(torch_layer.bias.detach().numpy())
for _ in range(10):
# forward
x = Tensor.randn(N, C, D, H, W)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
# TODO: is this numerical issue or a bug? RANGEIFY big reduce kernel amplifies numerical issue
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=1e-2, rtol=1e-3)
np.testing.assert_allclose(layer.bias.grad.numpy(), torch_layer.bias.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
def test_rmsnorm(self):
class TorchRMSNorm(torch.nn.Module):
# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L34C1-L77C36
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
super().__init__()
self.eps = eps
self.elementwise_affine = elementwise_affine
self.weight = torch.nn.Parameter(torch.ones(dim)) if elementwise_affine else None
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
output = self._norm(x.float()).type_as(x)
return output if self.weight is None else output * self.weight
B, T, embed_size = 4, 10, 20
torch_layer = TorchRMSNorm(embed_size)
layer = RMSNorm(embed_size)
for _ in range(10):
# forward
x = Tensor.randn(B, T, embed_size)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
np.testing.assert_allclose(layer.weight.grad.numpy(), torch_layer.weight.grad.detach().numpy(), atol=2e-3, rtol=1e-3)
torch_layer = TorchRMSNorm(embed_size, elementwise_affine=False)
layer = RMSNorm(embed_size, elementwise_affine=False)
for _ in range(10):
# forward
x = Tensor.randn(B, T, embed_size)
z = layer(x)
z.sum().backward()
torch_x = torch.tensor(x.numpy(), requires_grad=True)
torch_z = torch_layer(torch_x)
torch_z.sum().backward()
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-6, rtol=5e-6)
np.testing.assert_allclose(x.grad.numpy(), torch_x.grad.detach().numpy(), atol=1e-3, rtol=1e-3)
def test_embedding(self):
B, T, embed_size, vocab_size = 4, 10, 20, 28
# create in tinygrad
layer = Embedding(vocab_size, embed_size)
with torch.no_grad():
torch_layer = torch.nn.Embedding(vocab_size, embed_size).eval()
torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
# test
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# test with empty input length
x = Tensor(np.random.randint(0, vocab_size, (B, 0), dtype=np.int32))
z = layer(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
# test with jit enabled
@TinyJit
def layer_jit(x):
return layer(x).realize()
for _ in range(3):
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer_jit(x)
torch_x = torch.tensor(x.numpy())
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
def test_embedding_one_kernel(self, ops=612000, kcount=2):
GlobalCounters.reset()
layer = Embedding(20, 30)
layer.weight = Tensor.zeros_like(layer.weight).contiguous()
a = Tensor([[1, 5, 9, 11],
[12, 19, 8, 1]])
result = layer(a)
linear, var_vals = result.linear_with_vars()
self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount,
"first run realizes weight and embedding")
run_linear(linear, var_vals)
b = Tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
result = layer(b)
linear, var_vals = result.linear_with_vars()
self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]),
"second run realizes embedding only")
run_linear(linear, var_vals)
print(f"Embedding used {GlobalCounters.global_ops} ops")
self.assertLessEqual(GlobalCounters.global_ops, ops)
# TODO: fused with opts uses more ops
def test_embedding_one_kernel_fused(self):
with Context(NOOPT=0):
self.test_embedding_one_kernel(ops=612_000, kcount=2)
def test_embedding_one_kernel_fused_noopt(self):
with Context(NOOPT=1):
self.test_embedding_one_kernel(ops=0, kcount=2)
def test_embedding_shape(self):
vocab_size, embed_size = 10, 16
layer = Embedding(vocab_size, embed_size)
for rank in range(5):
shp = (1,) * rank
a = Tensor([3]).reshape(shp)
result = layer(a)
self.assertEqual(result.shape, shp + (embed_size,))
def test_embedding_regression(self):
# used to fail bounds check
embedding = Embedding(100, 1024)
input_ids = Tensor.empty(16, 16, dtype=dtypes.int)
embedding(input_ids).realize()
def test_load_state_dict(self):
layer = Conv2d(3, 5, kernel_size=3)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3),
'bias': Tensor.randn(5),
}
load_state_dict(layer, state_dict)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
#https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/torch/nn/modules/module.py#L2425
def test_load_conv_num_batches_tracked(self):
layer = BatchNorm(sz=1, track_running_stats=False)
state_dict = {
'weight': Tensor.ones(1),
'bias': Tensor.ones(1),
'num_batches_tracked': Tensor.ones(1),
}
load_state_dict(layer, state_dict)
state_dict['num_batches_tracked'] = Tensor.empty()
load_state_dict(layer, state_dict)
layer.num_batches_tracked = Tensor.ones(1)
load_state_dict(layer, state_dict)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_load_state_dict_sharded_model(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
layer = Conv2d(3, 5, kernel_size=3)
layer.weight.shard_(devices, 3)
layer.bias.shard_(devices, None)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).realize(),
'bias': Tensor.randn(5).realize(),
}
load_state_dict(layer, state_dict)
# sharded model shards the state_dict
self.assertEqual(layer.weight.device, devices)
self.assertEqual(layer.weight.uop.axis, 3)
self.assertEqual(layer.bias.device, devices)
self.assertEqual(layer.bias.uop.axis, None)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
@unittest.skipIf(not_support_multi_device, "no multi")
def test_load_state_dict_sharded_dict(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
layer = Conv2d(3, 5, kernel_size=3)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).shard(devices, 3),
'bias': Tensor.randn(5).shard(devices, None),
}
load_state_dict(layer, state_dict)
# NOTE: model is not sharded, still not sharded after load_state_dict
self.assertEqual(layer.weight.device, Device.DEFAULT)
self.assertEqual(layer.bias.device, Device.DEFAULT)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_load_state_dict_sharded_model_dict_same_axis(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
layer = Conv2d(3, 5, kernel_size=3)
layer.weight.shard_(devices, 3)
layer.bias.shard_(devices, None)
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).shard(devices, 3),
'bias': Tensor.randn(5).shard(devices, None),
}
load_state_dict(layer, state_dict)
self.assertEqual(layer.weight.device, devices)
self.assertEqual(layer.weight.uop.axis, 3)
self.assertEqual(layer.bias.device, devices)
self.assertEqual(layer.bias.uop.axis, None)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
@unittest.skipIf(not_support_multi_device, "no multi")
def test_load_state_dict_sharded_model_dict_different_axis(self):
devices = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3")
devices5 = (f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3", f"{Device.DEFAULT}:4", f"{Device.DEFAULT}:5")
layer = Conv2d(3, 5, kernel_size=3)
layer.weight.shard_(devices, 3)
layer.bias.shard_(devices, None)
# different shard axis
state_dict = {
'weight': Tensor.randn(5, 3, 3, 3).shard(devices, None),
'bias': Tensor.randn(5).shard(devices5, 0),
}
load_state_dict(layer, state_dict)
# NOTE: model and state_dict shard differently, use the state_dict sharding # TODO: revisit this?
self.assertEqual(layer.weight.device, devices)
self.assertEqual(layer.weight.uop.axis, None)
self.assertEqual(layer.bias.device, devices5)
self.assertEqual(layer.bias.uop.axis, 0)
np.testing.assert_allclose(layer.weight.numpy(), state_dict['weight'].numpy())
np.testing.assert_allclose(layer.bias.numpy(), state_dict['bias'].numpy())
def test_load_state_dict_shape_mismatch(self):
d1, d2 = 2, 4
layer = Linear(d1, d1, bias=False)
state_dict = {'weight': Tensor.randn(d2, d2)}
with self.assertRaisesRegex(ValueError, r'Shape mismatch in layer `weight`: Expected shape \(2, 2\), but found \(4, 4\) in state dict.'):
load_state_dict(layer, state_dict)
def test_lstm_cell(self):
layer = LSTMCell(32, 16)
with torch.no_grad():
torch_layer = torch.nn.LSTMCell(32, 16)
layer.weight_hh.assign(torch_layer.weight_hh.numpy())
layer.weight_ih.assign(torch_layer.weight_ih.numpy())
layer.bias_hh.assign(torch_layer.bias_hh.numpy())
layer.bias_ih.assign(torch_layer.bias_ih.numpy())
inp = Tensor.randn(1, 32)
out_h, out_c = layer(inp)
torch_out_h, torch_out_c = torch_layer(torch.tensor(inp.numpy()))
np.testing.assert_allclose(out_h.numpy(), torch_out_h.numpy(), atol=1e-6)
np.testing.assert_allclose(out_c.numpy(), torch_out_c.numpy(), atol=1e-6)
out_h, out_c = layer(inp, (out_h, out_c))
torch_out_h, torch_out_c = torch_layer(torch.tensor(inp.numpy()), (torch_out_h, torch_out_c))
np.testing.assert_allclose(out_h.numpy(), torch_out_h.numpy(), atol=1e-6)
np.testing.assert_allclose(out_c.numpy(), torch_out_c.numpy(), atol=1e-6)
def test_lstm_cell_no_bias(self):
layer = LSTMCell(32, 16, bias=False)
inp = Tensor.randn(1, 32)
out_h, out_c = layer(inp)
out_h.realize()
out_c.realize()
h = Tensor.randn(1, 16)
c = Tensor.randn(1, 16)
out_h, out_c = layer(inp, (h, c))
out_h.realize()
out_c.realize()
assert layer.bias_hh is None
assert layer.bias_ih is None
if __name__ == '__main__':
unittest.main()

3414
tinygrad_repo/test/backend/test_ops.py Normal file
View File

File diff suppressed because it is too large Load Diff

46
tinygrad_repo/test/backend/test_opt_gemm.py Normal file
View File

@@ -0,0 +1,46 @@
import numpy as np
import unittest
from tinygrad import Tensor
from tinygrad.helpers import get_single_element
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import run_linear
from tinygrad.uop.ops import Ops, UOp
from test.helpers import replace_opts
class TestOptGemm(unittest.TestCase):
@classmethod
def setUpClass(cls):
N = 64
cls.a = Tensor.randn(N, N).contiguous().realize()
cls.b = Tensor.randn(N, N).contiguous().realize()
with np.errstate(all='ignore'):
cls.res = cls.a.T.numpy() @ cls.b.T.numpy()
def _test_gemm_unrolled_permute_l(self, opts=[]):
t = self.a.T @ self.b.T
# TODO: this should be a generic test helper
call = get_single_element(t.schedule_linear().src)
new_call = call.replace(src=(replace_opts(call.src[0], opts), *call.src[1:]))
run_linear(UOp(Ops.LINEAR, src=(new_call,)))
test = call.src[1].buffer.numpy().reshape(self.res.shape)
np.testing.assert_allclose(self.res, test, atol=1e-4)
def test_gemm_unrolled_permute_l_44(self):
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4)]
self._test_gemm_unrolled_permute_l(opts)
def test_gemm_unrolled_permute_l_424(self):
# was failing with LLVM
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
self._test_gemm_unrolled_permute_l(opts)
def test_gemm_unrolled_permute_l_42(self):
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
self._test_gemm_unrolled_permute_l(opts)
def test_gemm_unrolled_permute_l_22(self):
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=2)]
self._test_gemm_unrolled_permute_l(opts)
if __name__ == '__main__':
unittest.main()

195
tinygrad_repo/test/backend/test_optim.py Normal file
View File

@@ -0,0 +1,195 @@
import numpy as np
import torch
import unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
from test.helpers import needs_second_gpu, slow
np.random.seed(1337)
x_init = np.random.randn(1,4).astype(np.float32)
W_init = np.random.randn(4,4).astype(np.float32)
m_init = np.random.randn(1,4).astype(np.float32)
def _param(tensor, val):
return tensor(val, requires_grad=True) if tensor is torch.tensor else tensor(val)
class TeenyNet:
def __init__(self, tensor):
self.x = _param(tensor, x_init.copy())
self.W = _param(tensor, W_init.copy())
def forward(self):
return (self.x * self.W).sum()
class TinyNet:
def __init__(self, tensor):
self.x = _param(tensor, x_init.copy())
self.W = _param(tensor, W_init.copy())
self.m = tensor(m_init.copy())
def forward(self):
out = self.x.matmul(self.W).relu()
# print(out.detach().numpy())
out = out.log_softmax(1)
out = out.mul(self.m).add(self.m).sum()
return out
def step(tensor, optim, steps=1, teeny=False, **kwargs):
net = TeenyNet(tensor) if teeny else TinyNet(tensor)
optim = optim([net.x, net.W], **kwargs)
for _ in range(steps):
out = net.forward()
optim.zero_grad()
out.backward()
optim.step()
return net.x.detach().numpy(), net.W.detach().numpy()
@slow
class TestOptim(unittest.TestCase):
def setUp(self):
self.old_training = Tensor.training
Tensor.training = True
def tearDown(self):
Tensor.training = self.old_training
def _test_optim(self, tinygrad_optim, torch_optim, steps, opts, atol, rtol):
for x,y in zip(step(Tensor, tinygrad_optim, steps, **opts),
step(torch.tensor, torch_optim, steps, **opts)):
np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
def _test_sgd(self, steps, opts, atol, rtol): self._test_optim(SGD, torch.optim.SGD, steps, opts, atol, rtol)
def _test_adam(self, steps, opts, atol, rtol): self._test_optim(Adam, torch.optim.Adam, steps, opts, atol, rtol)
def _test_adamw(self, steps, opts, atol, rtol): self._test_optim(AdamW, torch.optim.AdamW, steps, opts, atol, rtol)
def _test_muon(self, steps, opts, atol, rtol): self._test_optim(Muon, torch.optim.Muon, steps, opts, atol, rtol)
def test_multistep_sgd_high_lr_teeny(self): self._test_sgd(2, {'lr': 1.1, 'teeny': True}, 1e-6, 1e-5)
def test_multistep_adam_high_lr_teeny(self): self._test_adam(2, {'lr': 1.1, 'teeny': True}, 2e-4, 5e-4)
def test_multistep_muon_high_lr_teeny(self): self._test_muon(2, {'lr': 1.1, 'teeny': True}, 1e-2, 5e-4)
def test_sgd(self): self._test_sgd(1, {'lr': 0.001}, 1e-6, 0)
def test_sgd_high_lr(self): self._test_sgd(1, {'lr': 10}, 1e-6, 1e-5)
def test_sgd_wd(self): self._test_sgd(1, {'lr': 0.001, 'weight_decay': 0.1}, 1e-6, 0)
def test_sgd_high_lr_wd(self): self._test_sgd(1, {'lr': 10, 'weight_decay': 0.1}, 1e-6, 1e-5)
def test_multistep_sgd(self): self._test_sgd(10, {'lr': 0.001}, 1e-6, 0)
def test_multistep_sgd_high_lr(self): self._test_sgd(10, {'lr': 10}, 1e-6, 3e-4)
def test_multistep_sgd_wd(self): self._test_sgd(10, {'lr': 0.001, 'weight_decay': 0.1}, 1e-6, 0)
def test_multistep_sgd_high_lr_wd(self): self._test_sgd(10, {'lr': 9, 'weight_decay': 0.1}, 1e-6, 3e-4)
def test_multistep_sgd_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9}, 1e-6, 0)
def test_multistep_sgd_high_lr_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9}, 1e-5, 3e-4)
def test_multistep_sgd_momentum_wd(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'weight_decay': 0.1}, 1e-6, 0)
def test_multistep_sgd_high_lr_momentum_wd(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'weight_decay': 0.1}, 1e-5, 3e-4)
def test_multistep_sgd_nesterov_momentum(self): self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True}, 1e-5, 0)
def test_multistep_sgd_high_lr_nesterov_momentum(self): self._test_sgd(10, {'lr': 10, 'momentum': 0.9, 'nesterov': True}, 1e-5, 3e-4)
def test_multistep_sgd_nesterov_momentum_wd(self):
self._test_sgd(10, {'lr': 0.001, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 0)
def test_multistep_sgd_high_lr_nesterov_momentum_wd(self):
self._test_sgd(10, {'lr': 9, 'momentum': 0.9, 'nesterov': True, 'weight_decay': 0.1}, 1e-5, 3e-4)
def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
# TODO: disabled due to big atol
# def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
# TODO: disabled due to big atol
# def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)
# NOTE: momentum set to 0.95 by default, nesterov set to True by default
def test_multistep_muon_momentum_wd(self): self._test_muon(10, {'lr': 0.001, 'weight_decay': 0.01}, 3e-3, 0)
# ns defaults are numerically unstable, but it is tolerable in real training (see nsteps/nparam tests)
# TODO: disabled due to big atol
# def test_multistep_muon_high_lr_momentum_wd(self): self._test_muon(10, {'lr': 10, 'weight_decay': 0.01}, 1e-1, 3e-4)
def test_multistep_muon_no_nesterov_momentum(self): self._test_muon(10, {'lr': 0.001, 'nesterov': False}, 1e-3, 0)
# TODO: disabled due to big atol
# def test_multistep_muon_high_lr_no_nesterov_momentum(self): self._test_muon(10, {'lr': 10, 'nesterov': False}, 5e-2, 1e-1)
def test_muon_ns_steps(self): self._test_muon(1, {'lr': 0.001, 'ns_steps': 3}, 1e-4, 0)
# TODO: disabled due to big atol
# def test_muon_high_lr_ns_steps(self): self._test_muon(1, {'lr': 10, 'ns_steps': 3}, 1e-5, 3e-4)
def test_muon_ns_coefficients(self): self._test_muon(1, {'lr': 0.001,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
# TODO: disabled due to big atol
# def test_muon_high_lr_ns_coefficients(self): self._test_muon(1, {'lr': 10,'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
def test_muon_momentum_wd_ns_steps_ns_coefficients(self):
self._test_muon(10, {'lr': 0.001, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-4, 0)
# TODO: disabled due to big atol
# def test_multistep_muon_high_lr_momentum_wd_ns_steps_ns_coefficients(self):
# self._test_muon(10, {'lr': 10, 'momentum': 0.90, 'weight_decay': 0.01, 'ns_steps': 3, 'ns_coefficients': (2.0,-1.5,0.5)}, 1e-5, 3e-4)
def test_adam(self): self._test_adam(1, {'lr': 0.001}, 1e-5, 0)
def test_adam_high_lr(self): self._test_adam(1, {'lr': 10}, 1e-4, 1e-4)
def test_adamw(self): self._test_adamw(1, {'lr': 0.001}, 1e-5, 0)
def test_adamw_high_lr(self): self._test_adamw(1, {'lr': 10}, 1e-4, 1e-4)
def test_multistep_adam(self): self._test_adam(10, {'lr': 0.001}, 1e-5, 0)
def test_multistep_adam_high_lr(self): self._test_adam(10, {'lr': 10}, 2e-3, 5e-4)
def test_multistep_adamw(self): self._test_adamw(10, {'lr': 0.001}, 1e-5, 0)
def test_multistep_adamw_high_lr(self): self._test_adamw(10, {'lr': 10}, 5e-4, 2e-3)
def test_duped_weights(self):
for Opt in [Adam, AdamW, SGD]:
losses = []
for i in range(2):
w = Tensor(x_init.copy())
opt = Opt([w], lr=0.1) if i == 0 else Opt([w, w], lr=0.1)
loss = None
for _ in range(3):
loss = w.sum()
opt.zero_grad()
loss.backward()
opt.step()
losses.append(loss.numpy())
np.testing.assert_allclose(losses[0], losses[1], atol=1e-4, rtol=0)
@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
def test_mixed_precision(self):
old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half
# weight update would overflow without upcasting
self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4)
self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4)
self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4)
dtypes.default_float = old_default_float
def test_assert_tensor_train(self):
t = Tensor.ones((1,1))
optimizer = Adam([t])
optimizer.zero_grad()
old_state = Tensor.training
t.sum().backward()
Tensor.training = False
self.assertRaises(RuntimeError, optimizer.step)
Tensor.training = True
optimizer.step()
Tensor.training = old_state
def test_lamb_cpu_offload(self):
# test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
t = Tensor(x_init.copy())
opt = LAMB([t])
# move optimizer state to CPU
for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
# run a step
t.sum().backward()
opt.step()
self.assertEqual(t.device, Device.DEFAULT)
self.assertEqual(opt.m[0].device, "CPU")
@needs_second_gpu
def test_lamb_cpu_offload_multi(self):
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
t = Tensor(x_init.copy()).shard(ds, axis=1)
ds = t.device
opt = LAMB([t])
# move optimizer state to CPU
for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
# run a step
t.sum().backward()
opt.step()
self.assertEqual(t.device, ds)
self.assertEqual(opt.m[0].device, "CPU")
if __name__ == '__main__':
unittest.main()

180
tinygrad_repo/test/backend/test_pickle.py Normal file
View File

@@ -0,0 +1,180 @@
import unittest, pickle, types
import numpy as np
from tinygrad import Tensor, TinyJit, Variable, dtypes
from tinygrad.helpers import GlobalCounters, ContextVar, Context
from tinygrad.uop.ops import PatternMatcher, UPat, UOp
class TestPickle(unittest.TestCase):
def test_pickle_code_object(self):
y = lambda x: x*2 # noqa: E731
code_str = pickle.dumps(y.__code__)
fxn = types.FunctionType(pickle.loads(code_str), globals())
self.assertEqual(fxn(2), 4)
def test_pickle_pattern_matcher(self):
pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
sink = UOp.const(dtypes.int, 2)
tt = pm.rewrite(sink)
pm_str = pickle.dumps(pm)
pm2 = pickle.loads(pm_str)
self.assertEqual(pm2.rewrite(sink).key, tt.key)
def test_pickle_main_pattern_matcher(self):
from tinygrad.uop.symbolic import sym
ssym = pickle.dumps(sym)
dsym = pickle.loads(ssym)
self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)
def test_pickle_realized_tensor(self):
print("** init")
t = Tensor.rand(10, 10).realize()
st = pickle.dumps(t)
t_values = t.numpy()
del t # free buffers
print("** post pickle")
GlobalCounters.reset()
t2:Tensor = pickle.loads(st)
np.testing.assert_equal(t_values, t2.numpy())
# expect at most one COPY kernel
self.assertLessEqual(GlobalCounters.kernel_count, 1)
def test_pickle_realized_tensor_alt(self):
print("** init")
t = Tensor.rand(10, 10).to("CPU").realize()
st = pickle.dumps(t)
t_values = t.numpy()
del t # free buffers
print("** post pickle")
t2:Tensor = pickle.loads(st)
assert t2.uop.is_realized
np.testing.assert_equal(t_values, t2.numpy())
def test_pickle_realized_tensor_alt2(self):
print("** init")
t = Tensor.rand(10, 10).to("CPU").realize()
tensor_uop = t.uop
assert tensor_uop.is_realized, f"expected {tensor_uop} to be realized"
t_values = t.numpy()
# pickle
st = pickle.dumps(t)
# free buffers
del t
del tensor_uop
print("** post pickle")
t2:Tensor = pickle.loads(st)
assert t2.uop.is_realized, f"expected {t2.uop} to be realized"
np.testing.assert_equal(t_values, t2.numpy())
# NOTE: currently Buffer exists on the uop, not tensor
def test_pickle_buffer_uop(self):
t = Tensor.arange(4).clone().realize()
a = t.uop
assert a.is_realized
self.assertIsNotNone(buffer:=a.base.realized)
s = pickle.dumps(a)
# free buffers
del a
del buffer
a2:UOp = pickle.loads(s)
self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])
def test_pickle_unrealized_tensor(self):
t = Tensor.ones(10, 10)
st = pickle.dumps(t)
t2:Tensor = pickle.loads(st)
np.testing.assert_equal(t.numpy(), t2.numpy())
def test_pickle_variable(self):
v = Variable("i", 1, 20).bind(10)
t1 = Tensor.ones(10, v).contiguous()
t2 = Tensor.ones(10, v).contiguous()
ret = (t1+t2).sum(1)
st = pickle.dumps(ret)
del ret
vt2 = pickle.loads(st)
np.testing.assert_equal(vt2.numpy(), 20)
def test_pickle_buffer_view(self):
t = Tensor.arange(10).clone(device="CPU").realize()
vt = t[3:5].contiguous().realize()
assert hasattr(vt.uop.buffer, 'base')
ref_value = vt.tolist()
st = pickle.dumps(vt)
del t, vt
vt2 = pickle.loads(st)
assert hasattr(vt2.uop.buffer, 'base')
assert ref_value == vt2.tolist()
def test_pickle_numpy(self):
t = Tensor(np.array([1,2,3,4.]), dtype=dtypes.float32)
st = pickle.dumps(t)
t2:Tensor = pickle.loads(st)
np.testing.assert_equal(t.numpy(), t2.numpy())
def test_pickle_jit(self):
@TinyJit
def add(a, b): return a.sum()+b+1
for _ in range(3): add(Tensor.rand(10, 10), Tensor.rand(10, 10))
st = pickle.dumps(add)
del add
add_fxn = pickle.loads(st)
x = Tensor.ones(10, 10).contiguous().realize()
y = Tensor.ones(10, 10).contiguous().realize()
print("post jit")
out = add_fxn(x, y)
np.testing.assert_equal(out.numpy(), 102)
def test_pickle_jit_no_del(self):
@TinyJit
def fn(x): return x + 1.0
for _ in range(3): fn(Tensor.randn(4))
loaded = pickle.loads(pickle.dumps(fn))
self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0])
def test_pickle_context_var(self):
v = ContextVar("test_var", 0)
with Context(test_var=1):
vs = pickle.dumps(v)
v2 = pickle.loads(vs)
self.assertEqual(v2.value, 1)
def test_pickle_schedule(self):
a = Tensor([1,2])
out = a + 2
sched = out.schedule_linear()
pk = pickle.dumps(sched)
sched_pk = pickle.loads(pk)
self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0])
def test_pickle_renderer(self):
from tinygrad.device import Device
pk = pickle.dumps(Device.default.renderer)
pickle.loads(pk)
class TestPickleJIT(unittest.TestCase):
@classmethod
def setUpClass(cls):
N = 10
@TinyJit
def add(a, b): return a.sum()+b+1
for _ in range(3): add(Tensor.rand(N, N), Tensor.rand(N, N))
cls.st = pickle.dumps(add)
del add
def test_inspect(self):
import io
class FakeClass:
def __init__(self, *args, **kwargs):
print(self.module, self.name)
class InspectUnpickler(pickle.Unpickler):
def find_class(self, module, name): return type("SpecializedFakeClass", (FakeClass,), {"name": name, "module": module})
InspectUnpickler(io.BytesIO(self.st)).load()
@unittest.skip("we are still saving intermediate buffers")
def test_size(self):
# confirm no intermediate buffers are saved
self.assertLess(len(self.st), 1_000_000)
if __name__ == '__main__':
unittest.main()

227
tinygrad_repo/test/backend/test_profiler.py Normal file
View File

@@ -0,0 +1,227 @@
import unittest, struct, contextlib, statistics, gc
from tinygrad import Device, Tensor, dtypes, TinyJit
from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
from tinygrad.runtime.support.hcq import HCQCompiled
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
MOCKGPU = DEV.interface.startswith("MOCK")
def _dev_base(d):
p = d.split(":")
return p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"
@contextlib.contextmanager
def helper_collect_profile(*devs):
for dev in devs: dev.synchronize()
saved = [x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent) and x.device.startswith("METAL")]
Compiled.profile_events.clear()
for x in saved: Compiled.profile_events.append(x)
cpu_events.clear()
profile_list = []
with Context(PROFILE=1):
yield profile_list
for dev in devs: dev.synchronize()
for dev in devs: dev._at_profile_finalize()
for x in Compiled.profile_events: profile_list.append(x)
profile_list.extend(cpu_events)
def helper_profile_filter_device(profile, device:str):
assert any(getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent) for x in profile), f"device {device} is not registred"
dev_events = [x for x in profile if getattr(x, "device", None) == device and isinstance(x, ProfileDeviceEvent)]
assert len(dev_events) == 1, "only one device registration event is expected"
return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]
# TODO: support in HCQCompiled
is_cpu_hcq = Device.DEFAULT in {"CPU"}
@unittest.skipUnless((issubclass(type(Device[Device.DEFAULT]), HCQCompiled) and not is_cpu_hcq) or Device.DEFAULT in {"METAL"}, "Dev not supported")
class TestProfiler(unittest.TestCase):
@classmethod
def setUpClass(self):
TestProfiler.d0 = Device[Device.DEFAULT]
TestProfiler.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
TestProfiler.b = self.a + 1
si = self.b.schedule_linear().src[-1]
TestProfiler.prg = to_program(si.src[0], TestProfiler.d0.renderer)
TestProfiler.runtime = get_runtime(TestProfiler.d0.device, TestProfiler.prg)
TestProfiler.b.uop.buffer.allocate()
def test_profile_kernel_run(self, wait=False):
runner_name = TestProfiler.runtime.name
with helper_collect_profile(TestProfiler.d0) as profile:
gs, ls = TestProfiler.prg.arg.launch_dims({})
TestProfiler.runtime(TestProfiler.b.uop.buffer._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls, wait=wait)
profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent)]
assert len(kernel_runs) == 1, "one kernel run is expected"
assert kernel_runs[0].name == runner_name, "kernel name is not correct"
assert _dev_base(kernel_runs[0].device) == kernel_runs[0].device, "kernel should not be on a sub-device"
def test_profile_kernel_run_wait(self):
self.test_profile_kernel_run(wait=True)
def test_profile_copyin(self):
buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
with helper_collect_profile(TestProfiler.d0) as profile:
buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
assert len(kernel_runs) == 1, "one kernel run is expected"
def test_profile_multiops(self):
runner_name = TestProfiler.runtime.name
buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
with helper_collect_profile(TestProfiler.d0) as profile:
buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
gs, ls = TestProfiler.prg.arg.launch_dims({})
TestProfiler.runtime(buf1._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
buf1.copyout(memoryview(bytearray(buf1.nbytes)))
evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
assert len(evs) == 3, "3 kernel runs are expected"
# NOTE: order of events does not matter, the tool is responsible for sorting them
prg_events = [e for e in evs if e.device == TestProfiler.d0.device]
assert any(e.name == runner_name for e in prg_events), "kernel name is not correct"
#for i in range(1, 3):
# assert evs[i].st > evs[i-1].en, "timestamp not aranged"
def test_profile_multidev(self):
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception as e: self.skipTest(f"second device not available {e}")
buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
with helper_collect_profile(TestProfiler.d0, d1) as profile:
buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
for dev in [TestProfiler.d0.device, d1.device]:
evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
assert len(evs) == 1, "one kernel runs are expected"
def test_profile_multidev_transfer(self):
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception as e: self.skipTest(f"second device not available {e}")
buf1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:0").realize()
with helper_collect_profile(TestProfiler.d0, d1) as profile:
buf1.to(f"{Device.DEFAULT}:1").realize()
kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
assert len(kernel_runs) == 1, "one kernel run is expected"
@unittest.skipIf(Device.DEFAULT in "METAL" or (MOCKGPU and Device.DEFAULT == "AMD"), "AMD mockgpu does not support queue wait interrupts")
def test_profile_graph(self):
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception as e: self.skipTest(f"second device not available {e}")
def f(a):
x = (a + 1).realize()
return x, x.to(d1.device).realize()
a = Tensor.randn(10, 10, device=TestProfiler.d0.device).realize()
with helper_collect_profile(TestProfiler.d0, d1) as profile:
jf = TinyJit(f)
for _ in range(3): jf(a)
del jf
graph_evs = [x for x in profile if isinstance(x, ProfileGraphEvent)]
_, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
_, _ = helper_profile_filter_device(profile, d1.device)
assert len(graph_evs) == 2, "2 graph events are expected"
assert len(graph_evs[0].ents) == 2, "two entities are expected"
@unittest.skipIf(MOCKGPU, "skip MOCKGPU")
@unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "must be HCQ")
def test_dev_jitter_matrix(self):
dev_cnt = 6
try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
except Exception as e: self.skipTest(f"multiple devices not available {e}")
for dev in devs: dev.synchronize()
for dev in devs: dev._at_profile_finalize()
def _sync_d2d(d1:HCQCompiled, d2:HCQCompiled):
d1.hw_compute_queue_t().signal(d1.timeline_signal, d1.timeline_value).wait(d2.timeline_signal, d2.timeline_value) \
.timestamp(d1.timeline_signal).signal(d1.timeline_signal, d1.timeline_value+1).submit(d1)
d2.hw_compute_queue_t().signal(d2.timeline_signal, d2.timeline_value).wait(d1.timeline_signal, d1.timeline_value) \
.timestamp(d2.timeline_signal).signal(d2.timeline_signal, d2.timeline_value+1).submit(d2)
d1.timeline_value += 2
d2.timeline_value += 2
d1.timeline_signal.wait(d1.timeline_value - 1)
d2.timeline_signal.wait(d2.timeline_value - 1)
return d2.timeline_signal.timestamp - d1.timeline_signal.timestamp
# then test it by timing the GPU to GPU times
dev_evs = {x.device:x for x in Compiled.profile_events if isinstance(x, ProfileDeviceEvent)}
jitter_matrix = [[float('nan')] * len(devs) for _ in range(len(devs))]
pairs = [(p1, p2) for p1 in enumerate(devs) for p2 in enumerate(devs) if p1 != p2]
for (i1, d1), (i2, d2) in pairs:
cpu_diff = dev_evs[d1.device].tdiff - dev_evs[d2.device].tdiff
jitter_matrix[i1][i2] = statistics.median(_sync_d2d(d1, d2) - _sync_d2d(d2, d1) for _ in range(20)) / 2 - cpu_diff
print("pairwise clock jitter matrix (us):\n" + '\n'.join([''.join([f'{float(item):8.3f}' for item in row]) for row in jitter_matrix]))
for (i1, d1), (i2, d2) in pairs:
assert abs(jitter_matrix[i1][i2]) < 0.5, "jitter should be less than 0.5us"
def test_cpu_profile(self):
def test_fxn(err=False):
if err: raise Exception()
with helper_collect_profile(dev:=TestProfiler.d0) as profile:
with cpu_profile("test_1", dev):
test_fxn(err=False)
with self.assertRaises(Exception):
with cpu_profile("test_2", dev):
test_fxn(err=True)
range_events = [p for p in profile if isinstance(p, ProfileRangeEvent) and p.device == dev]
self.assertEqual(len(range_events), 2)
@unittest.skip("this test is flaky")
@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
def test_graph(self):
from test.backend.test_graph import helper_alloc_rawbuffer, helper_exec_op, helper_test_graphs
device = TestProfiler.d0.device
bufs = [helper_alloc_rawbuffer(device, fill=True) for _ in range(5)]
graphs = [[helper_exec_op(device, bufs[0], [bufs[1], bufs[2]]), helper_exec_op(device, bufs[0], [bufs[3], bufs[4]]),]]
with helper_collect_profile(dev:=TestProfiler.d0) as profile:
helper_test_graphs(dev.graph, graphs, runs:=2)
# NOTE: explicitly trigger deletion of all graphs
graphs.clear()
gc.collect()
graphs = [e for e in profile if isinstance(e, ProfileGraphEvent)]
self.assertEqual(len(graphs), runs)
for ge in graphs:
self.assertEqual(len(ge.ents), len(graphs))
@unittest.skip("this test is flaky")
def test_trace_metadata(self):
with Context(TRACEMETA=1):
a = Tensor.empty(1)+2
b = Tensor.empty(1)+2
with helper_collect_profile(TestProfiler.d0) as profile:
Tensor.realize(a, b)
profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and _dev_base(e.device) == e.device]
self.assertEqual(len(exec_points), len(range_events), 2)
self.assertEqual(len(dedup(e.arg['name'] for e in exec_points)), 1)
self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
if __name__ == "__main__":
unittest.main()

237
tinygrad_repo/test/backend/test_quantize_onnx.py Normal file
View File

@@ -0,0 +1,237 @@
# ruff: noqa: E501
import numpy as np
import unittest
from tinygrad import Tensor, Context, Device, dtypes, UOp
from tinygrad.uop.ops import Ops
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from test.helpers import replace_opts
N = 512
def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=False):
import onnx
from onnx import helper, numpy_helper, TensorProto
# Define input and output
input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [batch_size, in_size])
output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [batch_size, out_size])
# Create random weights and bias
W_data = np.random.randn(in_size, out_size).astype(np.float32)
W_init = numpy_helper.from_array(W_data, name="W")
if bias:
B_data = np.random.randn(out_size).astype(np.float32)
B_init = numpy_helper.from_array(B_data, name="B")
gemm_node = helper.make_node("Gemm", inputs=["input", "W", "B"], outputs=["output"], alpha=1.0, beta=1.0, transB=0)
graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init, B_init])
else:
gemm_node = helper.make_node("Gemm", inputs=["input", "W"], outputs=["output"], alpha=1.0, beta=1.0, transB=0)
graph_def = helper.make_graph([gemm_node], "SingleGemmGraph", [input_tensor], [output_tensor], initializer=[W_init])
# Create and save the model
#model_def = helper.make_model(graph_def, producer_name="single_gemm_example")
# TODO remove this once ORT supports 1.18.0
model_def = helper.make_model(graph_def, producer_name="single_gemm_example", ir_version=10, opset_imports=[helper.make_opsetid("", 22)])
onnx.save_model(model_def, model_path)
return model_path
def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
linear = out.schedule_linear()
call = linear.src[-1]
prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer)
if replace_src is not None:
old_name = prg.src[3].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
new_src = replace_src + "/* DSP boilerplate */" + prg.src[3].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
# drop BINARY and replace SOURCE so run_linear recompiles
prg = prg.replace(src=prg.src[:3] + (UOp(Ops.SOURCE, arg=new_src),))
linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),))
for _ in range(run_count): run_linear(linear)
def get_quantized_model(sz):
from onnxruntime.quantization import quantize_static, QuantFormat, QuantType, CalibrationDataReader
class FakeDataReader(CalibrationDataReader):
def __init__(self): self.cnt = 0
def get_next(self) -> dict:
self.cnt += 1
if self.cnt == 100: return None
return {"input": np.random.uniform(size=(sz, sz)).astype(np.float32)}
out_file = "/tmp/test_out.onnx"
quantize_static(create_gemm_model("/tmp/test_in.onnx", sz, sz, sz), out_file,
FakeDataReader(), quant_format=QuantFormat.QDQ, per_channel=False, reduce_range=False,
activation_type=QuantType.QUInt8, weight_type=QuantType.QInt8,
extra_options={"ActivationSymmetric": False})
return out_file
@unittest.skip("this is broken")
@unittest.skipIf(Device.DEFAULT != "CPU", "only tests for CPU")
class TestQuantizeOnnxCPU(unittest.TestCase):
def test_quant_128(self, sz=128):
try:
import onnx # noqa: F401 # pylint: disable=unused-import
except ImportError:
raise unittest.SkipTest()
from tinygrad.nn.onnx import OnnxRunner
out_file = get_quantized_model(sz)
run_onnx = OnnxRunner(out_file)
inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
with Context(QUANTIZE=1):
linear = run_onnx({"input":inp})["output"].schedule_linear()
prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
daccs = [u for u in tuple(prg.src[2].src) if u.op is Ops.DEFINE_REG]
assert all(u.dtype.scalar() is dtypes.int for u in daccs)
@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
class TestQuantizeOnnx(unittest.TestCase):
def test_quant_128(self): self.test_quant(128)
def test_quant(self, sz=512):
from examples.benchmark_onnx import load_onnx_model
# divide is ~1500-2000 without reduce_range, 750-900 with it
out_file = get_quantized_model(sz)
run_onnx_jit, _ = load_onnx_model(out_file)
run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))
def test_prequant_conv2d_1x1(self):
X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8))
out = X.conv2d(W, dtype=X.dtype)
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
def test_prequant_gemm(self):
N = 512
X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8))
out = X.matmul(W, dtype=X.dtype)
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
# TODO: this has to work
def test_prequant_gemm_intacc_early(self, xi=np.int8, wi=np.int8):
N = 512
X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(xi))
W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
# this divide is interesting and forces the accumulator to actually be an int
out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
def test_prequant_gemm_handcode(self):
src = """typedef int int128 __attribute__((aligned(512),vector_size(512)));
typedef int int32 __attribute__((aligned(128),vector_size(128)));
typedef int int64 __attribute__((aligned(256),vector_size(256)));
typedef unsigned char unsigned_char4 __attribute__((aligned(4),vector_size(4)));
typedef signed char signed_char128 __attribute__((aligned(128),vector_size(128)));
typedef unsigned char unsigned_char128 __attribute__((aligned(128),vector_size(128)));
typedef unsigned char unsigned_char256 __attribute__((aligned(256),vector_size(256)));
union V256 {
unsigned_char256 vec256;
struct {
unsigned_char128 lo128;
unsigned_char128 hi128;
};
};
__attribute__((noinline)) void fxn(unsigned char* restrict __attribute__((align_value(128))) data0,
unsigned char* restrict __attribute__((align_value(128))) data1,
signed char* restrict __attribute__((align_value(128))) data2) {
for (int ridx0 = 0; ridx0 < 512; ridx0++) {
int alu0 = (ridx0<<9);
for (int ridx1 = 0; ridx1 < 4; ridx1++) {
int alu1 = (ridx1<<7);
int32 acc0 = __builtin_HEXAGON_V6_vd0_128B();
int32 acc1 = __builtin_HEXAGON_V6_vd0_128B();
int32 acc2 = __builtin_HEXAGON_V6_vd0_128B();
int32 acc3 = __builtin_HEXAGON_V6_vd0_128B();
for (int ridx2 = 0; ridx2 < 128; ridx2++) {
unsigned_char4 val0 = *((unsigned_char4*)((data1+(alu0+(ridx2<<2)))));
int alu2 = (alu1+(ridx2<<11));
signed_char128 x0 = *((signed_char128*)((data2+alu2)));
signed_char128 x1 = *((signed_char128*)((data2+(alu2+512))));
signed_char128 x2 = *((signed_char128*)((data2+(alu2+1024))));
signed_char128 x3 = *((signed_char128*)((data2+(alu2+1536))));
union V256 ss01;
// ss01.lo128 = (x0[0], x1[0], x0[2], x1[2], x0[4], x1[4], ...)
// ss01.hi128 = (x0[1], x1[1], x0[3], x1[3], x0[5], x1[5], ...)
ss01.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x1, x0);
union V256 ss23;
// ss23.lo128 = (x2[0], x3[0], x2[2], x3[2], x2[4], x3[4], ...)
// ss23.hi128 = (x2[1], x3[1], x2[3], x3[3], x2[5], x3[5], ...)
ss23.vec256 = __builtin_HEXAGON_V6_vshufoeb_128B(x3, x2);
union V256 sslo;
// sslo.lo128 = (x0[0], x1[0], x2[0], x3[0], x0[4], x1[4], ...)
// sslo.hi128 = (x0[2], x1[2], x2[2], x3[2], x0[6], x1[6], ...)
sslo.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(ss23.lo128, ss01.lo128, 2);
union V256 sshi;
// sshi.lo128 = (x0[1], x1[1], x2[1], x3[1], x0[5], x1[5], ...)
// sshi.hi128 = (x0[3], x1[3], x2[3], x3[3], x0[7], x1[7], ...)
sshi.vec256 = __builtin_HEXAGON_V6_vdealvdd_128B(ss23.hi128, ss01.hi128, 2);
//unsigned_char128 w0 = (unsigned_char128){val0[0],val0[1],val0[2],val0[3],val0[0],val0[1],val0[2],val0[3],...
unsigned_char128 w0 = __builtin_HEXAGON_V6_lvsplatw_128B(*((unsigned int*)&val0));
acc0 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc0, w0, sslo.lo128);
acc1 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc1, w0, sshi.lo128);
acc2 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc2, w0, sslo.hi128);
acc3 = __builtin_HEXAGON_V6_vrmpybusv_acc_128B(acc3, w0, sshi.hi128);
}
acc0 /= 1000;
acc1 /= 1000;
acc2 /= 1000;
acc3 /= 1000;
// ','.join([f"acc{j}[{i}]" for i in range(32) for j in range(4)])
// acc0[0], acc0[1], acc0[2], ..... acc3[30], acc3[31]
unsigned_char128 packed = __builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B(acc3, acc2),
__builtin_HEXAGON_V6_vpackwh_sat_128B(acc1, acc0));
packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
packed = __builtin_HEXAGON_V6_vshuffb_128B(packed);
// acc0[0], acc1[0], acc2[0], ..... acc2[31], acc3[31]
*((unsigned_char128*)((data0+(alu0+alu1)))) = packed;
}
}
}"""
self.test_prequant_gemm_intacc(np.uint8, np.int8, src)
def test_prequant_gemm_intacc_32(self):
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=0)]
self.test_prequant_gemm_intacc(np.uint8, np.int8, N=32, opts=opts)
def test_prequant_gemm_intacc_128(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=128)
def test_prequant_gemm_intacc_256(self): self.test_prequant_gemm_intacc(np.uint8, np.int8, N=256)
def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
out = (X.int().matmul(W.int())//1000)
if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
out = out.cast(tg_dtype)
opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
sexec(out, opts, replace_src, run_count=1)
tout = out.numpy()
mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
mout = mout.astype(xi)
print(tout)
print(mout)
np.testing.assert_equal(tout, mout)
def test_prequant_gemm_intacc_wi(self): self.test_prequant_gemm_intacc(wi=np.int8)
def test_prequant_gemm_intacc_xiwi(self): self.test_prequant_gemm_intacc(xi=np.int8, wi=np.int8)
def test_prequant_gemm_intacc_xiwi_noclip(self): self.test_prequant_gemm_intacc(xi=np.int8, wi=np.int8, clip=False)
def test_prequant_gemv(self):
N = 2048
X = Tensor(np.random.uniform(0, 255, size=(1,N)).astype(np.uint8)).realize()
W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)).realize()
#out = X.cast(dtypes.int) @ W.cast(dtypes.int)
#out = X @ W
out = X.matmul(W, dtype=X.dtype)
opts = [Opt(op=OptOps.UPCAST, axis=0, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
sexec(out, opts)
if __name__ == "__main__":
unittest.main()

262
tinygrad_repo/test/backend/test_randomness.py Normal file
View File

@@ -0,0 +1,262 @@
import unittest, math
from tinygrad import dtypes, Tensor, Device
from tinygrad.helpers import getenv, DEV
from tinygrad.codegen import to_program
from tinygrad.uop.ops import Ops
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
from tinygrad.renderer.isa.x86 import X86Renderer
from test.helpers import not_support_multi_device, needs_second_gpu
from test.unit.test_randomness import equal_distribution, normal_test
import numpy as np
import torch
from hypothesis import given, settings, strategies as strat
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
class TestRandomness(unittest.TestCase):
def test_rand(self):
self.assertFalse(normal_test(Tensor.rand))
self.assertTrue(equal_distribution(Tensor.rand, torch.rand, lambda x: np.random.rand(*x)))
def test_rand_is_lazy(self):
Tensor.manual_seed(0)
r1 = Tensor.rand(10)
self.assertFalse(r1.uop.is_realized, "rand should be lazy - tensor should not be realized")
counter = Tensor._device_rng_counters[Device.DEFAULT]
self.assertFalse(counter.uop.is_realized, "rand should be lazy - counter should not be realized")
# second rand triggers assign path
r2 = Tensor.rand(10)
self.assertFalse(r2.uop.is_realized, "rand should be lazy - tensor should not be realized after second rand")
self.assertFalse(counter.uop.is_realized, "rand should be lazy - counter should not be realized after second rand")
Tensor.realize(r1, r2)
self.assertTrue(r1.uop.is_realized, "tensor should be realized after .realize()")
self.assertTrue(r2.uop.is_realized, "tensor should be realized after .realize()")
@unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support")
def test_rand_float16(self):
N = 128
x = Tensor.rand((2, N, N), dtype=dtypes.float16)
assert x.dtype == dtypes.float16
nx = x.numpy()
# seed dependant, check output range is [0, 1)
assert nx[nx == 1].size == 0
assert nx[nx == 0].size > 0
equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.float16), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry")
def test_threefry_against_reference(self):
Tensor.manual_seed(1337)
# reference generated using
"""
key0 = 1337
key1 = 0
values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.arange(20, dtype=np.uint32))
print(f"[{', '.join(f'{v}' for v in values)}]")
"""
jr = np.array([2221762175, 1752107825, 653745012, 1967534793, 1395205442, 3840423848, 2159346757,
603508235, 3319473678, 3363866483, 3544324138, 1436466838, 2169858556, 2570072943,
2387150698, 3678370550, 2911697663, 403244401, 2560861638, 1692360114])
counts = Tensor.arange(20, dtype=dtypes.uint32)
counts0, counts1 = counts.chunk(2)
r = Tensor._threefry_random_bits(Tensor([0, 1337], dtype='uint32'), counts0, counts1).numpy()
np.testing.assert_allclose(jr, r)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "X86 callee saved registers have ulong dtype")
def test_threefry_doesnt_use_long(self):
linear = Tensor.rand(20).schedule_linear()
for call in linear.src:
ast = call.src[0]
if ast.op is Ops.SINK:
prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer)
for u in tuple(prg.src[2].src):
self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}")
def test_threefry_against_reference_full(self):
Tensor.manual_seed(1337)
# reference generated using
"""
key0 = 1337
key1 = int.from_bytes(hashlib.sha256(int(0).to_bytes(4)).digest(), "big") & 0xffffffff
# derive new key for the counter offset (c_low=0, c_high=0 for first call)
new_key_values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.array([0, 0], dtype=np.uint32))
new_key = (np.uint32(new_key_values[0]), np.uint32(new_key_values[1]))
values = jax.extend.random.threefry_2x32(new_key, np.arange(20, dtype=np.uint32))
values = (values >> (32 - 23)) | np.array(1, dtype=np.float32).view(np.uint32)
values = values.view(np.float32) - 1
print(f"[{', '.join(f'{v}' for v in values)}]")
"""
jr = np.array([0.45735931396484375, 0.6311527490615845, 0.15571284294128418, 0.8149417638778687, 0.7862188816070557,
0.8008807897567749, 0.568588376045227, 0.9852620363235474, 0.42314577102661133, 0.9811755418777466,
0.38059568405151367, 0.09186363220214844, 0.9497315883636475, 0.5826880931854248, 0.3796330690383911,
0.5610522031784058, 0.16122901439666748, 0.3732343912124634, 0.9795231819152832, 0.3280656337738037], dtype=np.float32)
r = Tensor.rand(20).numpy()
np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
# next 20 (c_low=20, c_high=0)
jr = np.array([0.09199333190917969, 0.9130761623382568, 0.7048608064651489, 0.22254979610443115, 0.0014830827713012695,
0.37023448944091797, 0.7790107727050781, 0.7484984397888184, 0.7524604797363281, 0.19875383377075195,
0.48537540435791016, 0.10002851486206055, 0.5369305610656738, 0.3294715881347656, 0.5246957540512085,
0.7659651041030884, 0.7949080467224121, 0.34988296031951904, 0.9798505306243896, 0.2599533796310425], dtype=np.float32)
r = Tensor.rand(20).numpy()
np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
# next 10 (c_low=40, c_high=0)
jr = np.array([0.3198714256286621, 0.7984923124313354, 0.320881724357605, 0.4716068506240845, 0.7323365211486816,
0.9663800001144409, 0.13873648643493652, 0.16062307357788086, 0.49300849437713623, 0.10077548027038574], dtype=np.float32)
r = Tensor.rand(10).numpy()
np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_threefry_tensors_cnt(self):
Tensor.manual_seed(1337)
Tensor.rand(20).realize()
assert len(Tensor._device_rng_counters) == 1
assert len(Tensor._device_seeds) == 1
Tensor.rand(20, device=f"{Device.DEFAULT}:1").realize()
assert len(Tensor._device_rng_counters) == 2
assert len(Tensor._device_seeds) == 2
Tensor.manual_seed(2)
assert len(Tensor._device_rng_counters) == 0
assert len(Tensor._device_seeds) == 0
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_threefry_same_kernels(self):
Tensor.manual_seed(0)
Tensor.rand(1).realize()
s = Tensor.rand(20).schedule_linear().src
s2 = Tensor.rand(20).schedule_linear().src
assert len(s) == len(s2), f"{len(s)} != {len(s2)}"
for x,y in zip(s, s2):
if not (x.src[0] == y.src[0]):
print(f"{x.src[0]} != {y.src[0]}")
Tensor.rand(1, device=f"{Device.DEFAULT}:1").realize()
s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
assert len(s3) == len(s4), f"{len(s3)} != {len(s4)}"
assert len(s2) == len(s4), f"{len(s)} != {len(s3)}"
for x,y in zip(s3, s4):
if not (x.src[0] == y.src[0]):
print(f"{x.src[0]} != {y.src[0]}")
@unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need bfloat16 support")
def test_rand_bfloat16(self):
N = 128
x = Tensor.rand((2, N, N), dtype=dtypes.bfloat16)
assert x.dtype == dtypes.bfloat16
nx = x.numpy()
assert nx[nx == 1].size == 0
assert nx[nx == 0].size > 0
equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.bfloat16).float(), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))
def test_rand_like(self):
empty = Tensor.empty((80, 44))
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_randn_like(self):
empty = Tensor.empty((80, 44))
rand = Tensor.randn_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_rand_like_zero_shape(self):
empty = Tensor.empty(0, 20)
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_rand_like_more_dims(self):
empty = Tensor.empty((1, 2, 3, 4, 5, 6))
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
def test_rand_like_dtype(self):
empty = Tensor.empty((80, 44), dtype=dtypes.float16)
rand = Tensor.rand_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
empty = Tensor.empty((80, 44))
rand = Tensor.rand_like(empty, dtype=dtypes.float16)
assert rand.shape == empty.shape
assert rand.dtype == dtypes.float16
assert rand.device == empty.device
def test_randn_like_dtype(self):
empty = Tensor.empty((80, 44), dtype=dtypes.float16)
rand = Tensor.randn_like(empty)
assert rand.shape == empty.shape
assert rand.dtype == empty.dtype
assert rand.device == empty.device
empty = Tensor.empty((80, 44))
rand = Tensor.randn_like(empty, dtype=dtypes.float16)
assert rand.shape == empty.shape
assert rand.dtype == dtypes.float16
assert rand.device == empty.device
def test_randn_device(self):
self.assertEqual(Tensor.randn(3,3,device="CPU").device, "CPU")
@given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
def test_randn_finite(self, default_float):
if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
old_default_float = dtypes.default_float
# low precision can result in inf from randn
dtypes.default_float = default_float
t = Tensor.randn(64, 64)
mx = t.max().numpy().item()
mn = t.min().numpy().item()
print(f"testing with {default_float=}")
assert math.isfinite(mx), mx
assert math.isfinite(mn), mn
dtypes.default_float = old_default_float
def test_random_counter_overflow(self):
device = Device.DEFAULT
Tensor.manual_seed(1337)
Tensor.rand(1).realize()
Tensor._device_rng_counters[device].assign(Tensor([dtypes.uint32.max - 5, 0], device=device, dtype=dtypes.uint32)).realize()
Tensor.rand(10).realize()
c = Tensor._device_rng_counters[device].numpy()
np.testing.assert_allclose(c, [4, 1])
Tensor.rand(10).realize()
c = Tensor._device_rng_counters[device].numpy()
np.testing.assert_allclose(c, [14, 1])
if __name__ == "__main__":
unittest.main()

222
tinygrad_repo/test/backend/test_rangeify.py Normal file
View File

@@ -0,0 +1,222 @@
import unittest
from tinygrad import Tensor, nn, Device
from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops
from tinygrad.codegen.opt import OptOps, Opt
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.nir import NIRRenderer
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
class TestDoubleMatmul(unittest.TestCase):
def setUp(self):
with Context(DEBUG=0):
self.a, self.b, self.c = [Tensor.randn(16, 16).contiguous().realize() for _ in range(3)]
self.ref = (self.a @ self.b @ self.c).realize()
def _test(self, opts):
with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()
with Context(DEBUG=0):
err = (out-self.ref).square()
self.assertLess(err.max().item(), 1e-4)
self.assertLess(err.mean().item(), 1e-6)
def test_baseline(self): self._test(())
def test_upcast_0(self): self._test((Opt(OptOps.UPCAST, 0, 4),))
def test_upcast_1(self): self._test((Opt(OptOps.UPCAST, 1, 4),))
def test_upcast_2(self): self._test((Opt(OptOps.UPCAST, 2, 4),))
def test_upcast_01(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4)))
def test_upcast_01_mismatch(self): self._test((Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 4)))
def test_upcast_02(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 2, 4)))
def test_upcast_12(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4)))
def test_unroll_0(self): self._test((Opt(OptOps.UNROLL, 0, 4),))
def test_unroll_1(self): self._test((Opt(OptOps.UNROLL, 1, 4),))
def test_unroll_01(self): self._test((Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_0_unroll_0(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)))
def test_upcast_1_unroll_0(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)))
def test_upcast_2_unroll_0(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4)))
def test_upcast_0_unroll_1(self): self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_1_unroll_1(self): self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_2_unroll_1(self): self._test((Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_1_unroll_1_small(self): self._test((Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UNROLL, 1, 2)))
def test_upcast_1_unroll_1_rev(self): self._test((Opt(OptOps.UNROLL, 1, 2), Opt(OptOps.UPCAST, 1, 2)))
def test_upcast_01_unroll_01(self):
self._test((Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
def test_upcast_12_unroll_01(self):
self._test((Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UPCAST, 2, 4), Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UNROLL, 1, 4)))
class TestRangeifyAssign(unittest.TestCase):
def test_assign_permuted(self):
A = Tensor.empty(4, 4, dtype='int')
B = Tensor.arange(16).reshape(4,4)
ret = A.permute(1,0).assign(B)
lst = ret.tolist()
lst2 = A.tolist()
lst3 = B.tolist()
print(lst)
print(lst2)
print(lst3)
self.assertListEqual(lst, lst3)
self.assertListEqual(lst2, B.permute(1, 0).tolist())
class TestRangeifyEdgeCase(unittest.TestCase):
def test_matmul_relu_cat(self):
a = Tensor.ones(100, 512).contiguous().realize()
c = Tensor.ones(1, 512).contiguous().realize()
cm = Tensor.ones(512, 512)
c = c @ cm
c = c.relu()
res = Tensor.cat(a, c, dim=0)
self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
def test_pcontig_multi_gather(self):
# regression test: local bufferize must have device set for const_like to work
with Context(PCONTIG=2):
# NOTE: with uint type, this will become a long and fail on WEBGPU
forest = Tensor(list(range(8)), dtype='int')
idx = Tensor([0, 0], dtype='int')
node_val = forest.gather(0, idx)
idx2 = idx * 2 + 1
node_val2 = forest.gather(0, idx2)
result = (node_val + node_val2).numpy()
self.assertEqual(result.tolist(), [1, 1])
if getenv("BIG") > 2:
# llama 8B (8192)
BS, HEADS, SEQLEN, EMB = 4, 32, 8192, 128
elif getenv("BIG") > 1:
# llama 8B
BS, HEADS, SEQLEN, EMB = 4, 32, 2048, 128
elif getenv("BIG") > 0:
# bigger
BS, HEADS, SEQLEN, EMB = 4, 32, 128, 128
else:
BS, HEADS, SEQLEN, EMB = 4, 2, 16, 8
def fa():
Tensor.manual_seed(1337)
with Context(DEBUG=0): q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
GlobalCounters.reset()
return q.scaled_dot_product_attention(k, v)
def fa_bw():
Tensor.manual_seed(1337)
with Context(DEBUG=0):
q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
attn_output.weight.realize()
target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
GlobalCounters.reset()
attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
out = attn_output(attn)
loss = (out - target).square().mean()
loss.backward()
#ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
#ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
ret = [out, q.grad, k.grad, v.grad]
Tensor.realize(*ret)
return ret
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
class TestPcontig(unittest.TestCase):
def test_flash_attention_bw(self):
with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
grads = fa_bw()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(PCONTIG=0, DEBUG=2):
cmp_grads = fa_bw()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(DEBUG=0):
mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
mse = sum(mses)
print(f"mse: {mse}")
self.assertLessEqual(mse, 1e-6)
def test_flash_attention(self, opts=None):
with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(DEBUG=2):
cmp = fa().realize()
print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
with Context(DEBUG=0):
mse = ((cmp-ret)**2).sum().item()
print(f"mse: {mse}")
self.assertLessEqual(mse, 1e-6)
def test_flash_attention_opt(self):
opts = ()
# columns in top matrix
opts += (Opt(OptOps.UPCAST, 0, 4),)
# columns in bottom matrix
opts += (Opt(OptOps.UPCAST, 3, 4),)
# rows in all the matrix
opts += (Opt(OptOps.UPCAST, 4, 4),)
self.test_flash_attention(opts)
# contiguous + reduce can support ranges?
@unittest.skip("pm_rangeify no longer exists. test this in a different way")
class TestRangeifyPM(unittest.TestCase):
def setUp(self): self.base = Tensor.empty(10*10).reshape(10, 10).contiguous()
def assert_same(self, a, b):
def run_pm_rangeify(t:Tensor):
from tinygrad.schedule.rangeify import pm_rangeify, RangeifyContext
sink = t.uop.sink()
pm_realize = PatternMatcher([(UPat(Ops.CONTIGUOUS, name="x"), lambda x: x.replace(op=Ops.REALIZE))])
sink = graph_rewrite(sink, pm_realize)
return graph_rewrite(sink, pm_rangeify, ctx=RangeifyContext())
self.assertIs(run_pm_rangeify(a.contiguous()), run_pm_rangeify(b.contiguous()))
def test_nothing_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.pad(((0,0),(0,1)))
self.assert_same(a, b)
def test_reshape_match(self):
a = self.base
b = self.base.reshape(100).reshape(10, 10)
self.assert_same(a, b)
def test_permute_reshape_match(self):
a = self.base
b = self.base.permute(1,0).reshape(100).reshape(10, 10).permute(1,0)
self.assert_same(a, b)
def test_padded_permute_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.permute(1,0).pad(((0,1),(0,0))).permute(1,0)
self.assert_same(a, b)
@unittest.expectedFailure
def test_padded_reshape_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.reshape(100).reshape(10, 10).pad(((0,0),(0,1)))
self.assert_same(a, b)
@unittest.expectedFailure
def test_padded_permute_reshape_match(self):
a = self.base.pad(((0,0),(0,1)))
b = self.base.permute(1,0).reshape(100).reshape(10, 10).pad(((0,1),(0,0))).permute(1,0)
self.assert_same(a, b)
# why is this failing?
@unittest.expectedFailure
def test_cross_pad_match(self):
a = self.base.pad(((0,0),(0,1))).pad(((0,1),(0,0)))
b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
self.assert_same(a, b)
if __name__ == '__main__':
unittest.main()

108
tinygrad_repo/test/backend/test_renderer_failures.py Normal file
View File

@@ -0,0 +1,108 @@
import unittest
import numpy as np
from tinygrad.device import Device
from tinygrad.dtype import dtypes, ConstType
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.helpers import prod
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.wgsl import WGSLRenderer
from tinygrad.runtime.ops_python import PythonRenderer
from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
from tinygrad.tensor import Tensor, _to_np_dtype
def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
for x in inputs: x.realize()
sz = 1 if local_size is None else prod(local_size)
outs = [UOp.new_buffer(Device.DEFAULT, sz, u.src[1].dtype) for u in sink.src if u.op is Ops.STORE]
for u in outs: u.buffer.allocate().copyin(np.zeros(sz, dtype=_to_np_dtype(u.dtype)).data)
run_linear(UOp(Ops.LINEAR, src=(sink.call(*outs, *(x.uop.base for x in inputs)),)))
return [u.buffer.numpy() for u in outs]
def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
dtype = alu_src_uops[0].dtype
a = UOp.param(0, dtype.ptr())
b = UOp.param(1, dtype.ptr())
idx = UOp.const(dtypes.int, 0)
ld = b.index(idx)
alu = ld.alu(alu_op, *alu_src_uops)
store = UOp.store(a.index(idx), alu)
return _test_uop_result([Tensor([input_val])], UOp(Ops.SINK, dtypes.void, (store,), arg=KernelInfo()))[0]
class TestRendererFailures(unittest.TestCase):
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
def test_gated_store_with_alu(self):
a = UOp.param(0, dtypes.int.ptr())
gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
np.testing.assert_equal(ret, [0, 1, 1, 1])
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
def test_gated_store_with_alu_2d(self):
a = UOp.param(0, dtypes.int.ptr())
gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage), "uops are for cstyle")
class TestCStyleFailures(unittest.TestCase):
def test_inline_const_alu(self):
# CPU doesn't use the max function
ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
self.assertEqual(ret[0], 1)
def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
ret = Tensor.empty(1, dtype=dtype)
for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
linear = ret.schedule_linear()
assert len(linear.src) == 1
src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[3].arg
self.assertEqual("("*5 not in src, should_strip_paren)
def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
def test_repeat_mul(self): self._test_src_strip_paren(Ops.MUL)
def test_repeat_xor(self): self._test_src_strip_paren(Ops.XOR)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
def test_repeat_or(self): self._test_src_strip_paren(Ops.OR)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
def test_repeat_and(self): self._test_src_strip_paren(Ops.AND)
def test_repeat_sub(self): self._test_src_strip_paren(Ops.SUB, should_strip_paren=False)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
class TestWGSLFailures(unittest.TestCase):
def test_multiply_infinity(self):
# multiplying a positive constant by infinity should return infinity
# WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(dtypes.float32, float("inf")))
self.assertEqual(ret[0], float("inf"))
@unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "tests for ptx renderer")
class TestPTXFailures(unittest.TestCase):
@unittest.skip("INDEX can only have a gate ALU parent, not an IF")
def test_gated_store_with_if(self):
a = UOp.param(0, dtypes.int.ptr())
gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
val = UOp.const(dtypes.int, 1)
if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0, if_uop), val))
sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
np.testing.assert_equal(ret, [0, 1, 1, 1])
@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
def test_gated_define_acc_with_half_dtype(self):
a = Tensor.randn(32, 32, dtype=dtypes.half).realize()
b = Tensor.randn(34, 32, dtype=dtypes.half).realize()
result = a.pad((1,1)).matmul(b, dtype=dtypes.half).numpy()
reference = a.pad((1,1)).matmul(b, dtype=dtypes.float).numpy()
np.testing.assert_allclose(result, reference, atol=1e-2, rtol=1e-2)
if __name__ == '__main__':
unittest.main()

1408
tinygrad_repo/test/backend/test_schedule.py Normal file
View File

File diff suppressed because it is too large Load Diff

378
tinygrad_repo/test/backend/test_setitem.py Normal file
View File

@@ -0,0 +1,378 @@
import unittest
from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
from tinygrad.helpers import Context
import numpy as np
class TestSetitem(unittest.TestCase):
def test_simple_setitem(self):
cases = (
((6,6), (slice(2,4), slice(3,5)), Tensor.ones(2,2)),
((6,6), (slice(2,4), slice(3,5)), Tensor([1.,2.])),
((6,6), (slice(2,4), slice(3,5)), 1.0),
((6,6), (3, 4), 1.0),
((6,6), (3, None, 4, None), 1.0),
((4,4,4,4), (Ellipsis, slice(1,3), slice(None)), Tensor(4.0)),
((4,4,4,4), (Ellipsis, slice(1,3)), 4),
((4,4,4,4), (2, slice(1,3), None, 1), 4),
((4,4,4,4), (slice(1,3), slice(None), slice(0,4,2)), 4),
((4,4,4,4), (slice(1,3), slice(None), slice(None), slice(0,3)), 4),
((6,6), (slice(1,5,2), slice(0,5,3)), 1.0),
((6,6), (slice(5,1,-2), slice(5,0,-3)), 1.0),
)
for shp, slc, val in cases:
t = Tensor.zeros(shp).contiguous()
t[slc] = val
n = np.zeros(shp)
n[slc] = val.numpy() if isinstance(val, Tensor) else val
np.testing.assert_allclose(t.numpy(), n)
def test_padded_setitem(self):
t = Tensor.arange(10)
t[4:1:-2] = 11
self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])
def test_setitem_inplace_mul(self):
t = Tensor.arange(10).clone().realize()
t[:3] *= 10
self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])
@unittest.skip("crashed in LLVM CI")
def test_setitem_fancy_on_unrealized_view(self):
# fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
sub = base[1:3]
flat = sub.reshape(sub.numel()).contiguous()
idx = Tensor([0, 3, 7, 9])
flat[idx] = Tensor([99, 98, 97, 96], dtype=dtypes.float)
sub.assign(flat.reshape(2, 5))
np.testing.assert_allclose(sub.numpy(), [[99, 6, 7, 98, 9], [10, 11, 97, 13, 96]])
def test_setitem_dtype(self):
for dt in (dtypes.int, dtypes.float, dtypes.bool):
for v in (5., 5, True):
t = Tensor.ones(6,6, dtype=dt).contiguous()
t[1] = v
self.assertEqual(t.dtype, dt)
def test_setitem_dtype_mismatch(self):
t = Tensor.zeros(6, dtype=dtypes.float).contiguous().realize()
with self.assertRaises(RuntimeError): t[2:4] = Tensor([1, 2], dtype=dtypes.int)
def test_setitem_chained_indexing(self):
# N[i][j] must work the same as N[i, j]
N1 = Tensor.zeros((3, 3)).contiguous().realize()
N1[1, 2] = 5
N2 = Tensor.zeros((3, 3)).contiguous().realize()
N2[1][2] = 5
np.testing.assert_equal(N1.numpy(), N2.numpy())
def test_setitem_detach(self):
# setitem on detached tensor should work
t = Tensor.zeros((3, 3)).contiguous().realize()
t.detach()[1, 2] = 5
self.assertEqual(t[1, 2].item(), 5.0)
def test_setitem_permute(self):
# setitem on permuted tensor should modify original
t = Tensor.zeros((2, 3)).contiguous().realize()
t.T[1, 0] = 5 # t.T is (3, 2), so [1, 0] maps to t[0, 1]
self.assertEqual(t[0, 1].item(), 5.0)
def test_setitem_flip(self):
# setitem on flipped tensor should modify original
t = Tensor.zeros((3,)).contiguous().realize()
t[::-1][0] = 5 # flip, then set first element (which is last in original)
self.assertEqual(t[2].item(), 5.0)
def test_setitem_inplace_operator(self):
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] += 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 5]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] -= 1
np.testing.assert_allclose(t.numpy(), [[0, 1], [1, 2]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] *= 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 6]])
# NOTE: have to manually cast setitem target to least_upper_float for div
t = Tensor.arange(4, dtype=dtypes.float).reshape(2, 2).contiguous()
t[1] /= 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [1, 1.5]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] **= 2
np.testing.assert_allclose(t.numpy(), [[0, 1], [4, 9]])
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] ^= 5
np.testing.assert_allclose(t.numpy(), [[0, 1], [7, 6]])
def test_setitem_consecutive_inplace_operator(self):
t = Tensor.arange(4).reshape(2, 2).contiguous()
t[1] += 2
t[1] -= 1
np.testing.assert_allclose(t.numpy(), [[0, 1], [3, 4]])
def test_setitem_overlapping_indices(self):
t = Tensor([1,2,3,4])
# regular overlapping indices
t[[1,1]] = Tensor([5,6])
np.testing.assert_allclose(t.numpy(), [1,6,3,4])
# overlapping indices with zero value overlapped
t[[1,1]] = Tensor([0,1])
np.testing.assert_allclose(t.numpy(), [1,1,3,4])
def test_setitem_overlapping_indices_with_0(self):
t = Tensor([1,2,3,4])
t[[1,1]] = Tensor([1,0])
np.testing.assert_allclose(t.numpy(), [1,0,3,4])
def test_setitem_with_1_in_shape(self):
t = Tensor([[1],[2],[3]])
t[[0,0]] = Tensor([[1],[2]])
np.testing.assert_allclose(t.numpy(), [[2],[2],[3]])
def test_fancy_setitem(self):
t = Tensor.zeros(6,6).contiguous()
t[[1,2], [3,2]] = 3
n = np.zeros((6,6))
n[[1,2], [3,2]] = 3
np.testing.assert_allclose(t.numpy(), n)
def test_simple_jit_setitem(self):
@TinyJit
def f(t:Tensor, a:Tensor):
t[2:4, 3:5] = a
# NOTE: without return t or an explicit realize, it's lazy and not captured
return t
for i in range(1, 6):
t = Tensor.zeros(6, 6).contiguous().realize()
a = Tensor.full((2, 2), fill_value=i, dtype=dtypes.float).contiguous()
f(t, a)
n = np.zeros((6, 6))
n[2:4, 3:5] = np.full((2, 2), i)
np.testing.assert_allclose(t.numpy(), n)
def test_jit_setitem_variable_offset(self):
with Context(CHECK_OOB=0):
@TinyJit
def f(t:Tensor, a:Tensor, v:Variable):
t.shrink(((v,v+1), None)).assign(a).realize()
t = Tensor.zeros(6, 6).contiguous().realize()
n = np.zeros((6, 6))
for i in range(6):
v = Variable("v", 0, 6).bind(i)
a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
n[i, :] = i+1
f(t, a, v)
np.testing.assert_allclose(t.numpy(), n)
np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
def test_setitem_overlapping_inplace1(self):
t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
t[1:] = t[:-1]
self.assertEqual(t.tolist(), [[3.0], [3.0], [2.0]])
def test_setitem_overlapping_inplace2(self):
t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
t[:-1] = t[1:]
self.assertEqual(t.tolist(), [[2.0], [1.0], [1.0]])
# TODO: WEBGPU pipeline validation error. this generates (1==gidx0)|(2==gidx0)|(3==gidx0)|(4==gidx0)|(5==gidx0) ...
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU pipeline validation error")
def test_setitem_big(self):
idx_size, val = 256, 4
t = Tensor.arange(0, idx_size+1)
idx = Tensor.arange(0, idx_size)
t[idx] = val
self.assertEqual(t.tolist(), [val]*idx_size+[idx_size])
def test_setitem_advanced_indexing(self):
# Example from https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
t = Tensor.zeros(10,20,30,40,50, dtype=dtypes.int).contiguous()
ind_1 = Tensor([5,3,7,8])
ind_2 = Tensor([[[0],[1],[2]],[[3],[4],[5]]])
v = Tensor.arange(2*3*4*10*30*50).reshape(2,3,4,10,30,50)
t[:, ind_1, :, ind_2, :] = v
n = np.zeros((10,20,30,40,50), dtype=np.int32)
n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
np.testing.assert_equal(t.numpy(), n)
def test_setitem_tensor_int_indexing(self):
t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
t[Tensor([0, 2]), 0] = Tensor([99, 88], dtype=dtypes.int)
n = np.zeros((4, 3), dtype=np.int32)
n[[0, 2], 0] = [99, 88]
np.testing.assert_equal(t.numpy(), n)
def test_setitem_tensor_slice_indexing(self):
t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
t[Tensor([0, 2]), :2] = Tensor([[10, 20], [30, 40]], dtype=dtypes.int)
n = np.zeros((4, 3), dtype=np.int32)
n[[0, 2], :2] = [[10, 20], [30, 40]]
np.testing.assert_equal(t.numpy(), n)
def test_setitem_2d_tensor_indexing(self):
t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
index = Tensor([[0, 1], [1,0]])
v = Tensor.arange(2*2).reshape(2, 2).contiguous()
t[index] = v
n = np.zeros((2,), dtype=np.int32)
n[index.numpy()] = v.numpy()
np.testing.assert_equal(t.numpy(), n)
def test_setitem_swap_rows(self):
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
tmp = t[0]
t[0] = t[1]
t[2] = tmp
# NOTE: not [[2, 3], [2, 3], [0, 1]], same with eager
np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
# eager version
t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
tmp = t[0].realize()
t[0] = t[1].realize()
t[2] = tmp.realize()
np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])
def test_lazy_sum_between_writes(self):
# lazy sums should capture buffer state at the time they were created
t = Tensor.zeros(6).contiguous().realize()
s0 = t.sum()
t[:3].assign(1.0)
s1 = t.sum()
t[3:].assign(2.0)
s2 = t.sum()
try:
np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
except AssertionError:
# TODO: broken now, lazy sums all see final buffer state
np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])
# eager version
t = Tensor.zeros(6).contiguous().realize()
s0 = t.sum().realize()
t[:3].assign(1.0).realize()
s1 = t.sum().realize()
t[3:].assign(2.0).realize()
s2 = t.sum().realize()
np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
def test_cross_assign_independence(self):
# when assigning to two tensors using computations from both,
# both assigns should see the OLD values of both tensors
a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
new_a = a + b # [4, 6, 8, 10]
new_b = a * 2 # [0, 2, 4, 6] -- should use OLD a
a.assign(new_a)
b.assign(new_b)
np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
try:
np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
except AssertionError:
# TODO: broken now, new_b sees mutated a
np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
# eager version
a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
new_a = (a + b).realize()
new_b = (a * 2).realize()
a.assign(new_a).realize()
b.assign(new_b).realize()
np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
def test_setitem_multiple_disjoint_on_invalid(self):
z = Tensor.invalids(10, dtype="int").realize()
z[2:5] = 2
z[6:7] = 3
z.realize()
self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
self.assertListEqual(z[6:7].tolist(), [3])
class TestWithGrad(unittest.TestCase):
def test_basic_setitem_works(self):
z = Tensor.rand(8, 8)
x = Tensor.rand(8)
z[:3] = x
def test_set_backward(self):
z = Tensor.ones(8, 8)
x = Tensor.rand(8, 8)
z[:] = x
z.sum().backward()
np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))
def test_set_nonleaf_backward(self):
x = Tensor([1.0, 2.0, 3.0, 4.0])
z = x * 2
z[:2] = Tensor([10.0, 20.0])
z.sum().backward()
np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])
def test_set_overlapping_backward(self):
z = Tensor.zeros(6)
x = Tensor.ones(4).contiguous()
y = Tensor.ones(4) * 2
z[:4] = x
z[2:] = y
z.sum().backward()
np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
np.testing.assert_allclose(y.grad.numpy(), np.ones(4))
def test_set_iadd_backward(self):
z = Tensor([1.0, 2.0, 3.0, 4.0])
x = Tensor([10.0, 20.0])
z[:2] += x
z.sum().backward()
np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
np.testing.assert_allclose(x.grad.numpy(), np.ones(2))
def test_set_used_before_setitem(self):
z = Tensor([1.0, 2.0, 3.0, 4.0])
_ = z.sum()
with self.assertRaises(RuntimeError):
z[:2] = Tensor([0.0, 0.0])
def test_setitem_raises_with_unrealized_downstream(self):
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
_y = x * 2.0
with self.assertRaises(RuntimeError):
x[0] = 99.0
def test_setitem_raises_on_unrealized_compute_base(self):
# y has a compute (unrealized) base; tmp is a view of y. eager: tmp would follow y's mutation. lazy: tmp keeps the old MUL graph.
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
y = x * 2.0
_tmp = y[:1]
with self.assertRaises(RuntimeError):
y[0] = 99.0
def test_setitem_raises_on_aliased_uop(self):
# two Tensor objects sharing the exact same unrealized uop. setitem on one updates its uop, the other keeps the stale graph reference.
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
y = x * 2.0
_z = Tensor(y.uop)
with self.assertRaises(RuntimeError):
y[0] = 99.0
class TestSetitemLoop(unittest.TestCase):
def test_arange(self):
N = 10
cmp = Tensor.empty(N)
for i in range(N): cmp[i] = i
self.assertListEqual(Tensor.arange(N).tolist(), cmp.tolist())
if __name__ == '__main__':
unittest.main()

206
tinygrad_repo/test/backend/test_softmax_fusion.py Normal file
View File

@@ -0,0 +1,206 @@
import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters, Context, Device
from tinygrad.dtype import DTypeLike, dtypes
from tinygrad.engine.realize import run_linear
from tinygrad.helpers import DEBUG, get_single_element
def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
# only support axis =-1
x = x_in.reshape(-1, x_in.shape[-1])
nr_dim, r_dim = x.shape
inp = x.reshape(nr_dim, 1, 1, r_dim).expand(nr_dim, r_dim, 1, r_dim)
imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, r_dim).max(axis=-2, keepdim=True)
m = inp - imx.detach()
if dtype is not None: m = m.cast(dtype)
e = m.exp()
ss = e.sum(axis=-1, keepdim=True)
inp = x.reshape(nr_dim, r_dim, 1, 1)
imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, 1).max(axis=-2, keepdim=True)
m = inp - imx.detach()
if dtype is not None: m = m.cast(dtype)
e = m.exp()
out = e.div(ss).reshape(x_in.shape)
return out
def run_one_schedule_item(out):
linear = out.schedule_linear()
get_single_element(linear.src)
run_linear(linear)
class TestFuse(unittest.TestCase):
def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
GlobalCounters.reset()
out_single = fxn(*args, **kwargs)
if not allow_multiple: run_one_schedule_item(out_single)
np_single = out_single.numpy()
GlobalCounters.reset()
np_multi = fxn(*args, **kwargs).numpy()
np.testing.assert_allclose(np_single, np_multi, atol=atol)
@unittest.skip("needs RANGEIFY>1")
def test_fuse_norm(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a / a.mean(axis=1), a)
@unittest.skip("needs RANGEIFY>1")
def test_fuse_argmax(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.argmax(axis=-1), a)
@unittest.skip("needs RANGEIFY>1")
def test_fuse_softmax(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.softmax(axis=-1), a)
def test_fuse_gemm_softmax(self):
a = Tensor.rand(50,50).realize()
b = Tensor.rand(50,50).realize()
self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)
@unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no float16 on {Device.DEFAULT}")
@unittest.skip("needs RANGEIFY>1")
def test_fuse_softmax_dtype(self):
a = Tensor.rand(50,50).realize()
self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)
def test_fuse_arange_eye(self):
self._test_fuse(lambda: Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10))
@unittest.skip("needs RANGEIFY>1")
def test_double_gemm(self):
N = 32
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
a = (Tensor.rand(N,N)-0.5).realize()
b = (Tensor.rand(N,N)-0.5).realize()
c = (Tensor.rand(N,N)-0.5).realize()
self._test_fuse(lambda a,b,c: a@b@c, a, b, c, atol=1e-5)
def test_embedding(self):
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
vocab_sz = 123
embed_sz = 16
weight = (Tensor.rand(vocab_sz, embed_sz)-0.5).realize()
a = Tensor([1, 1, 2, 3]).realize()
def embedding(idx:Tensor):
arange = Tensor.arange(vocab_sz).unsqueeze(-1)
big_shp = idx.shape + (vocab_sz, embed_sz)
arange, vals = arange.expand(big_shp), weight.expand(big_shp)
idx = idx.reshape(idx.shape+(1, 1)).expand(big_shp)
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
self._test_fuse(embedding, a, atol=1e-5)
@unittest.skip("needs RANGEIFY>1")
def test_attention_kernel_count(self):
wq = Tensor.empty(32, 32)
wk = Tensor.empty(32, 32)
wv = Tensor.empty(32, 32)
x = Tensor.empty(2, 100, 32)
q = (x @ wq).contiguous()
k = (x @ wk).contiguous()
v = (x @ wv).contiguous()
attn = q.scaled_dot_product_attention(k, v)
s = attn.schedule_linear()
self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention
@unittest.skip("needs RANGEIFY>1")
def test_flash_attention(self):
BS = 4
HEADS = 2
MATDIM = 16
EMB = 8
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
q = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
k = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
v = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
# TODO: OPT is breaking things. NOOPT isn't linearizing
with Context(NOOPT=1):
self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)
def test_mismatch_reduce(self):
a = Tensor.ones(16, 10).contiguous().realize()
b = Tensor.ones(16, 20).contiguous().realize()
c = (a.sum(axis=1) + b.sum(axis=1))
self.assertListEqual(c.tolist(), [30]*16)
@unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
def test_fuse_and_tc_opt(self):
A = Tensor.randn(8, 8).realize()
B = Tensor.randn(8, 8).realize()
C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
out = (C + (A @ B))
out.realize()
class TestSoftmaxFusion(unittest.TestCase):
@classmethod
def setUpClass(cls):
with Context(TRACK_MATCH_STATS=0): cls.test = Tensor.rand(32, 10).contiguous().realize()
def setUp(self):
GlobalCounters.reset()
def test_norm(self):
print("*** norm ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
# NOTE: there's an implied expand on the mean here
sout = self.test / self.test.mean(-1, keepdim=True)
sout.realize()
print("*** single kernel norm ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
inp = self.test.reshape(32, 10, 1)
div = self.test.reshape(32, 1, 10).expand(32, 10, 10).mean(axis=-1, keepdim=True)
out = (inp / div).reshape(32, 10)
out.realize()
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
def test_softmax(self):
# this is the softmax from scaled_dot_product_attention
# it becomes 3 kernels
print("*** softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
sout = self.test.softmax(-1)
sout.realize()
print("*** single kernel softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
out = single_kernel_softmax(self.test)
out.realize()
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
@unittest.skip("needs RANGEIFY>1")
def test_auto_softmax(self):
print("*** softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
sout = self.test.softmax(-1)
sout.realize()
print("*** auto single kernel softmax ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
out = self.test.contiguous().softmax(-1)
run_one_schedule_item(out)
np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)
def test_softmax_bw(self):
print("*** softmax bw ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
self.test.softmax(-1).sum().backward()
sg = self.test.grad.realize()
self.test.grad = None
print("*** single kernel softmax bw ***")
with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
single_kernel_softmax(self.test).sum().backward()
g = self.test.grad.realize()
np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)
if __name__ == '__main__':
unittest.main()

59
tinygrad_repo/test/backend/test_stunning.py Normal file
View File

@@ -0,0 +1,59 @@
import unittest
from tinygrad import nn, Tensor, Variable, Context, Device
from tinygrad.helpers import trange
class Model:
def __init__(self): self.layer = nn.Linear(28*28, 10)
def __call__(self, x:Tensor) -> Tensor: return self.layer(x.flatten(1))
class TestStunning(unittest.TestCase):
def test_indexing_variable(self):
a = Tensor.arange(100*10).reshape(100, 10).contiguous()
# index without variable
nv = a[12].tolist()
# index with variable
vi = Variable('i', 0, a.shape[0]-1)
wv = a[vi.bind(12)].tolist()
self.assertListEqual(nv, wv)
def test_indexing_two_bind(self):
a = Tensor.arange(100*10).reshape(100, 10).contiguous()
nv = a[12].cat(a[76]).tolist()
vi = Variable('i', 0, a.shape[0]-1)
with self.assertRaisesRegex(RuntimeError, "bind mismatch on"):
wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
self.assertListEqual(nv, wv)
@unittest.skipIf(Device.DEFAULT in {"WEBGPU", "NV", "CUDA"}, "Too many buffers / too slow")
@unittest.skip("This is binding a Variable to two different values")
def test_simple_train(self, steps=6, bs=4, adam=True):
X_train, Y_train, _, _ = nn.datasets.mnist()
model = Model()
if adam: opt = nn.optim.Adam(nn.state.get_parameters(model))
else: opt = nn.optim.SGD(nn.state.get_parameters(model), momentum=0.1)
samples = Tensor.randint(steps, bs, high=X_train.shape[0])
Y_train = Y_train.one_hot(10)
X_samp, Y_samp = X_train[samples], Y_train[samples]
vi = Variable('i', 0, samples.shape[0]-1)
with Context(SPLIT_REDUCEOP=0):
with Tensor.train():
losses = []
for i in range(samples.shape[0]):
vib = vi.bind(i)
opt.zero_grad()
pred = model(X_samp[vib].realize())
loss = (pred - Y_samp[vib]).square().mean()
losses.append(loss.backward())
opt.schedule_step()
#losses = Tensor.stack(*losses)
# run
for i in (t:=trange(len(losses))): t.set_description(f"loss: {losses[i].item():6.2f}")
if __name__ == '__main__':
unittest.main()

184
tinygrad_repo/test/backend/test_subbuffer.py Normal file
View File

@@ -0,0 +1,184 @@
import unittest
from tinygrad import Device, dtypes, Tensor
from tinygrad.device import Buffer
from tinygrad.helpers import Context, DEV
from test.helpers import needs_second_gpu
@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
class TestSubBuffer(unittest.TestCase):
def setUp(self):
self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
self.buf.copyin(memoryview(bytearray(range(10))))
self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)
def test_subbuffer(self):
vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
tst = vbuf.as_memoryview().tolist()
assert tst == [3, 4]
def test_subbuffer_cast(self):
# NOTE: bitcast depends on endianness
vbuf = self.buf.view(2, dtypes.uint16, offset=3).ensure_allocated()
tst = vbuf.as_memoryview().cast("H").tolist()
assert tst == [3|(4<<8), 5|(6<<8)]
def test_subbuffer_double(self):
vbuf = self.buf.view(4, dtypes.uint8, offset=3).ensure_allocated()
vvbuf = vbuf.view(2, dtypes.uint8, offset=1).ensure_allocated()
tst = vvbuf.as_memoryview().tolist()
assert tst == [4, 5]
def test_subbuffer_len(self):
vbuf = self.buf.view(5, dtypes.uint8, 2).ensure_allocated()
mv = vbuf.as_memoryview()
assert len(mv) == 5
mv = vbuf.as_memoryview(allow_zero_copy=True)
assert len(mv) == 5
def test_subbuffer_used(self):
t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
vt = t[2:4].realize()
out = (vt + 100).tolist()
assert out == [102, 103]
@needs_second_gpu
@unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
def test_subbuffer_transfer(self):
t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
vt = t[2:5].contiguous().realize()
out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
assert out == [2, 3, 4]
def test_subbuffer_deallocate(self):
with Context(LRU=0):
vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
self.buf.deallocate()
vbuf.deallocate()
# Allocate a fake one on the same place
_ = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
self.buf.ensure_allocated()
self.buf.copyin(memoryview(bytearray(range(10, 20))))
vbuf.ensure_allocated()
tst = vbuf.as_memoryview().tolist()
assert tst == [13, 14]
def test_subbuffer_is_allocated(self):
buf = self.buf_unalloc
sub_buf = buf.view(3, dtypes.uint8, offset=4)
self.assertFalse(buf.is_allocated())
self.assertFalse(buf.is_initialized())
self.assertFalse(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# base buffer alloc
buf.allocate()
self.assertTrue(buf.is_allocated())
self.assertTrue(buf.is_initialized())
self.assertTrue(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# sub buffer alloc
sub_buf.allocate()
self.assertTrue(sub_buf.is_initialized())
# sub buffer dealloc
sub_buf.deallocate()
self.assertTrue(buf.is_allocated())
self.assertTrue(buf.is_initialized())
self.assertTrue(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# base buffer dealloc
buf.deallocate()
self.assertFalse(buf.is_allocated())
self.assertFalse(buf.is_initialized())
self.assertFalse(sub_buf.is_allocated())
self.assertFalse(sub_buf.is_initialized())
# sub buffer alloc
sub_buf.ensure_allocated()
self.assertTrue(buf.is_allocated())
self.assertTrue(buf.is_initialized())
self.assertTrue(sub_buf.is_allocated())
self.assertTrue(sub_buf.is_initialized())
def test_subbuffer_copy_in_out(self):
sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
data_out_sub = bytearray([0]*3)
sub_buf.copyout(memoryview(data_out_sub))
assert data_out_sub == bytearray(range(3, 6))
sub_buf.copyin(memoryview(bytearray(range(3))))
assert sub_buf.as_memoryview().tolist() == list(range(3))
assert self.buf.as_memoryview().tolist()[3:6] == list(range(3))
sub_buf.copyout(memoryview(data_out_sub))
assert data_out_sub == bytearray(range(3))
data_out_base = bytearray([0]*10)
self.buf.copyout(memoryview(data_out_base))
assert data_out_base[0:3] == bytearray(range(0, 3))
assert data_out_base[3:6] == data_out_sub
assert data_out_base[6:10] == bytearray(range(6, 10))
def test_subbuffer_copy_in_out_view_of_view(self):
view1 = self.buf.view(7, dtypes.uint8, offset=2).ensure_allocated() # [2:9]
view2 = view1.view(3, dtypes.uint8, offset=2).ensure_allocated() # [4:7]
self.assertTrue(view1.is_allocated())
self.assertTrue(view2.is_allocated())
data_in = bytearray([7, 8, 9])
view2.copyin(memoryview(data_in))
data_out_v2 = bytearray([0]*3)
view2.copyout(memoryview(data_out_v2))
assert data_in == data_out_v2
expected_base_data = memoryview(bytearray(range(10)))
expected_base_data[4:7] = data_in
data_out_base = bytearray([0]*10)
self.buf.copyout(memoryview(data_out_base))
assert expected_base_data == data_out_base
def test_subbuffer_alloc(self):
sub_buf = self.buf.view(4, dtypes.int8, offset=3)
sub_buf.allocate()
sub_buf.copyin(memoryview(bytearray(range(10, 14))))
assert self.buf.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
sub_buf.allocate()
sub_buf.copyin(memoryview(bytearray(range(10, 14))))
assert self.buf_unalloc.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
def test_subbuffer_dealloc(self):
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
sub_buf.deallocate()
assert self.buf.as_memoryview().tolist() == list(range(10))
def test_subbuffer_double_dealloc(self):
sub_buf = self.buf.view(3, dtypes.uint8, offset=4).ensure_allocated()
self.buf.deallocate()
with self.assertRaises(AssertionError):
self.buf.deallocate()
sub_buf.deallocate()
with self.assertRaises(AssertionError):
sub_buf.deallocate()
def test_subbuffer_uaf(self):
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
assert self.buf.as_memoryview().tolist(), list(range(10))
sub_buf.deallocate()
with self.assertRaises(AssertionError):
sub_buf.as_memoryview().tolist()
assert self.buf.as_memoryview().tolist(), list(range(10))
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
assert sub_buf.as_memoryview().tolist(), list(range(3, 7))
self.buf.deallocate()
with self.assertRaises(AssertionError):
sub_buf.as_memoryview().tolist()
if __name__ == '__main__':
unittest.main()

332
tinygrad_repo/test/backend/test_symbolic_jit.py Normal file
View File

@@ -0,0 +1,332 @@
import unittest
from test.helpers import assert_jit_cache_len
from tinygrad import Variable, Tensor, TinyJit
from tinygrad.engine.jit import JitError
import numpy as np
class TestSymbolicJit(unittest.TestCase):
def test_plus1(self):
def f(a): return (a+1).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi])[:3, :i].numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_plus1_pad(self):
# TODO: without contiguous, the pad is not captured in jit
def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi]).numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_add(self):
def f(a, b): return (a+b).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b[:, :vi])
symbolic = symbolic[:3, :i].numpy()
expected = f(a[:, :i], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_matmul(self):
def f(a, b): return (a@b).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(10, 5)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
expected = f(a[:, :i], b[:i, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_mixed_with_no_symbol_kernel(self):
def f(a, b):
s = (a@b).realize()
s = (s+s).realize() # this one does not have symbols in input
return s
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(10, 5)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b[:vi, :]).numpy()
expected = f(a[:, :i], b[:i, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 2)
def test_attention(self):
def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).realize()
jf = TinyJit(f)
q = Tensor.rand(2, 1, 4, 8)
k = Tensor.rand(2, 10, 4, 8)
v = Tensor.rand(2, 10, 4, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy()
expected = f(q, k[:, :i], v[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 5)
def test_cat_dim0(self):
def f(a, b): return a.cat(b, dim=0).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(2, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:vi], b)[:i+2, :3].numpy()
expected = f(a[:i], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_cat_dim1(self):
def f(a, b): return a.cat(b, dim=1).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 2)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(a[:, :vi], b)[:3, :i+2].numpy()
expected = f(a[:, :i], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_cat_dim0_two_vars(self):
def f(a, b): return a.cat(b, dim=0).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:vi], b[:vj])[:i+j, :3].numpy()
expected = f(a[:i], b[:j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_cat_dim1_two_vars(self):
def f(a, b): return a.cat(b, dim=1).realize()
jf = TinyJit(f)
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
expected = f(a[:, :i], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_two_vars_plus1_ij(self):
def f(a, b): return (a@b+1).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:vi, :], b[:, :vj])[:i, :j].numpy()
expected = f(a[:i, :], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_two_vars_plus1_ji(self):
def f(a, b): return (a@b+1).realize()
jf = TinyJit(f)
a = Tensor.rand(10, 3)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = jf(a[:vj, :], b[:, :vi])[:j, :i].numpy()
expected = f(a[:j, :], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_jit_symbolic_shape_mismatch(self):
@TinyJit
def add(a, b): return (a+b).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
add(a[:, :vi], b[:, :vi])
vi2 = Variable("i", 1, 10).bind(7)
a = Tensor.rand(3, 7)[:, :vi2]
bad = Tensor.rand(4, 7)[:, :vi2]
with self.assertRaises(JitError):
add(a, bad)
def test_shrink(self):
# shrink is a movement, so we pair it with a simple function to test the JIT interaction
def f(a): return (a+1).realize()
jf = TinyJit(f)
a = Tensor.rand(7, 11)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = a.shrink(((3,5),(vi,vi+2)))
symbolic = jf(symbolic).numpy()
expected = f(a.shrink(((3,5),(i,i+2)))).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_slice(self):
# slice is a movement, so we pair it with a simple function to test the JIT interaction
def f(a): return (a+1).realize()
jf = TinyJit(f)
a = Tensor.rand(7, 11)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = a[3:5, vi:vi+2]
symbolic = jf(symbolic).numpy()
expected = f(a[3:5, i:i+2]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_slice_var_shape(self):
def f(a): return (a+1).realize()
jf = TinyJit(f)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.ones(vi, 11).contiguous()
symbolic = a[:, 1:2]
symbolic = jf(symbolic)[:i, :1].numpy()
expected = f(a[:i, :][:, 1:2]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_ones_sum(self):
def f(a): return a.sum().realize()
jf = TinyJit(f)
t = Tensor.ones(10).contiguous()
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = jf(t[:vi]).item()
expected = f(t[:i]).item()
np.testing.assert_equal(symbolic, expected)
def test_mean(self):
def f(a): return a.mean().realize()
def f0(a): return a.mean(0).realize()
def f1(a): return a.mean(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
c = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# axis = None
symbolic = jf(a[:vi]).numpy()
expected = a[:i].mean().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi]).numpy()
expected = b[:i].mean(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi])[:i].numpy()
expected = c[:i].mean(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_mean_2d(self):
def f(a): return a.mean().realize()
def f0(a): return a.mean(0).realize()
def f1(a): return a.mean(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 10)
b = Tensor.rand(10, 10)
c = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
# axis = None
symbolic = jf(a[:vi, :vj]).numpy()
expected = a[:i, :j].mean().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi, :vj])[:j].numpy()
expected = b[:i, :j].mean(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi, :vj])[:i].numpy()
expected = c[:i, :j].mean(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var(self):
def f(a): return a.var().realize()
def f0(a): return a.var(0).realize()
def f1(a): return a.var(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
c = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# axis = None
symbolic = jf(a[:vi]).numpy()
expected = a[:i].var().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi]).numpy()
expected = b[:i].var(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi])[:i].numpy()
expected = c[:i].var(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var_2d(self):
def f(a): return a.var().realize()
def f0(a): return a.var(0).realize()
def f1(a): return a.var(1).realize()
jf = TinyJit(f)
jf0 = TinyJit(f0)
jf1 = TinyJit(f1)
a = Tensor.rand(10, 10)
b = Tensor.rand(10, 10)
c = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
# axis = None
symbolic = jf(a[:vi, :vj]).numpy()
expected = a[:i, :j].var().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 0
symbolic = jf0(b[:vi, :vj])[:j].numpy()
expected = b[:i, :j].var(0).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
# axis = 1
symbolic = jf1(c[:vi, :vj])[:i].numpy()
expected = c[:i, :j].var(1).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
if __name__ == '__main__':
unittest.main()

358
tinygrad_repo/test/backend/test_symbolic_ops.py Normal file
View File

@@ -0,0 +1,358 @@
import unittest
from tinygrad import Tensor, Variable, GlobalCounters
from tinygrad.uop.ops import sym_infer
from tinygrad.dtype import dtypes
from examples.gpt2 import Attention
import numpy as np
class TestSymbolicOps(unittest.TestCase):
def test_plus1(self):
def f(a): return (a+1).realize()
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi])[:3, :i].numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_plus1_pad(self):
def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi]).numpy()
expected = f(a[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_add(self):
def f(a, b): return (a+b).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi], b[:, :vi])[:, :i].numpy()
expected = f(a[:, :i], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_matmul(self):
def f(a, b): return (a@b).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(10, 5)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = f(a[:, :vi], b[:vi, :]).numpy()
expected = f(a[:, :i], b[:i, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_attention(self, dropout_p=0.0, imin=1, imax=5, use_symbolic=True):
def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p).realize()
q = Tensor.rand(2, 1, 4, 8)
k = Tensor.rand(2, 10, 4, 8)
v = Tensor.rand(2, 10, 4, 8)
for i in range(imin, imax):
vi = Variable("i", 1, 10).bind(i) if use_symbolic else i
Tensor.realize(q, k, v)
GlobalCounters.reset()
symbolic = f(q, k[:, :vi, :, :], v[:, :vi, :, :])[:2, :4, :1, :8].numpy()
expected = f(q, k[:, :i, :, :], v[:, :i, :, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_attention_cmp_symbolic(self):
# symbolic isn't seeing if i == i, so it's not putting them on the same axis
self.test_attention(imin=4, imax=5, use_symbolic=False)
self.test_attention(imin=4, imax=5, use_symbolic=True)
def test_attention_training(self):
with Tensor.train():
self.test_attention(dropout_p=0.0)
with self.assertRaises(ValueError):
# symbolic shape dropout is not supported
self.test_attention(dropout_p=0.5)
def test_sdpa_symbolic_seq_len(self):
# symbolic seq_len on all of q/k/v (dim -2 after transpose)
q = Tensor.rand(2, 10, 4, 8)
k = Tensor.rand(2, 10, 4, 8)
v = Tensor.rand(2, 10, 4, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_sdpa_symbolic_seq_len_query_only(self):
# symbolic seq_len on query only (dim -2 after transpose)
q = Tensor.rand(2, 10, 4, 8)
k = Tensor.rand(2, 5, 4, 8)
v = Tensor.rand(2, 5, 4, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
k.transpose(1, 2), v.transpose(1, 2)).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_sdpa_symbolic_batch(self):
# symbolic batch dim (dim 0)
q = Tensor.rand(10, 4, 3, 8)
k = Tensor.rand(10, 4, 3, 8)
v = Tensor.rand(10, 4, 3, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy()
expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_sdpa_symbolic_heads(self):
# symbolic heads dim (dim -3)
q = Tensor.rand(2, 10, 3, 8)
k = Tensor.rand(2, 10, 3, 8)
v = Tensor.rand(2, 10, 3, 8)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
Tensor.realize(q, k, v)
symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy()
expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_attention_pos_0_sz_0(self):
Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None)
def test_attention_pos_0_sz_1(self):
Attention(128, 8)(Tensor.ones(1, 1, 128), Variable("start_pos", 0, 128).bind(0), None)
def test_attention_pos_0_sz_2(self):
Attention(128, 8)(Tensor.ones(1, 2, 128), Variable("start_pos", 0, 128).bind(0), None)
def test_cat_dim0(self):
def f(a, b): return a.cat(b, dim=0).realize()
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
b = Tensor.rand(2, 3)
symbolic = f(a[:vi, :], b)[:i+2, :3].numpy()
expected = f(a[:i, :], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_cat_dim1(self):
def f(a, b): return a.cat(b, dim=1).realize()
a = Tensor.rand(3, 10)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
b = Tensor.rand(3, 2)
symbolic = f(a[:, :vi], b)[:3, :i+2].numpy()
expected = f(a[:, :i], b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_cat_dim0_two_vars(self):
def f(a, b): return a.cat(b, dim=0).realize()
a = Tensor.rand(10, 3)
b = Tensor.rand(10, 3)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:vi, :], b[:vj, :])[:i+j, :3].numpy()
expected = f(a[:i, :], b[:j, :]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_cat_dim1_two_vars(self):
def f(a, b): return a.cat(b, dim=1).realize()
a = Tensor.rand(3, 10)
b = Tensor.rand(3, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:, :vi], b[:, :vj])[:3, :i+j].numpy()
expected = f(a[:, :i], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_two_vars_plus1_ij(self):
def f(a, b): return (a@b+1).realize()
a = Tensor.rand(10, 3).realize()
b = Tensor.rand(3, 10).realize()
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:vi, :], b[:, :vj])[:i, :j].numpy()
expected = f(a[:i, :], b[:, :j]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_two_vars_plus1_ji(self):
# reverse the order of variables
def f(a, b): return (a@b+1).realize()
a = Tensor.rand(10, 3).realize()
b = Tensor.rand(3, 10).realize()
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
symbolic = f(a[:vj, :], b[:, :vi])[:j, :i].numpy()
expected = f(a[:j, :], b[:, :i]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_invalid_symbolic_reshape(self):
a = Tensor.rand(30)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
# Cannot reshape into symbolic from non-symbolic
with self.assertRaises(ValueError): a.reshape((3, vi))
def test_shrink(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.rand(7, 11)
symbolic = a.shrink(((3,5),(vi,vi+2)))
symbolic = symbolic.numpy()
expected = a.shrink(((3,5),(i,i+2))).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_slice(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.rand(7, 11)
symbolic = a[3:5, vi:vi+2]
symbolic = symbolic.numpy()
expected = a[3:5, i:i+2].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_slice_no_start(self):
a = Tensor.rand(7, 11)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = a[3:5, :vi:1][:2, :i].numpy()
expected = a[3:5, :i:1].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_expand_padded(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor(1).unsqueeze(0).pad((0, 1)).unsqueeze(0)
symbolic = a.expand(vi, 2)[:i, :2].numpy()
expected = a.expand(i, 2).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_slice_var_shape(self):
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.ones(vi, 11).contiguous()
symbolic = a[:, 1:2][:i, :1].numpy()
expected = Tensor.ones(i, 11)[:, 1:2].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_ones_sum(self):
t = Tensor.ones(10).contiguous()
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
symbolic = t[:vi].sum().item()
expected = t[:i].sum().item()
np.testing.assert_equal(symbolic, expected)
def test_mean(self):
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
for axis in [None, 0, 1]:
expected = a[:i].mean(axis).numpy()
symbolic = a[:vi].mean(axis)
if axis is None:
symbolic = symbolic.numpy()
else:
symbolic = symbolic[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_mean_2d(self):
a = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
for axis in [None, 0, 1]:
expected = a[:i, :j].mean(axis).numpy()
symbolic = a[:vi, :vj].mean(axis)
if axis is None:
symbolic = symbolic.numpy()
else:
symbolic = symbolic[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var(self):
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
for axis in [None, 0, 1]:
expected = a[:i].var(axis).numpy()
symbolic = a[:vi].var(axis)
if axis is None:
symbolic = symbolic.numpy()
else:
symbolic = symbolic[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_var_2d(self):
a = Tensor.rand(10, 10)
for i in range(2, 5):
for j in range(2, 5):
vi = Variable("i", 1, 10).bind(i)
vj = Variable("j", 1, 10).bind(j)
for axis in [None, 0, 1]:
expected = a[:i, :j].var(axis).numpy()
symbolic_result = a[:vi, :vj].var(axis)
if axis is None:
symbolic = symbolic_result.numpy()
else:
symbolic = symbolic_result[:expected.shape[0]].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
def test_bitcast_down(self):
a = Tensor.rand(10, 3)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
expected = a[:i].bitcast(dtypes.uint8).numpy()
symbolic_result = a[:vi].bitcast(dtypes.uint8)
if len(expected.shape) == 2:
symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
else:
symbolic = symbolic_result[:].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
def test_bitcast_up(self):
a = Tensor.rand(10, 4)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
expected = a[:i].bitcast(dtypes.uint64).numpy()
symbolic_result = a[:vi].bitcast(dtypes.uint64)
if len(expected.shape) == 2:
symbolic = symbolic_result[:expected.shape[0], :expected.shape[1]].numpy()
else:
symbolic = symbolic_result[:].numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=0)
def test_conv2d_ceildiv_edge_case(self):
# tests symbolic ceildiv in conv2d output shape calculation
# val=79 triggers the edge case where old ceildiv simplifies incorrectly: old gives floor=12, correct ceildiv=13
v = Variable('v', 11, 100)
val = 79
x_full = Tensor.randn(1, 8, 100)
weight = Tensor.randn(16, 8, 12)
# symbolic version
result = x_full[:, :, :v.bind(val)].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
var_val = {v.expr: val}
shape = tuple(sym_infer(s, var_val) for s in result.shape)
self.assertEqual(shape, (1, 16, 13))
# concrete version for comparison
expected = x_full[:, :, :val].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
np.testing.assert_allclose(result[:, :, :13].numpy(), expected.numpy(), atol=1e-5, rtol=1e-5)
if __name__ == '__main__':
unittest.main()

776
tinygrad_repo/test/backend/test_tensor.py Normal file
View File

@@ -0,0 +1,776 @@
import numpy as np
import torch
import unittest, copy, mmap, random, math, array
from tinygrad import Tensor, Device, dtypes, nn
from tinygrad.helpers import getenv, temp, mv_address
from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
from hypothesis import given, settings, strategies as strat
from tinygrad.dtype import DTYPES_DICT
from tinygrad.uop.ops import UOp
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
x_init = np.random.randn(1,3).astype(np.float32)
U_init = np.random.randn(3,3).astype(np.float32)
V_init = np.random.randn(3,3).astype(np.float32)
W_init = np.random.randn(3,3).astype(np.float32)
m_init = np.random.randn(1,3).astype(np.float32)
gradient = np.random.randn(1,3).astype(np.float32)
class TestTinygrad(unittest.TestCase):
def test_zerodim_initialization(self):
self.assertEqual(Tensor(55).shape, ())
self.assertEqual(Tensor(3.14).shape, ())
def test_deviceless_const_construct_device_repr(self):
t = Tensor(UOp.const(dtypes.float, 2.0))
self.assertIsNone(t.uop.device)
self.assertIsNone(t.device)
self.assertIn("<UOp None", repr(t))
def test_deviceless_const_realize_noop(self):
t = Tensor(UOp.const(dtypes.float, 2.0))
uop = t.uop
t.realize()
self.assertIs(t.uop, uop)
self.assertIsNone(t.uop.device)
def test_plus_equals(self):
a = Tensor.randn(10,10)
b = Tensor.randn(10,10)
c = a + b
val1 = c.numpy()
a += b
val2 = a.numpy()
np.testing.assert_allclose(val1, val2)
def test_backward_pass(self):
def test_tinygrad():
x = Tensor(x_init)
W = Tensor(W_init)
m = Tensor(m_init)
out = x.dot(W).relu()
out = out.log_softmax()
out = out.mul(m).add(m).sum()
out.backward()
return out.numpy(), x.grad.numpy(), W.grad.numpy()
def test_pytorch():
x = torch.tensor(x_init, requires_grad=True)
W = torch.tensor(W_init, requires_grad=True)
m = torch.tensor(m_init)
out = x.matmul(W).relu()
out = torch.nn.functional.log_softmax(out, dim=1)
out = out.mul(m).add(m).sum()
out.backward()
return out.detach().numpy(), x.grad, W.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
# A simple test is to check that we can accumulate gradients (run backward twice or more times)
def test_accumulate_gradients(self):
x = Tensor(x_init)
W = Tensor(W_init)
m = Tensor(m_init)
out = x.dot(W).relu()
out = out.log_softmax()
out = out.mul(m).add(m).sum()
out.backward()
xgrad, wgrad = x.grad.numpy(), W.grad.numpy()
out.backward()
xgrad2, wgrad2 = x.grad.numpy(), W.grad.numpy()
out.backward() # no need to retain again since we will not re-run backward
xgrad3, wgrad3 = x.grad.numpy(), W.grad.numpy()
np.testing.assert_allclose(xgrad3, xgrad * 3., atol=1e-6)
np.testing.assert_allclose(wgrad3, wgrad * 3., atol=1e-6)
np.testing.assert_allclose(xgrad2, xgrad * 2., atol=1e-6)
np.testing.assert_allclose(wgrad2, wgrad * 2., atol=1e-6)
def test_second_order_backward_pass(self):
def test_pytorch():
x_val = torch.tensor([2.0], requires_grad=True)
f = x_val**3
first_derivative = torch.autograd.grad(outputs=f, inputs=x_val, create_graph=True)[0]
second_derivative = torch.autograd.grad(outputs=first_derivative, inputs=x_val)[0]
# d^2f/dx^2 = 6x = 6*2 = 12
return second_derivative.numpy()
def test_tinygrad():
x_val = Tensor([2.0])
f = x_val**3
first_derivative = f.sum().gradient(x_val)[0]
second_derivative = first_derivative.sum().gradient(x_val)[0]
return second_derivative.numpy()
np.testing.assert_allclose(test_tinygrad(), test_pytorch(), atol=1e-5)
# passing `gradient` to backward
def test_backward_pass_vjp(self):
def test_tinygrad():
x = Tensor(x_init)
W = Tensor(W_init)
m = Tensor(m_init)
out = x.dot(W).relu()
out = out.log_softmax()
out = out.mul(m).add(m)
out.backward(Tensor(gradient))
return out.numpy(), x.grad.numpy(), W.grad.numpy()
def test_pytorch():
x = torch.tensor(x_init, requires_grad=True)
W = torch.tensor(W_init, requires_grad=True)
m = torch.tensor(m_init)
out = x.matmul(W).relu()
out = torch.nn.functional.log_softmax(out, dim=1)
out = out.mul(m).add(m)
out.backward(torch.tensor(gradient))
return out.detach().numpy(), x.grad, W.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
def test_backward_pass_diamond_model(self):
def test_tinygrad():
u = Tensor(U_init)
v = Tensor(V_init)
w = Tensor(W_init)
x = u.mul(v).relu()
y = u.mul(w).relu()
out = x.add(y).mul(y).relu()
out = out.log_softmax()
out = out.sum()
out.backward()
return out.numpy(), u.grad.numpy(), v.grad.numpy(), w.grad.numpy()
def test_pytorch():
u = torch.tensor(U_init, requires_grad=True)
v = torch.tensor(V_init, requires_grad=True)
w = torch.tensor(W_init, requires_grad=True)
x = u.mul(v).relu()
y = u.mul(w).relu()
out = x.add(y).mul(y).relu()
out = torch.nn.functional.log_softmax(out, dim=1)
out = out.sum()
out.backward()
return out.detach().numpy(), u.grad, v.grad, w.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5, rtol=1e-6)
def test_const_backward_pass(self):
init = 3.5
def test_pytorch():
w1 = torch.tensor(init, requires_grad=True)
w2 = torch.tensor(init, requires_grad=True)
out = w1.add(w2)
out.backward()
return w1.grad, w2.grad
def test_tinygrad():
w1 = Tensor(init).clone()
w2 = Tensor(init).clone()
out = w1.add(w2)
out.backward()
return w1.grad.numpy(), w2.grad.numpy()
for x, y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
def test_const_backward_pass_optimizer(self):
init = 3.5
def test_pytorch():
w1 = torch.tensor(init, requires_grad=True)
w2 = torch.tensor(init, requires_grad=True)
out = w1.add(w2)
out.backward()
return w1.grad.numpy(), w2.grad.numpy()
def test_tinygrad():
w1 = Tensor(init).clone()
w2 = Tensor(init).clone()
assert w1.is_param is True and w2.is_param is True
nn.optim.SGD([w1, w2], lr=0.01)
assert w1.is_param is True and w2.is_param is True
out = w1.add(w2)
out.backward()
return w1.grad.numpy(), w2.grad.numpy()
for x, y in zip(test_tinygrad(), test_pytorch()):
np.testing.assert_allclose(x, y, atol=1e-5)
def test_dropout(self):
with Tensor.train():
n, rate = 1_000_000, 0.1
w = Tensor.ones(n).dropout(rate)
non_zeros = np.count_nonzero(w.numpy())
expected = n * (1 - rate)
np.testing.assert_allclose(non_zeros, expected, rtol=2e-3)
def test_jacobian(self):
W = np.random.RandomState(42069).random((10, 5)).astype(np.float32)
x = np.random.RandomState(69420).random((1, 10)).astype(np.float32)
torch_x = torch.tensor(x, requires_grad=True)
torch_W = torch.tensor(W, requires_grad=True)
def torch_func(x): return torch.nn.functional.log_softmax(x.matmul(torch_W).relu(), dim=1)
PJ = torch.autograd.functional.jacobian(torch_func, torch_x).squeeze().numpy()
tiny_x = Tensor(x)
tiny_W = Tensor(W)
def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
J = jacobian(tiny_func, tiny_x)
NJ = numerical_jacobian(tiny_func, tiny_x)
np.testing.assert_allclose(PJ, J, atol = 1e-5)
np.testing.assert_allclose(PJ, NJ, atol = 1e-3)
def test_gradcheck(self):
W = np.random.RandomState(1337).random((10, 5)).astype(np.float32)
x = np.random.RandomState(7331).random((1, 10)).astype(np.float32)
tiny_x = Tensor(x)
tiny_W = Tensor(W)
def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
self.assertTrue(gradcheck(tiny_func, tiny_x, eps = 1e-3))
# coarse approx. since a "big" eps and the non-linearities of the model
self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 1e-5))
def test_random_fns_are_deterministic_with_seed(self):
for random_fn in [Tensor.randn, Tensor.normal, Tensor.uniform, Tensor.scaled_uniform, Tensor.glorot_uniform, Tensor.kaiming_normal]:
with self.subTest(msg=f"Tensor.{random_fn.__name__}"):
Tensor.manual_seed(1337)
a = random_fn(10,10).realize()
Tensor.manual_seed(1337)
b = random_fn(10,10).realize()
np.testing.assert_allclose(a.numpy(), b.numpy())
def test_randperm(self):
Tensor.manual_seed(0)
a = Tensor.randperm(10).realize()
np.testing.assert_equal(a.numpy(), [8, 9, 4, 3, 6, 1, 7, 5, 2, 0])
b = Tensor.randperm(1000).realize()
np.testing.assert_equal(set(b.numpy()), set(range(1000)))
def test_rand_rejects_unknown_kwargs(self):
with self.assertRaises(TypeError): Tensor.rand(5, generator="foo")
def test_randn_isnt_inf_on_zero(self):
# simulate failure case of rand handing a zero to randn
original_rand, Tensor.rand = Tensor.rand, Tensor.zeros
try: self.assertNotIn(np.inf, Tensor.randn(16).numpy())
except: raise
finally: Tensor.rand = original_rand
def test_zeros_like_has_same_dtype_and_shape(self):
for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]:
a = Tensor([1, 2, 3], dtype=datatype)
b = Tensor.zeros_like(a)
assert a.dtype == b.dtype, f"dtype mismatch {a.dtype=} != {b.dtype}"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
a = Tensor([1, 2, 3])
b = Tensor.zeros_like(a, dtype=dtypes.int8)
assert a.dtype == dtypes.default_int and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
def test_ones_like_has_same_dtype_and_shape(self):
for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]:
a = Tensor([1, 2, 3], dtype=datatype)
b = Tensor.ones_like(a)
assert a.dtype == b.dtype, f"dtype mismatch {a.dtype=} != {b.dtype}"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
a = Tensor([1, 2, 3])
b = Tensor.ones_like(a, dtype=dtypes.int8)
assert a.dtype == dtypes.default_int and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char"
assert a.shape == b.shape, f"shape mismatch {a.shape} != {b.shape}"
def test_rand_like_device(self):
a = Tensor.ones(3, 3, device="CPU")
b = Tensor.rand_like(a)
self.assertEqual(b.device, a.device)
def test_ndim(self):
assert Tensor(1).ndim == 0
assert Tensor.randn(1).ndim == 1
assert Tensor.randn(2,2,2).ndim == 3
assert Tensor.randn(1,1,1,1,1,1).ndim == 6
def test_argfix(self):
for f in [Tensor.zeros, Tensor.ones, Tensor.rand, Tensor.randn, Tensor.empty]:
self.assertEqual(f().shape, ())
self.assertEqual(f(1).shape, (1,))
self.assertEqual(f(10,20,40).shape, (10,20,40))
self.assertEqual(f([]).shape, ())
self.assertEqual(f([1]).shape, (1,))
self.assertEqual(f([10,20,40]).shape, (10,20,40))
self.assertEqual(f(()).shape, ())
self.assertEqual(f((1,)).shape, (1,))
self.assertEqual(f((10,20,40)).shape, (10,20,40))
with self.assertRaises(ValueError): f((2, 2), 2, 2)
with self.assertRaises(ValueError): f((2, 2), (2, 2))
with self.assertRaises(ValueError): f((128, 128), 0.0, 0.01)
def test_numel(self):
assert Tensor.randn(10, 10).numel() == 100
assert Tensor.randn(1,2,5).numel() == 10
assert Tensor.randn(1,1,1,1,1,1).numel() == 1
assert Tensor([]).numel() == 0
assert Tensor.randn(1,0,2,5).numel() == 0
assert Tensor(3).numel() == 1
def test_len(self):
assert len(torch.zeros(7)) == len(Tensor.zeros(7))
assert len(torch.zeros(10,20)) == len(Tensor.zeros(10,20))
assert len(torch.zeros(10,20)) == len(Tensor.zeros(10,20,30))
assert len(torch.zeros(1).flatten()) == len(Tensor.zeros(1).flatten())
with self.assertRaises(TypeError): len(Tensor(3))
def test_size(self):
t1, t2 = torch.zeros(10,20), Tensor.zeros(10,20)
assert t1.size() == t2.size()
assert t1.size(0) == t2.size(0)
assert t1.size(1) == t2.size(1)
assert t1.size(-1) == t2.size(-1)
assert t1.size(-2) == t2.size(-2)
with self.assertRaises(IndexError): t2.size(2)
def test_tolist(self):
# NOTE: float16 Tensor.tolist() requires python 3.12
for arr in [[1,2,3], [1.5,2,3], [[1,2,3], [4,5,6]], 3]:
assert Tensor(arr).tolist() == torch.tensor(arr).tolist() == arr
def test_element_size(self):
for _, dtype in DTYPES_DICT.items():
assert dtype.itemsize == Tensor.randn(3, dtype=dtype).element_size(), f"Tensor.element_size() not matching Tensor.dtype.itemsize for {dtype}"
def test_deepwalk_ctx_check(self):
layer = Tensor.uniform(1, 1)
x = Tensor.randn(1, 1, 1)
x.dot(layer).mean().backward()
x = Tensor.randn(1, 1, 1)
x.dot(layer).mean().backward()
def test_zerosized_tensors(self):
np.testing.assert_equal(Tensor([]).numpy(), np.array([]))
np.testing.assert_equal(Tensor(None).numpy(), np.array([]))
def test_tensor_ndarray_dtype(self):
arr = np.array([1]) # where dtype is implicitly int64
assert Tensor(arr).dtype == dtypes.int64
assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32 # check if ndarray correctly casts to Tensor dtype
assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64 # check that it works for something else
def test_tensor_from_blob(self):
x = memoryview(bytearray(16)).cast('I')
t = Tensor.from_blob(mv_address(x), (4,), dtype=dtypes.int, device="CPU")
z = (t+1)
np.testing.assert_equal(z.numpy(), [1, 1, 1, 1])
x[:] = array.array('I', [0, 1, 2, 3])
z = (t+1)
np.testing.assert_equal(z.numpy(), [1, 2, 3, 4])
def test_tensor_list_dtype(self):
for arr in ([1], [[[1]]], [[1,1],[1,1]], [[[1,1],[1,1]],[[1,1],[1,1]]]):
assert Tensor(arr).dtype == dtypes.default_int
assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32
assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64
for arr in ([True], [[[False]]], [[True,False],[True,False]], [[[False,True],[False,False]],[[True,True],[False,True]]]):
assert Tensor(arr).dtype == dtypes.bool
assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32
assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64
# empty tensor defaults
for arr in ([], [[[]]], [[],[]]):
t = Tensor(arr)
assert t.dtype == dtypes.default_float
np.testing.assert_allclose(t.numpy(), np.array(arr))
# mixture of bool and int
for arr in ([True, 3], [[True],[3]], [[[True]], [[3]]], [[True, 3], [3, True]]):
t = Tensor(arr)
assert t.dtype == dtypes.default_int
np.testing.assert_allclose(t.numpy(), np.array(arr))
# mixture of bool, int and float
for arr in ([[True,True],[3.,True]], [[0,1],[3.,4]], [[[0],[1]],[[3.],[4]]], [[[True],[1]],[[3.],[4]]]):
t = Tensor(arr)
assert t.dtype == dtypes.default_float
np.testing.assert_allclose(t.numpy(), np.array(arr))
def test_tensor_list_shapes(self):
self.assertEqual(Tensor([[[]]]).shape, (1,1,0))
self.assertEqual(Tensor([[],[]]).shape, (2,0))
self.assertEqual(Tensor([[[[]],[[]]], [[[]],[[]]], [[[]],[[]]]]).shape, (3,2,1,0))
def test_tensor_list_errors(self):
# inhomogeneous shape
with self.assertRaises(ValueError): Tensor([[],[[]]])
with self.assertRaises(ValueError): Tensor([[1],[]])
with self.assertRaises(ValueError): Tensor([[1],[1],1])
with self.assertRaises(ValueError): Tensor([[[1,1,1],[1,1]]])
with self.assertRaises(ValueError): Tensor([[1,1,1],[[1,1,1]]])
def test_tensor_mixed_list_tuple(self):
def _list_or_tuple(): return list if random.random() < 0.5 else tuple
def _generate_data(depth):
if depth == 0: return _list_or_tuple()()
if depth == 1: return _list_or_tuple()([random.random(), random.random()])
return _list_or_tuple()([_generate_data(depth-1), _generate_data(depth-1)])
for depth in range(7):
for _ in range(20):
data = _generate_data(depth)
np.testing.assert_allclose(Tensor(data).numpy(), np.array(data))
def test_tensor_list_implicit_cast(self):
data = [True, False]
np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
data = [-1, 0, 1, 2, 3]
np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
data = [-3.5, -2.5, -1.5, 0, 1.5, 2.5, 3.5]
np.testing.assert_equal(Tensor(data, dtype=dtypes.int).numpy(), torch.tensor(data, dtype=torch.int).numpy())
# NOTE: torch and jax raise OverflowError: Python integer -3 out of bounds for uint8
# np.testing.assert_equal(Tensor(data, dtype=dtypes.uint8).numpy(), torch.tensor(data, dtype=torch.uint8).numpy())
np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())
def test_tensor_list_special_values(self):
if dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes():
data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]
data = data + [-x for x in data]
with np.errstate(over='ignore'): np.testing.assert_allclose(Tensor(data, dtype=dtypes.float16).numpy(), np.array(data).astype(np.float16))
# uint32
data = [1 << 33, 1 << 32, 1 << 32 - 1, 1]
data = data + [-x for x in data]
np.testing.assert_allclose(Tensor(data, dtype=dtypes.uint32).numpy(), np.array(data).astype(np.uint32))
# int32
data = [1 << 33, 1 << 32, 1 << 32 - 1, 1]
data = data + [-x for x in data]
np.testing.assert_allclose(Tensor(data, dtype=dtypes.int32).numpy(), np.array(data).astype(np.int32))
def test_tensor_list_ndarray(self):
data = [np.array([1, 2, 3]), np.array([1, 2, 3]), np.array([1, 2, 3])]
np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
data = [np.array([1.0, 2.0, 3.0]), np.array([1, 2, 3]), np.array([1, 2, 3])]
np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
data = [np.array(1.0), np.array(2.0), np.array(3.0)]
np.testing.assert_equal(Tensor(data).numpy(), np.array(data))
def test_tensor_dtype_errors(self):
with self.assertRaises(AttributeError): Tensor([3], dtype="typo")
with self.assertRaises(AttributeError): Tensor([3], dtype=(dtypes.int,))
def test_tensor_bytes(self):
data = b"abc123"
t = Tensor(data)
assert t.dtype == dtypes.uint8
assert t.shape == (6,)
np.testing.assert_equal(t.numpy(), list(data))
def test_tensor_copy(self):
x = copy.deepcopy(Tensor.ones((3,3,3)))
np.testing.assert_allclose(x.numpy(), np.ones((3,3,3)))
def test_copy_from_disk(self):
t = Tensor.randn(30).to(f"disk:{temp('test_copy_from_disk')}")
a = t[10:20]
dev = a.to(Device.DEFAULT)
np.testing.assert_allclose(a.numpy(), dev.numpy())
def test_copy_from_numpy_dtype(self):
data = np.array([1.0, 2, 3], dtype=np.float32)
t = Tensor(data, dtype=dtypes.bfloat16)
try:
# TODO: fix dtype in tinygrad space
assert t.dtype == dtypes.bfloat16
except AssertionError:
assert t.dtype == dtypes.float32
np.testing.assert_equal(t.tolist(), data)
np.testing.assert_equal((t+1).tolist(), data+1)
# Regression test for https://github.com/tinygrad/tinygrad/issues/1751
def test_copy_from_numpy_unaligned(self):
# 2**15 is the minimum for repro
arr = np.random.randn(2**15).astype(np.float32)
fn = temp('test_copy_from_numpy_unaligned')
with open(fn, 'wb') as f: f.write(b't' + arr.tobytes())
with open(fn, "a+b") as f: memview = memoryview(mmap.mmap(f.fileno(), arr.nbytes + 1))
ua_arr = np.frombuffer(memview[1:], dtype=arr.dtype, count=arr.shape[0])
np.testing.assert_allclose(arr, ua_arr)
assert not ua_arr.flags.aligned
# force device copy - to() is opt'd away - Tensor(dev)/1 is ignored
np.testing.assert_allclose(ua_arr, (Tensor(ua_arr)/Tensor(1)).numpy())
def test_item_to_tensor_to_item(self):
for a in [0, 1, 2, 3, -1, -100, 100, -101.1, 2.345, 100.1, True, False]:
item = Tensor(a).item()
assert type(item) is type(a), a
np.testing.assert_allclose(item, a), a
buffered_item = Tensor([a]).item()
assert type(buffered_item) is type(a), a
np.testing.assert_allclose(buffered_item, a), a
reshaped_item = Tensor([a]).reshape((1, 1, 1, 1, 1)).item()
assert type(reshaped_item) is type(a), a
np.testing.assert_allclose(reshaped_item, a), a
def test_no_bool(self):
with self.assertRaises(TypeError):
if Tensor(3):
print("hi")
with self.assertRaises(TypeError):
_a = Tensor([3]) in [Tensor([3]), Tensor([4]), Tensor([5])]
def test_repr_with_grad(self):
a = Tensor([1.0])
b = Tensor([1])
c = (a + b).sum().backward()
print(a)
print(c)
def test_no_attributeerror_after_apply_uop_exception(self):
try:
Tensor.arange(4).reshape(3,2)
except ValueError:
Tensor.zeros(2, 2).realize()
def test_shrink(self):
t = Tensor.arange(32).clone().realize()
self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
t = t.reshape(4, 8).contiguous().realize()
self.assertListEqual(t.shrink_to(2, 2).tolist(), [[0, 1], [8, 9]])
self.assertListEqual(t.shrink_to(None, 2).tolist(), t.shrink_to(4, 2).tolist())
with self.assertRaises(ValueError): t.shrink_to(2)
with self.assertRaises(ValueError): t.shrink_to(2, 2, 2)
@unittest.skip("this test is just flaky, sync issue")
class TestMoveTensor(unittest.TestCase):
d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
@given(strat.sampled_from([d0, d1]), strat.sampled_from([d0, d1]),
strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
def test_to_preserves(self, src, dest, dtype, is_param):
if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes():
return
s = Tensor([1, 2, 3], device=src, dtype=dtype).is_param_(is_param)
if is_param: s.sum().backward()
t = s.to(dest)
np.testing.assert_equal(s.numpy(), t.numpy())
assert s.dtype == t.dtype
assert s.is_param == t.is_param
if is_param:
np.testing.assert_equal(s.grad.numpy(), t.grad.numpy())
@given(strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
def test_shard_preserves(self, dtype, is_param):
s = Tensor([1, 2, 3], dtype=dtype).is_param_(is_param)
t = s.shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"))
np.testing.assert_equal(s.numpy(), t.numpy())
assert s.dtype == t.dtype
assert s.is_param == t.is_param
@given(strat.sampled_from([d0, d1]))
def test_same_dev(self, dev):
x = Tensor([1,2,3], device=dev)
y = x.to(dev)
assert x is y
def test_to_grad(self):
x = Tensor.eye(3, device=self.d0)
y = Tensor([[2.0,0,-2.0]], device=self.d0)
z = y.matmul(x).to(self.d1).sum()
z.backward()
np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])
class TestZeroShapeTensor(unittest.TestCase):
def test_rand(self):
t = Tensor.rand(3, 2, 0)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
t = Tensor.rand(0)
assert t.shape == (0,)
np.testing.assert_equal(t.numpy(), np.zeros((0,)))
t = Tensor.rand(0, 0, 0)
assert t.shape == (0, 0, 0)
np.testing.assert_equal(t.numpy(), np.zeros((0, 0, 0)))
def test_full(self):
t = Tensor.zeros(3, 2, 0)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
t = Tensor.full((3, 2, 0), 12)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.full((3, 2, 0), 12))
def test_reshape(self):
t = Tensor.zeros(3, 2, 0)
a = t.reshape(7, 0)
assert a.shape == (7, 0)
np.testing.assert_equal(a.numpy(), np.zeros((7, 0)))
a = t.reshape(0)
assert a.shape == (0,)
np.testing.assert_equal(a.numpy(), np.zeros((0,)))
with self.assertRaises(ValueError):
# cannot reshape from size 0 to size 1
a = t.reshape(())
def test_expand(self):
t = Tensor.full((1, 2, 0), 12).expand((6, 2, 0))
assert t.shape == (6, 2, 0)
np.testing.assert_equal(t.numpy(), np.full((6, 2, 0), 12))
def test_pad(self):
t = Tensor.rand(3, 2, 0).pad((None, None, (1, 1)), value=1)
self.assertEqual(t.shape, (3, 2, 2))
np.testing.assert_equal(t.numpy(), np.ones((3, 2, 2)))
t = Tensor.rand(3, 2, 0).pad((None, (1, 1), None), value=1)
self.assertEqual(t.shape, (3, 4, 0))
np.testing.assert_equal(t.numpy(), np.ones((3, 4, 0)))
t = Tensor.rand(3, 2, 0).pad(((1, 1), None, None), value=1)
self.assertEqual(t.shape, (5, 2, 0))
np.testing.assert_equal(t.numpy(), np.ones((5, 2, 0)))
np.testing.assert_equal(Tensor([1, 2]).pad_to(4).numpy(), [1, 2, 0, 0])
np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
def test_shrink_into_zero(self):
t = Tensor.rand(3, 4).realize()
assert t.shrink((None, (2, 2))).realize().shape == (3, 0)
assert t.shrink(((2, 2), None)).realize().shape == (0, 4)
assert t.shrink(((2, 2), (2, 2))).realize().shape == (0, 0)
def test_cat(self):
a = Tensor.rand(3, 2, 2)
b = Tensor.rand(3, 2, 0)
t = a.cat(b, dim=2)
assert t.shape == (3, 2, 2)
np.testing.assert_equal(t.numpy(), a.numpy())
t = b.cat(a, dim=2)
assert t.shape == (3, 2, 2)
np.testing.assert_equal(t.numpy(), a.numpy())
t = b.cat(b, dim=0)
assert t.shape == (6, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((6, 2, 0)))
t = b.cat(b, dim=1)
assert t.shape == (3, 4, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 4, 0)))
t = b.cat(b, dim=2)
assert t.shape == (3, 2, 0)
np.testing.assert_equal(t.numpy(), np.zeros((3, 2, 0)))
def test_elementwise(self):
a = Tensor.rand(3, 2, 0)
a_exp = a.exp()
assert a_exp.shape == (3, 2, 0)
np.testing.assert_equal(a_exp.numpy(), np.exp(a.numpy()))
b = Tensor.rand(3, 2, 0)
assert b.shape == (3, 2, 0)
ab = a * b
assert ab.shape == (3, 2, 0)
np.testing.assert_equal(ab.numpy(), a.numpy() * b.numpy())
mask = (Tensor.rand(3, 2, 0) > 0.5)
assert mask.shape == (3, 2, 0)
c = mask.where(a, b)
assert c.shape == (3, 2, 0)
np.testing.assert_equal(c.numpy(), np.where(mask.numpy(), a.numpy(), b.numpy()))
def test_reduce_over_non_zero(self):
a = Tensor.ones(3, 2, 0).sum(axis=1)
assert a.shape == (3, 0)
np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=1))
def test_reduce_over_zero(self):
a = Tensor.ones(3, 2, 0).sum(axis=2)
assert a.shape == (3, 2)
np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=2))
a = Tensor.ones(3, 2, 0).sum(axis=2, keepdim=True)
assert a.shape == (3, 2, 1)
np.testing.assert_equal(a.numpy(), np.sum(np.zeros((3, 2, 0)), axis=2, keepdims=True))
def test_clone(self):
a = Tensor.rand(16, 16).realize()
b = a.clone()
np.testing.assert_allclose(a.numpy(), b.numpy())
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
a = Tensor.rand(16, 16).mul(5.0).add(5.0).realize()
b = a.clone()
np.testing.assert_allclose(a.numpy(), b.numpy())
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
def test_clone_deviceless_const(self):
t = Tensor(UOp.const(dtypes.float, 2.0)).clone()
np.testing.assert_equal(t.numpy(), 2.0)
self.assertTrue(t.uop.has_buffer_identity())
def test_numpy_deviceless_const(self):
np.testing.assert_equal(Tensor(UOp.const(dtypes.float, 2.0)).numpy(), 2.0)
def test_clone_with_shrink(self):
a = Tensor.rand(16, 16)
b = a.shrink(((2, 10), None)).clone()
b.realize()
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
def test_clone_with_shrink_realized(self):
a = Tensor.rand(16, 16).realize()
b = a.shrink(((2, 10), None)).clone()
b.realize()
self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
def test_clone_with_grad(self):
a = Tensor.rand(16, 16)
a.mul(5.0).add(5.0).mean().backward()
b = a.clone()
assert a.grad is not None
assert b.grad is not None
np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())
def test_clone_deviceless_const_to_cpu(self):
t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
self.assertEqual(t.device, "CPU")
np.testing.assert_equal(t.numpy(), 2.0)
def test_reduce_default(self):
np.testing.assert_equal(Tensor([]).max().numpy(), -float("inf"))
np.testing.assert_equal(Tensor([]).min().numpy(), float("inf"))
np.testing.assert_equal(Tensor([]).sum().numpy(), 0)
np.testing.assert_equal(Tensor([]).mean().numpy(), float("nan"))
class TestTensorCreationDevice(unittest.TestCase):
# test auxiliary tensors are created on the same device
def test_one_hot(self):
y = Tensor([1, 2, 3]).to("CPU")
x = y.one_hot(10)
x.realize()
if __name__ == '__main__':
unittest.main()

165
tinygrad_repo/test/backend/test_tensor_variable.py Normal file
View File

@@ -0,0 +1,165 @@
import unittest
import numpy as np
from tinygrad import Tensor, Variable
class TestTensorVariable(unittest.TestCase):
def test_add_tvar(self):
vv = Variable("a", 0, 10).bind(1)
ret = (Tensor(vv) + 3).item()
assert ret == 4
def test_inner_tvar_node(self):
vv = Variable("w", 0, 10).bind(2)
ret = Tensor(vv * 4).item()
assert ret == 8
def test_inner_tvar_mul(self):
vv = Variable("w", 0, 10).bind(2)
assert (Tensor(3) * vv).item() == 6
def test_inner_tvar_mul_node(self):
vv = Variable("w", 0, 10).bind(2)
assert (Tensor(3) * (vv * 4)).item() == 24
def test_symbolic_mean(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 10).contiguous()[:, :vv]
ret = t.mean().item()
assert ret == 1
def test_symbolic_mean_2d(self):
vv = Variable("a", 1, 10).bind(2)
vv2 = Variable("b", 1, 10).bind(2)
t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
ret = t.mean().item()
assert ret == 1
def test_symbolic_mean_2d_axis_1(self):
vv = Variable("a", 1, 10).bind(2)
vv2 = Variable("b", 1, 10).bind(2)
t = Tensor.ones(10, 10).contiguous()[:vv2, :vv]
ret = t.mean(axis=1)[:2].reshape(2, 1).numpy()
assert np.all(ret == 1)
def test_symbolic_mean_2d_add(self):
add_term = Variable("c", 0, 10).bind(1)
vv = Variable("a", 1, 10).bind(1)
vv2 = Variable("b", 1, 10).bind(1)
t = Tensor.ones(20, 20).contiguous()[:vv2+add_term, :vv+add_term]
ret = t.mean().item()
assert ret == 1
def test_symbolic_var(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 10).contiguous()[:, :vv]
ret = t.var().item()
assert ret == 0
def test_symbolic_pad(self):
vv = Variable("a", 1, 10).bind(2)
t = Tensor.ones(2, 2).contiguous()
t = t.pad([vv, vv, vv, vv]).mean()
ones = 4
zeros = 6+6+4+4+6+6
self.assertAlmostEqual(t.item(), ones/(ones+zeros))
def test_symbolic_arange(self):
vv = Variable("a", 1, 10)
ret = Tensor.arange(0, vv.bind(4))
self.assertListEqual(ret[:4].tolist(), [0,1,2,3])
def test_symbolic_arange_sym_start(self):
vv = Variable("a", 1, 6)
ret = Tensor.arange(vv.bind(4), 7)
self.assertListEqual(ret[:3].tolist(), [4,5,6])
def test_symbolic_arange_sym_step(self):
vv = Variable("step", 1, 3)
ret = Tensor.arange(0, 10, vv.bind(2))
self.assertListEqual(ret[:5].tolist(), [0,2,4,6,8])
def test_symbolic_arange_two_vars(self):
begin = Variable("b", 1, 5)
end = Variable("e", 6, 10)
ret = Tensor.arange(begin.bind(4), end.bind(7))
self.assertListEqual(ret[:3].tolist(), [4,5,6])
def test_symbolic_arange_three_vars(self):
begin = Variable("b", 0, 5)
end = Variable("e", 10, 20)
step = Variable("s", 1, 3)
ret = Tensor.arange(begin.bind(2), end.bind(14), step.bind(3))
self.assertListEqual(ret[:4].tolist(), [2,5,8,11])
def test_symbolic_full(self):
vv = Variable("x", 1, 10).bind(5)
t = Tensor.full((3,), vv)
self.assertListEqual(t.tolist(), [5,5,5])
def test_variable_empty(self):
v = Variable("i", 1, 10)
# TODO: Tensor creation from unbound variable should assert
# with self.assertRaises(AssertionError): t = Tensor.empty(3, v)
vb = v.bind(3)
t = Tensor.empty(3, vb)
assert t.uop.base.buffer.size == 30
assert t.uop.shape == (3, vb)
def test_symbolic_chunk(self):
# chunk should work when split dimension is concrete, even if other dims are symbolic
vv = Variable("a", 1, 10).bind(4)
t = Tensor.ones(10, 8).contiguous()[:vv, :] # shape (vv, 8)
chunks = t.chunk(2, dim=-1) # split along concrete dim 8
assert len(chunks) == 2
assert chunks[0].shape[1] == 4
assert chunks[1].shape[1] == 4
# verify the values by shrinking to concrete shape first
np.testing.assert_equal(chunks[0].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
np.testing.assert_equal(chunks[1].shrink(((0, 4), (0, 4))).numpy(), np.ones((4, 4)))
def test_symbolic_split(self):
# split should work when split dimension is concrete, even if other dims are symbolic
vv = Variable("a", 1, 10).bind(3)
t = Tensor.arange(30).reshape(10, 3).contiguous()[:, :vv] # shape (10, vv)
splits = t.split(5, dim=0) # split along concrete dim 10
assert len(splits) == 2
assert splits[0].shape[0] == 5
assert splits[1].shape[0] == 5
# verify the values by shrinking to concrete shape first
np.testing.assert_equal(splits[0].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[:5, :3])
np.testing.assert_equal(splits[1].shrink(((0, 5), (0, 3))).numpy(), np.arange(30).reshape(10, 3)[5:, :3])
def test_symbolic_chunk_error_on_symbolic_dim(self):
# chunk should fail when trying to split along a symbolic dimension
vv = Variable("a", 1, 10).bind(4)
t = Tensor.ones(10, 8).contiguous()[:vv, :] # shape (vv, 8)
with self.assertRaises(AssertionError):
t.chunk(2, dim=0) # can't split along symbolic dim
def test_symbolic_var_sum(self, var_name="u"):
t = Variable("t", 1, 10).bind(4)
v = Variable(var_name, 1, 5).bind(1)
mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous()
mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy()
def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s")
def test_symbolic_triu(self):
t = Variable("t", 1, 10).bind(4)
for start_pos in (0, 1, 3):
var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1)
out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
np.testing.assert_equal(out, expected)
def test_symbolic_tril(self):
t = Variable("t", 1, 10).bind(4)
for start_pos in (0, 1, 3):
var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1)
out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
np.testing.assert_equal(out, expected)
if __name__ == '__main__':
unittest.main()

17
tinygrad_repo/test/backend/test_to_numpy.py Normal file
View File

@@ -0,0 +1,17 @@
from tinygrad.tensor import Tensor
import numpy as np
import pickle
import unittest
class TestToNumpy(unittest.TestCase):
def test_numpy_is_numpy(self):
output = Tensor.ones((1, 3, 4096)).realize().numpy()
new = np.copy(output)
print(type(new))
serialized = pickle.dumps(new)
out = pickle.loads(serialized)
assert out.shape == (1,3,4096)
assert (out==1).all()
if __name__ == '__main__':
unittest.main()

205
tinygrad_repo/test/backend/test_transcendental.py Normal file
View File

@@ -0,0 +1,205 @@
import unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _to_np_dtype
from tinygrad.helpers import Context, getenv, DEV, OSX
from test.backend.test_schedule import check_schedule
from test.backend.test_dtype_alu import ht, dtypes_float
import numpy as np
import math
from hypothesis import given, settings, strategies as strat
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
class TestTranscendentalMath(unittest.TestCase):
@unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
@given(ht.float64, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)]))
def test_float64(self, x, op):
if op[0] == Tensor.sin:
# TODO: reduction does not work # 536870912.125 # 2914593.01171875 # 134217728.03125 # 230581075.65625 # 139216373.71875
if abs(x) > 100_000_000: return
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float64)).numpy(),
op[1](np.array([x], dtype=_to_np_dtype(dtypes.float64))),
atol=3e-2, rtol=1e-5) # sin can have bigger atol for very big x
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
@given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else [])))
def test_float32(self, x, op):
# wrong nan behavior on Vulkan
if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float32)).numpy(),
op[1](np.array([x], dtype=_to_np_dtype(dtypes.float32))),
atol=2e-5, rtol=1e-5)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@given(ht.float16, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
([(Tensor.sin, np.sin)] if dtypes.ulong in supported_dtypes else [])))
def test_float16(self, x, op):
# wrong nan behavior on Vulkan
if (math.isnan(x) or (x < 0 and op[0] == Tensor.log)) and Device.DEFAULT == "WEBGPU" and not OSX: return
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
np.testing.assert_allclose(op[0](Tensor([x], dtype=dtypes.float16)).numpy(),
op[1](np.array([x], dtype=_to_np_dtype(dtypes.float16))),
atol=1e-2, rtol=5e-3) # exp can have bigger rtol
# TODO: WEBGPU produces incorrect values near infinity
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU incorrect values near inf")
@given(strat.sampled_from([(dtypes.float64, 709.5), (dtypes.float32, 88.7), (dtypes.float16, 11)]))
def test_exp_near_inf(self, dtype_x):
# reordering compute might return inf
dtype, x = dtype_x
if dtype not in supported_dtypes: return
with Context(TRANSCENDENTAL=2):
y = Tensor([x], dtype=dtype).exp().numpy()
expected = np.exp(np.array([x], dtype=_to_np_dtype(dtype)))
np.testing.assert_allclose(y, expected, rtol=5e-3)
class TestFromFuzzer(unittest.TestCase):
@given(strat.sampled_from(dtypes_float))
@unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong")
def test_sin(self, dtype):
if dtype not in supported_dtypes: return
if dtype == dtypes.float64:
# crashes in CI CUDA
if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
def _test_value(n: float, unit: float=1.0):
next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
ulp = next_float - 1.0
ulp = unit * ulp
with Context(TRANSCENDENTAL=2):
np.testing.assert_allclose(Tensor([n], dtype=dtype).sin().numpy(), np.sin(np.array([n], dtype=_to_np_dtype(dtype))), atol=ulp, rtol=1e-5)
_test_value(-35.0)
_test_value(-25.0)
_test_value(25.0)
_test_value(30.0) # 30.0 == switch_over
_test_value(35.0)
_test_value(0.0)
_test_value(np.pi / 2)
# worst case of ulp 1.5
_test_value(np.pi * 2, unit=1.5)
@given(strat.sampled_from(dtypes_float))
def test_log2(self, dtype):
if dtype not in supported_dtypes: return
if dtype == dtypes.float64:
# crashes in CI CUDA
if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
def _test_value(n: float, unit: float=1.0):
next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
ulp = next_float - 1.0
ulp = unit * ulp
with Context(TRANSCENDENTAL=2):
np.testing.assert_allclose(Tensor([n], dtype=dtype).log2().numpy(), np.log2(np.array([n], dtype=_to_np_dtype(dtype))), atol=ulp, rtol=1e-5)
fmin = np.finfo(_to_np_dtype(dtype)).tiny
for scale in [1.0, 1e10, 1e20, 1e30]:
_test_value(fmin * scale)
_test_value(-fmin * scale)
_test_value(0)
_test_value(0.0000009)
class TestFloat16Log2(unittest.TestCase):
"""Tests for native float16 log2 implementation (no float32 cast)"""
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
def test_float16_log2_basic(self):
# basic values
test_values = [1.0, 2.0, 4.0, 0.5, 0.25, 10.0, 100.0, 1000.0]
with Context(TRANSCENDENTAL=2):
for val in test_values:
result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
expected = np.log2(np.float16(val))
np.testing.assert_allclose(result, expected, rtol=1e-3, err_msg=f"log2({val})")
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "Nan handling differs on Vulkan")
def test_float16_log2_special(self):
# special values: inf, -inf, nan, 0, negative
with Context(TRANSCENDENTAL=2), np.errstate(all='ignore'):
# log2(inf) = inf
assert np.isinf(Tensor([np.inf], dtype=dtypes.float16).log2().numpy()[0])
# log2(0) = -inf
assert Tensor([0.0], dtype=dtypes.float16).log2().numpy()[0] == -np.inf
# log2(negative) = nan
assert np.isnan(Tensor([-1.0], dtype=dtypes.float16).log2().numpy()[0])
# log2(nan) = nan
assert np.isnan(Tensor([np.nan], dtype=dtypes.float16).log2().numpy()[0])
@unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
def test_float16_log2_denormal(self):
# test values near and below float16 min normal (6.1e-5)
# these exercise the denormal handling path with 2^10 scaling
test_values = [1e-4, 6e-5, 1e-5]
with Context(TRANSCENDENTAL=2):
for val in test_values:
result = Tensor([val], dtype=dtypes.float16).log2().numpy()[0]
expected = np.log2(np.float16(val))
# denormals have lower precision due to float16 limitations
np.testing.assert_allclose(result, expected, rtol=5e-2, err_msg=f"log2({val})")
class TestTranscendentalSchedule(unittest.TestCase):
@unittest.skipUnless(dtypes.ulong in supported_dtypes, "Needs ulong")
def test_transcendental_sin_fusion(self):
with Context(TRANSCENDENTAL=2):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = a.sin() + b.sin()
c = c.sin()
check_schedule(c, 1)
def test_transcendental_log2_fusion(self):
with Context(TRANSCENDENTAL=2):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = a.log2() + b.log2()
c = c.log2()
check_schedule(c, 1)
def test_transcendental_exp2_fusion(self):
with Context(TRANSCENDENTAL=2):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = a.exp2() + b.exp2()
c = c.exp2()
check_schedule(c, 1)
class TestTranscendentalVectorized(unittest.TestCase):
def _vectorized_data(self, low, high, vec_size):
np_data = np.linspace(low, high, num=(128 // vec_size) * vec_size, dtype=np.float32).reshape(-1, vec_size)
data = Tensor(np_data, dtype=dtypes.float32.vec(vec_size))
return data, np_data
def _test_vectorized_op(self, fxn, np_fxn, data_range, vec_size, param_range=None):
data, np_data = self._vectorized_data(data_range[0], data_range[1], vec_size)
if param_range:
param, np_param = self._vectorized_data(param_range[0], param_range[1], vec_size)
out, np_out = fxn(data, param), np_fxn(np_data, np_param)
else:
out, np_out = fxn(data), np_fxn(np_data)
np.testing.assert_allclose(out.numpy(), np_out, rtol=1e-4)
def test_exp2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.exp2, np.exp2, (-100, 100), vec_size)
def test_log2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)
@unittest.skipIf(Device.DEFAULT == "DSP", "requires int division")
@unittest.skipIf(DEV.renderer == "NAK", "MUFU.SIN is not accurate enough")
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and OSX, "WEBGPU Metal backend is not accurate enough")
def test_sin_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)
def test_pow_vectorized(self):
# np.pow returns nan for negative values raised to a non-integral power
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.pow, np.pow, (0.001, 200), vec_size, param_range=(-10, 10))
def test_sqrt_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sqrt, np.sqrt, (0, 100), vec_size)
if __name__ == '__main__':
unittest.main()

346
tinygrad_repo/test/backend/test_uops.py Normal file
View File

@@ -0,0 +1,346 @@
from typing import Optional, Any
import unittest, math
import numpy as np
from tinygrad.tensor import Tensor, _to_np_dtype
from tinygrad.helpers import Context
from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat # noqa: F401
from tinygrad.device import Buffer, Device
from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.engine.realize import run_linear
from tinygrad.codegen import to_program
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import to_uops_list
def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
for u,b in zip(buf_uops, bufs): buffers[u] = b
run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))
def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype).replace(src=()))
else: uops.append(UOp(op, dtype, tuple(src), arg))
return uops[-1]
def _test_single_value(vals, op, dts):
uops = []
output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
buf_loads = [uop(uops, Ops.PARAM, dtype.ptr(), (), i+1) for i,dtype in enumerate(dts)]
loads = (buf_loads[i].index(uop(uops, Ops.CONST, dtypes.int32, (), 0)) for i, dtype in enumerate(dts))
alu = uop(uops, op, output_dtype, loads)
out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True), alu))
buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
buf2 = [Buffer(Device.DEFAULT, 1, dtype).allocate().copyin(np.array([a], dtype=_to_np_dtype(dtype)).data) for a,dtype in zip(vals, dts)]
run_uops([out], [buf]+buf2)
ret = np.empty(1, _to_np_dtype(output_dtype))
buf.copyout(ret.data)
return ret[0]
def _test_single_value_const(vals, op, dts):
uops = []
output_dtype = dtypes.bool if op in (Ops.CMPLT, Ops.CMPNE) else dts[-1]
buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
alu = uop(uops, op, output_dtype, loads)
out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
run_uops([out], [buf])
ret = np.empty(1, _to_np_dtype(output_dtype))
buf.copyout(ret.data)
return ret[0]
def _test_uops_result(output_dtype, uops, res):
# uops = []
buf_store = uop(uops, Ops.PARAM, output_dtype.ptr(), (), 0)
# res = output_fn(uops)
out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
run_uops([out], [buf])
ret = np.empty(1, _to_np_dtype(output_dtype))
buf.copyout(ret.data)
return ret[0]
class TestUOps(unittest.TestCase):
def _equal(self, v1, v2):
assert isinstance(v2, (float, int, bool))
if isinstance(v2, float):
np.testing.assert_allclose(v1, v2, rtol=2e-7)
else:
np.testing.assert_equal(v1, v2)
def _test_uop_fxn(self, op, fxn, dts=(dtypes.float32, )):
for f in [_test_single_value, _test_single_value_const]:
for a in [-2.0, 0.0, 1.0]:
a = dts[0].const(a)
self._equal(f([a], op, dts), fxn(a))
def _test_bop_fxn(self, op, fxn, dts=(dtypes.float32, )*2, no_b_zero=False, no_b_neg=False):
for f in [_test_single_value, _test_single_value_const]:
for a in [-2.0, 0.0, 1.0]:
for b in [-3.0, 1.0] + ([] if no_b_zero else [0.0]):
a = dts[0].const(a)
b = dts[1].const(abs(b) if no_b_neg else b)
self._equal(f([a,b], op, dts), fxn(a,b))
def _test_top_fxn(self, op, fxn, dts=(dtypes.float32, )*3):
for f in [_test_single_value, _test_single_value_const]:
for a in [-2.0, 0, 1]:
for b in [-3.0, 3.0]:
for c in [-4.0, 4.0]:
a = dts[0].const(a)
b = dts[1].const(b)
c = dts[2].const(c)
self._equal(f([a,b,c], op, dts), fxn(a,b,c))
class TestFloatUOps(TestUOps):
@unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
def test_exp2(self): self._test_uop_fxn(Ops.EXP2, lambda a: np.exp2(a))
@unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
def test_log2(self): self._test_uop_fxn(Ops.LOG2, lambda a: math.log2(a) if a > 0 else float('-inf' if a==0 else 'nan'))
@unittest.skipIf(Device.DEFAULT == "CPU", 'not supported as uop')
def test_sin(self): self._test_uop_fxn(Ops.SIN, lambda a: math.sin(a))
def test_recip(self): self._test_uop_fxn(Ops.RECIPROCAL, lambda a: 1/a if a != 0 else float('inf'))
def test_sqrt(self): self._test_uop_fxn(Ops.SQRT, lambda a: math.sqrt(a) if a >= 0 else float('nan'))
def test_add(self): self._test_bop_fxn(Ops.ADD, lambda a,b: a+b)
def test_mul(self): self._test_bop_fxn(Ops.MUL, lambda a,b: a*b)
def test_max(self): self._test_bop_fxn(Ops.MAX, lambda a,b: max(a,b))
def test_cmplt(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: a<b)
def test_cmpne(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: a!=b)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
def test_cmpne_nan(self): # NaN != x for any x (IEEE 754)
for a, b in [(math.nan, 1.0), (1.0, math.nan), (math.nan, math.nan)]:
self.assertTrue(_test_single_value(
[dtypes.float32.const(a), dtypes.float32.const(b)],
Ops.CMPNE, (dtypes.float32, dtypes.float32)))
# MOD isn't tested on floats
def test_where(self):
self._test_top_fxn(Ops.WHERE, lambda a,b,c: b if a!=0 else c, (dtypes.bool, dtypes.float, dtypes.float))
@unittest.skipUnless(Device.DEFAULT == "PYTHON", "only python supports MULACC")
def test_mulacc(self):
self._test_top_fxn(Ops.MULACC, lambda a,b,c: a*b+c, (dtypes.float, dtypes.float, dtypes.float))
class TestNonFloatUOps(TestUOps):
def test_add_int32(self): self._test_bop_fxn(Ops.ADD, lambda a,b: int(a)+int(b), (dtypes.int32, dtypes.int32))
def test_mul_int32(self): self._test_bop_fxn(Ops.MUL, lambda a,b: int(a)*int(b), (dtypes.int32, dtypes.int32))
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
def test_shr_int32(self): self._test_bop_fxn(Ops.SHR, lambda a,b: int(a)>>int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CStyleLanguage)), "only ptx and cstyle use bitshifts")
def test_shl_int32(self): self._test_bop_fxn(Ops.SHL, lambda a,b: int(a)<<int(b), (dtypes.int32, dtypes.int32), no_b_neg=True)
def test_div_int32(self):
self._test_bop_fxn(Ops.CDIV, lambda a,b: int(a/b), (dtypes.int32, dtypes.int32), no_b_zero=True)
def test_and_int32(self): self._test_bop_fxn(Ops.AND, lambda a,b: int(a)&int(b), (dtypes.int32, dtypes.int32))
def test_or_int32(self): self._test_bop_fxn(Ops.OR, lambda a,b: int(a)|int(b), (dtypes.int32, dtypes.int32))
def test_mod_int32(self):
self._test_bop_fxn(Ops.CMOD,
lambda a,b: abs(int(a))%abs(int(b))*(1,-1)[a<0], (dtypes.int32, dtypes.int32), no_b_zero=True)
def test_cmplt_int32(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: int(a)<int(b), (dtypes.int32, dtypes.int32))
def test_cmpne_int32(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: int(a)!=int(b), (dtypes.int32, dtypes.int32))
def test_mul_bool(self): self._test_bop_fxn(Ops.MUL, lambda a,b: bool(a) and bool(b), (dtypes.bool, dtypes.bool))
def test_where_float16(self):
self._test_top_fxn(Ops.WHERE, lambda a,b,c: b if a!=0 else c, (dtypes.bool, dtypes.float16, dtypes.float16))
class TestBoolUOps(TestUOps):
def _test_uop_bool_fxn(self, op, fxn):
for f in [_test_single_value, _test_single_value_const]:
for a in [False, True]:
self._equal(f([a], op, (dtypes.bool, )*1), fxn(a))
def _test_bop_bool_fxn(self, op, fxn):
for f in [_test_single_value, _test_single_value_const]:
for a in [False, True]:
for b in [False, True]:
self._equal(f([a,b], op, (dtypes.bool, )*2), fxn(a,b))
def _test_top_bool_fxn(self, op, fxn):
for f in [_test_single_value, _test_single_value_const]:
for a in [False, True]:
for b in [False, True]:
for c in [False, True]:
self._equal(f([a,b,c], op, (dtypes.bool, )*3), fxn(a,b,c))
def test_add_bool(self): self._test_bop_bool_fxn(Ops.ADD, lambda a,b: a or b)
def test_mul_bool(self): self._test_bop_bool_fxn(Ops.MUL, lambda a,b: a and b)
def test_xor_bool(self): self._test_bop_bool_fxn(Ops.XOR, lambda a,b: a != b)
def test_and_bool(self): self._test_bop_bool_fxn(Ops.AND, lambda a,b: a & b)
def test_or_bool(self): self._test_bop_bool_fxn(Ops.OR, lambda a,b: a | b)
def test_cmpne_bool(self): self._test_bop_bool_fxn(Ops.CMPNE, lambda a,b: a != b)
def test_cmplt_bool(self): self._test_bop_bool_fxn(Ops.CMPLT, lambda a,b: a < b)
def test_where_bool(self): self._test_top_bool_fxn(Ops.WHERE, lambda a,b,c: b if a else c)
class TestLocalAccess(unittest.TestCase):
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
def test_local_basic(self):
uops = []
smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.float32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.float32, (), 42.0)))
barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
sres = uop(uops, Ops.LOAD, dtypes.float32, (smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0), ptr=True),))
self.assertEqual(_test_uops_result(dtypes.float32, uops, sres), 42)
# NOTE: webgpu specific, since only webgpu performs bitpacking
@unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local access with packed data type")
def test_local_packed(self):
uops = []
smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
st = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), uop(uops, Ops.CONST, dtypes.uint8, (), 42)))
barr = uop(uops, Ops.BARRIER, dtypes.void, (st,))
sres = smem.after(barr).index(uop(uops, Ops.CONST, dtypes.int32, (), 0))
self.assertEqual(_test_uops_result(dtypes.uint8, uops, sres), 42)
# NOTE: webgpu specific, since only webgpu performs bitpacking
@unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
def test_packed_smem_size(self):
_dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
size = 16
for dtype in _dtypes:
temp = UOp(Ops.DEFINE_LOCAL, dtype.ptr(size=size, addrspace=AddrSpace.LOCAL), (), 'smem')
uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
out = Device[Device.DEFAULT].renderer.render(uops)
# half is supported in wgsl, so it doesn't have to be packed
corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
self.assertIn(f"temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;", out)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
@unittest.skip("tinygrad doesn't support this behavior")
def test_local_indirect(self):
uops = []
smem = uop(uops, Ops.DEFINE_LOCAL, dtypes.int32.ptr(size=16, addrspace=AddrSpace.LOCAL), (), 'smem')
st1 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), uop(uops, Ops.CONST, dtypes.int32, (), 2)))
st2 = uop(uops, Ops.STORE, dtypes.void, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 2)), uop(uops, Ops.CONST, dtypes.int32, (), 42)))
barr = uop(uops, Ops.BARRIER, dtypes.void, (st1,st2))
ofs = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(uop(uops, Ops.CONST, dtypes.int32, (), 1)), barr))
sres = uop(uops, Ops.LOAD, dtypes.int32, (smem.index(ofs),))
self.assertEqual(_test_uops_result(dtypes.int32, uops, sres), 42)
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
class TestAssembly(unittest.TestCase):
def test_bitshift_left(self):
g1 = UOp.param(0, dtypes.int32.ptr())
out = UOp.param(1, dtypes.int32.ptr())
c1 = UOp.const(dtypes.int, 2)
c2 = UOp.const(dtypes.int, 3)
l1 = g1.index(c1)
a1 = UOp(Ops.MUL, dtypes.int, (l1, c1))
a2 = UOp(Ops.MUL, dtypes.int, (l1, c2))
uops = to_uops_list([out.index(UOp.const(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)],
ren=Device[Device.DEFAULT].renderer)
Device[Device.DEFAULT].renderer.render(uops)
ops = [x.op for x in uops]
self.assertIn(Ops.SHL, ops)
self.assertIn(Ops.MUL, ops)
def test_mulacc_unrolled(self):
# test that acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
# is not acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)
a = Tensor.empty(1024)
b = Tensor.empty(1024)
c = (a*b).sum()
ast = c.schedule_linear().src[-1].src[0]
opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
program = to_program(ast, Device[Device.DEFAULT].renderer)
uops = tuple(program.src[2].src)
self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
def test_mulacc_shl(self):
g1 = UOp.param(0, dtypes.int32.ptr())
c1 = UOp.const(dtypes.int, 0)
c2 = UOp.const(dtypes.int, 1)
expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer)
Device[Device.DEFAULT].renderer.render(uops)
self.assertIn(Ops.MULACC, [x.op for x in uops])
def test_use_cmpeq(self):
g = UOp.param(0, dtypes.uint32.ptr())
c = UOp.const(dtypes.uint, 7)
comp = g.index(c).ne(c).ne(True)
uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
Device[Device.DEFAULT].renderer.render(uops)
ops = [x.op for x in uops]
self.assertIn(Ops.CMPEQ, ops)
self.assertNotIn(Ops.CMPNE, ops)
class TestZeroRange(unittest.TestCase):
def test_reduce_variable(self):
for i in range(3,-1,-1):
v = UOp.variable("i", 0, 5).bind(i)
out = Tensor.ones(10, dtype=dtypes.int).contiguous().shrink(((0,v),)).sum()
self.assertEqual(out.item(), i)
class TestUOpPrograms(unittest.TestCase):
def _run(self, prog:UOp, *tensors:Tensor):
run_linear(UOp(Ops.LINEAR, src=(prog.call(*[t.uop.buf_uop for t in tensors]),)), update_stats=False)
def test_simple(self):
out = Tensor.empty(10,10,dtype=dtypes.int)
ptr = UOp.placeholder(out.shape, out.dtype, slot=0)
i, j = UOp.range(10, axis_id=0), UOp.range(10, axis_id=1)
prog = ptr[i,j].set(42).end(i,j)
self._run(prog.sink(arg=KernelInfo()), out)
with Context(DEBUG=0): self.assertTrue((out == 42).all().item())
def test_matmul(self):
a = Tensor.randn(10,10)
b = Tensor.randn(10,10)
c = Tensor.empty(10,10)
ref = (a@b)
with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
# C[i,j] = sum_k A[i,k] * B[k,j]
# Shapes: A[M,K], B[K,N], C[M,N]
M = N = K = 10
DT = dtypes.float32
# Placeholders (bind slots explicitly)
A = UOp.placeholder((M, K), DT, slot=0)
B = UOp.placeholder((K, N), DT, slot=1)
C = UOp.placeholder((M, N), DT, slot=2)
# Axes: i,j are spatial; k is a reduction axis over the shared dim K
i = UOp.range(M, axis_id=0) # rows of A/C
j = UOp.range(N, axis_id=1) # cols of B/C
k = UOp.range(K, axis_id=2, axis_type=AxisType.REDUCE) # reduction over K
# Zero-init: write a scalar 0 to each (i,j).
C = C[i, j].set(0.0)
# Accumulate: C_after(k) enforces the dependency along the reduction axis
C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j])
# Finalize the loop nest / schedule in (i, j, k) order
prog = C.end(i, j, k)
# run program
# TODO: make this work with opts_to_apply
self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
def test_matmul_relu(self):
a, b, c = Tensor.randn(10,10), Tensor.randn(10,10), Tensor.empty(10,10)
ref = (a@b).relu()
with Context(DEBUG=0): Tensor.realize(a, b, c, ref)
A, B, C = a.uop.placeholder_like(0), b.uop.placeholder_like(1), c.uop.placeholder_like(2)
i, j, k = UOp.range(10, 0), UOp.range(10, 1), UOp.range(10, 2, axis_type=AxisType.REDUCE)
C = C[i, j].set(0.0)
C = C[i, j].set(C.after(k)[i, j] + A[i, k] * B[k, j], end=k)
C = C[i, j].set(C[i, j].maximum(0.0))
prog = C.end(i, j)
self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
if __name__ == '__main__':
unittest.main(verbosity=2)

27
tinygrad_repo/test/backend/test_zero_copy.py Normal file
View File

@@ -0,0 +1,27 @@
import unittest
from tinygrad import Tensor, Device
import time
def time_tensor_numpy(out:Tensor):
times = []
for _ in range(5):
st = time.perf_counter()
out.uop.base.realized.as_memoryview(allow_zero_copy=True)
et = time.perf_counter() - st
times.append(et)
return min(times)
N = 4096
class TestZeroCopy(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT not in {"CPU", "METAL"}, "device isn't zero copy")
def test_zero_copy_from_default_to_cpu(self):
demo = Tensor.rand(1).realize()
t1 = time_tensor_numpy(demo)
out = Tensor.rand(N, N).realize()
t2 = time_tensor_numpy(out)
gbps = out.nbytes()*1e-9/max(t2-t1, 1e-10)
print(f"time(base): {t1*1e3:.2f} ms, time(copy): {t2*1e3:.2f} ms : copy speed {gbps:.2f} GB/s")
self.assertGreater(gbps, 600) # more than 600 GB/s = no copy
if __name__ == '__main__':
unittest.main(verbosity=2)

0
tinygrad_repo/test/device/__init__.py Normal file
View File

42
tinygrad_repo/test/device/test_amd_llvm.py Normal file
View File

@@ -0,0 +1,42 @@
import unittest
from tinygrad import Device
from tinygrad.device import CompileError
if Device.DEFAULT == "AMD":
# NOTE: if you don't gate this, LVP fails on Mac
from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler
@unittest.skipUnless(Device.DEFAULT == "AMD", "Runs only on AMD")
class TestAMDLLVM(unittest.TestCase):
def test_compiler(self):
src = '''
; https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/imm.ll
define amdgpu_kernel void @i64_imm_inline_lo(ptr addrspace(1) %out) {
entry:
store i64 1311768464867721221, ptr addrspace(1) %out ; 0x1234567800000005
ret void
}
'''
compiler = AMDLLVMCompiler("gfx1100")
compiler.compile(src)
def test_compiler_diag_error(self):
src = """
@local_temp0 = internal unnamed_addr addrspace(3) global [{N} x float*] undef, align 16
define amdgpu_kernel void @test(float* noalias align 32 %data0, half* noalias align 32 %data1, float* noalias align 32 %data2) #0
{{
%local_temp0 = addrspacecast [{N} x float*] addrspace(3)* @local_temp0 to [{N} x float*]*
%v178 = getelementptr inbounds float, float* %local_temp0, i32 1
%v133 = getelementptr inbounds float, float* %data2, i32 1
%v134 = load float, float* %v133
store float %v134, float* %v178
ret void
}}
"""
compiler = AMDLLVMCompiler("gfx1100")
compiler.compile(src.format(N=65536//8))
with self.assertRaises(CompileError):
# llvm diagnostic: <unknown>:0:0: local memory (65544) exceeds limit (65536) in function 'test'
compiler.compile(src.format(N=65536//8+1))
if __name__ == '__main__':
unittest.main()

653
tinygrad_repo/test/device/test_hcq.py Normal file
View File

@@ -0,0 +1,653 @@
import unittest, ctypes, struct, os, random, numpy as np, time
from tinygrad import Device, Tensor, dtypes
from tinygrad.helpers import mv_address, DEBUG, DEV
from test.helpers import slow, replace_opts
from tinygrad.device import Buffer, BufferSpec
from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.support.system import PCIIfaceBase
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad import Variable
MOCKGPU = DEV.interface.startswith("MOCK")
@unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run")
class TestHCQ(unittest.TestCase):
@classmethod
def setUpClass(self):
TestHCQ.d0 = Device[Device.DEFAULT]
TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
TestHCQ.b = self.a + 1
si = self.b.schedule_linear().src[-1]
TestHCQ.prg = to_program(si.src[0], TestHCQ.d0.renderer)
TestHCQ.runtime = get_runtime(TestHCQ.d0.device, TestHCQ.prg)
TestHCQ.b.uop.buffer.allocate()
TestHCQ.kernargs_ba_ptr = TestHCQ.runtime.fill_kernargs([TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf])
TestHCQ.kernargs_ab_ptr = TestHCQ.runtime.fill_kernargs([TestHCQ.a.uop.buffer._buf, TestHCQ.b.uop.buffer._buf])
def setUp(self):
TestHCQ.d0.synchronize()
TestHCQ.a.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
TestHCQ.b.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 0))))
TestHCQ.d0.synchronize() # wait for copyins to complete
# Test signals
def test_signal(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
with self.subTest(name=str(queue_type)):
queue_type().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
def test_signal_update(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
virt_signal = TestHCQ.d0.signal_t(base_buf=HCQBuffer(Variable("sig_addr", 0, 0xffffffffffffffff, dtypes.uint64), 16))
with self.subTest(name=str(queue_type)):
q = queue_type().signal(virt_signal, virt_val)
var_vals = {virt_signal.base_buf.va_addr.expr: TestHCQ.d0.timeline_signal.base_buf.va_addr, virt_val.expr: TestHCQ.d0.timeline_value}
q.submit(TestHCQ.d0, var_vals)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
var_vals = {virt_signal.base_buf.va_addr.expr: TestHCQ.d0.timeline_signal.base_buf.va_addr, virt_val.expr: TestHCQ.d0.timeline_value}
q.submit(TestHCQ.d0, var_vals)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
# Test wait
def test_wait(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
with self.subTest(name=str(queue_type)):
fake_signal = TestHCQ.d0.new_signal()
fake_signal.value = 1
queue_type().wait(fake_signal, 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
@unittest.skipIf(Device.DEFAULT == "CPU" or (DEV.interface == "MOCKPCI" and DEV.device == "AMD"), "Can't handle async update on CPU/MOCKPCI device")
def test_wait_late_set(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
with self.subTest(name=str(queue_type)):
fake_signal = TestHCQ.d0.new_signal()
queue_type().wait(fake_signal, 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
with self.assertRaises(RuntimeError):
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value, timeout=500)
fake_signal.value = 1
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
def test_wait_update(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
with self.subTest(name=str(queue_type)):
virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
virt_signal = TestHCQ.d0.signal_t(base_buf=HCQBuffer(Variable("sig_addr", 0, 0xffffffffffffffff, dtypes.uint64), 16))
fake_signal = TestHCQ.d0.new_signal()
q = queue_type().wait(virt_signal, virt_val).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
fake_signal.value = 0x30
q.submit(TestHCQ.d0, {virt_signal.base_buf.va_addr.expr: fake_signal.base_buf.va_addr, virt_val.expr: fake_signal.value})
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
# Test exec
def test_exec_one_kernel(self):
TestHCQ.d0.hw_compute_queue_t().exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
def test_exec_2_kernels_100_times(self):
virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
q = TestHCQ.d0.hw_compute_queue_t()
q.wait(TestHCQ.d0.timeline_signal, virt_val - 1) \
.exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
.exec(TestHCQ.runtime, TestHCQ.kernargs_ab_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
.signal(TestHCQ.d0.timeline_signal, virt_val)
for _ in range(100):
q.submit(TestHCQ.d0, {virt_val.expr: TestHCQ.d0.timeline_value})
TestHCQ.d0.timeline_value += 1
val = TestHCQ.a.uop.buffer.as_memoryview().cast("f")[0]
assert val == 200.0, f"got val {val}"
@unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
def test_exec_update(self):
sint_global = (Variable("sint_global", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.prg.arg.global_size[1:])
sint_local = (Variable("sint_local", 0, 0xffffffff, dtypes.uint32),) + tuple(TestHCQ.prg.arg.local_size[1:])
q = TestHCQ.d0.hw_compute_queue_t()
q.exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, sint_global, sint_local) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.submit(TestHCQ.d0, {sint_global[0].expr: 1, sint_local[0].expr: 1})
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[0]
assert val == 1.0, f"got val {val}"
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 0.0, f"got val {val}, should not be updated"
@unittest.skipIf(Device.DEFAULT in {"CPU"}, "No globals/locals on LLVM/CPU")
def test_exec_update_fuzz(self):
virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
virt_local = [Variable(f"local_{i}", 0, 0xffffffff, dtypes.uint32) for i in range(3)]
a = Tensor.randint((3, 3, 3), dtype=dtypes.int, device=Device.DEFAULT).realize()
b = a + 1
si = b.schedule_linear().src[-1]
prg = to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer)
runtime = get_runtime(Device.DEFAULT, prg)
zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
ctypes.memset(zb._buf.va_addr, 0, zb.nbytes)
kernargs = runtime.fill_kernargs([zt._buf, zb._buf])
q = TestHCQ.d0.hw_compute_queue_t()
q.memory_barrier() \
.exec(runtime, kernargs, (1,1,1), virt_local) \
.signal(TestHCQ.d0.timeline_signal, virt_val)
for x in range(1, 4):
for y in range(1, 4):
for z in range(1, 4):
ctypes.memset(zt._buf.va_addr, 0, zb.nbytes)
q.submit(TestHCQ.d0, {virt_val.expr: TestHCQ.d0.timeline_value, virt_local[0].expr: x, virt_local[1].expr: y, virt_local[2].expr: z})
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
res_sum = sum(x for x in zt.as_memoryview().cast("I"))
assert x * y * z == res_sum, f"want {x * y * z}, got {res_sum}"
# Test copy
def test_copy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf, 8) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 1.0, f"got val {val}"
def test_copy_long(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
sz = 64 << 20
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 1, sz)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf, buf2._buf, sz) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
mv_buf1 = buf1.as_memoryview().cast('Q')
assert libc.memcmp(mv_address(mv_buf1), buf2._buf.va_addr, sz) == 0
@slow
def test_copy_64bit(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
# NOTE: these must be a multiple of 8 for .view(fmt='Q') to work
for sz in [(1 << 32) - 8, (1 << 32), (1 << 32) + 8, (5 << 30), (6 << 30) - 0x4642ee0]:
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 0x3e, sz)
buf2_q_view = buf2._buf.cpu_view().view(fmt='Q')
for i in range(0, sz//8, 0x1000):
for j in range(32): buf2_q_view[min(max(i + j - 16, 0), (sz // 8) - 1)] = random.randint(0, 0xffffffffffffffff)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf, buf2._buf, sz) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
mv_buf1 = buf1.as_memoryview()
assert libc.memcmp(mv_address(mv_buf1), buf2._buf.va_addr, sz) == 0
def test_update_copy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
virt_src_addr = Variable("virt_src_addr", 0, 0xffffffffffffffff, dtypes.uint64)
virt_dest_addr = Variable("virt_dest_addr", 0, 0xffffffffffffffff, dtypes.uint64)
q = TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(HCQBuffer(virt_dest_addr, 8), HCQBuffer(virt_src_addr, 8), 8) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.submit(TestHCQ.d0, {virt_src_addr.expr: TestHCQ.a.uop.buffer._buf.va_addr, virt_dest_addr.expr: TestHCQ.b.uop.buffer._buf.va_addr})
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
val = TestHCQ.b.uop.buffer.as_memoryview().cast("f")[1]
assert val == 1.0, f"got val {val}"
def test_update_copy_long(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
virt_src_addr = Variable("virt_src_addr", 0, 0xffffffffffffffff, dtypes.uint64)
virt_dest_addr = Variable("virt_dest_addr", 0, 0xffffffffffffffff, dtypes.uint64)
sz = 64 << 20
buf1 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
ctypes.memset(buf2._buf.va_addr, 1, sz)
q = TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(HCQBuffer(virt_dest_addr, sz), HCQBuffer(virt_src_addr, sz), sz) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.submit(TestHCQ.d0, {virt_src_addr.expr: buf2._buf.va_addr, virt_dest_addr.expr: buf1._buf.va_addr})
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
mv_buf1 = buf1.as_memoryview().cast('Q')
for i in range(sz//8): assert mv_buf1[i] == 0x0101010101010101, f"offset {i*8} differs, not all copied, got {hex(mv_buf1[i])}"
# Test bind api
def test_bind(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)
virt_signal = TestHCQ.d0.signal_t(base_buf=HCQBuffer(Variable("sig_addr", 0, 0xffffffffffffffff, dtypes.uint64), 16))
with self.subTest(name=str(queue_type)):
fake_signal = TestHCQ.d0.new_signal()
q = queue_type().wait(virt_signal, virt_val).signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
q.bind(TestHCQ.d0)
fake_signal.value = 0x30
q.submit(TestHCQ.d0, {virt_signal.base_buf.va_addr.expr: fake_signal.base_buf.va_addr, virt_val.expr: fake_signal.value})
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
# Test multidevice
def test_multidevice_signal_wait(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
try: d1 = Device[f"{Device.DEFAULT}:1"]
except Exception: self.skipTest("no multidevice, test skipped")
TestHCQ.d0.hw_copy_queue_t().signal(sig:=TestHCQ.d0.new_signal(value=0), value=0xfff) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
d1.hw_copy_queue_t().wait(sig, value=0xfff) \
.signal(d1.timeline_signal, d1.timeline_value).submit(d1)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
d1.timeline_signal.wait(d1.timeline_value)
d1.timeline_value += 1
# Test profile api
def test_speed_exec_time(self):
sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \
.exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
.timestamp(sig_en) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
et = float(sig_en.timestamp - sig_st.timestamp)
print(f"exec kernel time: {et:.2f} us")
assert 0.1 <= et <= (3000000 if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)
def test_speed_copy_bandwidth(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
# THEORY: the bandwidth is low here because it's only using one SDMA queue. I suspect it's more stable like this at least.
SZ = 200_000_000
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
.copy(a._buf, b._buf, SZ) \
.timestamp(sig_en) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
et = float(sig_en.timestamp - sig_st.timestamp)
et_ms = et / 1e3
gb_s = ((SZ / 1e9) / et_ms) * 1e3
print(f"same device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s")
assert (0.2 if MOCKGPU else 10) <= gb_s <= 1000
def test_speed_cross_device_copy_bandwidth(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
try: _ = Device[f"{Device.DEFAULT}:1"]
except Exception: self.skipTest("no multidevice, test skipped")
SZ = 200_000_000
b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
TestHCQ.d0.allocator.map(b._buf)
sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
.copy(a._buf, b._buf, SZ) \
.timestamp(sig_en) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
et = float(sig_en.timestamp - sig_st.timestamp)
et_ms = et / 1e3
gb_s = ((SZ / 1e9) / et_ms) * 1e3
print(f"cross device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s")
assert (0.2 if MOCKGPU else 2) <= gb_s <= 100
def test_timeline_signal_rollover(self):
for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
if queue_type is None: continue
with self.subTest(name=str(queue_type)):
TestHCQ.d0.timeline_value = (1 << 32) - 20 # close value to reset
queue_type().signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value - 1)
for _ in range(40):
queue_type().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_value += 1
TestHCQ.d0.synchronize()
def test_small_copies_from_host_buf(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
for i in range(256):
ctypes.memset(buf2._buf.va_addr, i, 1)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf, buf2._buf, 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
assert buf1.as_memoryview()[0] == i
def test_small_copies_from_host_buf_intercopy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
for i in range(256):
ctypes.memset(buf3._buf.va_addr, i, 1)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf, buf3._buf, 1) \
.copy(buf2._buf, buf1._buf, 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
assert buf2.as_memoryview()[0] == i
def test_small_copies_from_host_buf_transfer(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
try: _ = Device[f"{Device.DEFAULT}:1"]
except Exception: self.skipTest("no multidevice, test skipped")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(f"{Device.DEFAULT}:1", 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
TestHCQ.d0.allocator.map(buf2._buf)
for i in range(256):
ctypes.memset(buf3._buf.va_addr, i, 1)
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf, buf3._buf, 1) \
.copy(buf2._buf, buf1._buf, 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
assert buf2.as_memoryview()[0] == i
def test_memory_barrier(self):
a = Tensor([0, 1], device=Device.DEFAULT, dtype=dtypes.int8).realize()
b = a + 1
prg = to_program(b.schedule_linear().src[-1].src[0], TestHCQ.d0.renderer)
runtime = get_runtime(TestHCQ.d0.device, prg)
buf1 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
kernargs_ptr = runtime.fill_kernargs([buf1._buf, buf2._buf])
for i in range(255):
ctypes.memset(buf2._buf.va_addr, i, 2)
# Need memory_barrier after direct write to vram
TestHCQ.d0.hw_compute_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.memory_barrier() \
.exec(runtime, kernargs_ptr, prg.arg.global_size, prg.arg.local_size) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
assert buf1.as_memoryview()[0] == (i + 1), f"has {buf1.as_memoryview()[0]}, need {i + 1}"
def test_memory_barrier_before_copy(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf2 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
for i in range(256):
ctypes.memset(buf3._buf.va_addr, i, 1)
# Need memory_barrier after direct write to vram
TestHCQ.d0.hw_compute_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.memory_barrier() \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_value += 1
TestHCQ.d0.hw_copy_queue_t().wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1) \
.copy(buf1._buf, buf3._buf, 1) \
.copy(buf2._buf, buf1._buf, 1) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
assert buf2.as_memoryview()[0] == i
def test_write(self):
buf = Buffer(Device.DEFAULT, 4, dtypes.uint32, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
try:
TestHCQ.d0.hw_compute_queue_t().write(buf._buf, 0x42) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
except NotImplementedError: self.skipTest("write not implemented")
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
assert buf.as_memoryview().cast("I")[0] == 0x42
def test_poll_bit_set(self):
buf = Buffer(Device.DEFAULT, 4, dtypes.uint32, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
try:
TestHCQ.d0.hw_compute_queue_t().write(buf._buf, 0x01000000, b64=False) \
.poll_bit(buf._buf, 0x01000000, 0x01000000) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
except NotImplementedError: self.skipTest("write/poll_bit not implemented")
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
def test_poll_bit_clear(self):
buf = Buffer(Device.DEFAULT, 4, dtypes.uint32, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
try:
TestHCQ.d0.hw_compute_queue_t().write(buf._buf, 0xFE000000, b64=False) \
.poll_bit(buf._buf, 0, 0x01000000) \
.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
except NotImplementedError: self.skipTest("write/poll_bit not implemented")
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
TestHCQ.d0.timeline_value += 1
def test_map_cpu_buffer_to_device(self):
if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")
sz = 0x2000
cpu_buffer = Buffer("CPU", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
cpu_buffer._buf.cpu_view().view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
for devid in range(6):
if DEBUG >= 2: print(f"Testing map to device {Device.DEFAULT}:{devid}")
try: d = Device[f"{Device.DEFAULT}:{devid}"]
except Exception: break
local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
d.allocator.map(cpu_buffer._buf)
d.hw_copy_queue_t().wait(d.timeline_signal, d.timeline_value - 1) \
.copy(local_buf._buf, cpu_buffer._buf, sz) \
.signal(d.timeline_signal, d.timeline_value).submit(d)
d.timeline_signal.wait(d.timeline_value)
d.timeline_value += 1
np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
@unittest.skipUnless(MOCKGPU and not (DEV.device == "AMD" and DEV.interface == "MOCKPCI"), "Emulate this on MOCKGPU to check the path in CI")
def test_on_device_hang(self):
if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
os.environ["MOCKGPU_EMU_FAULTADDR"] = "0xDEADBEE1"
# Check api calls
with self.assertRaises(RuntimeError) as ctx:
self.d0.on_device_hang()
assert "0xDEADBEE1" in str(ctx.exception)
os.environ.pop("MOCKGPU_EMU_FAULTADDR")
def test_multidevice(self):
try: amd_dev = Device["AMD"]
except Exception: self.skipTest("no AMD device, test skipped")
try: nv_dev = Device["NV"]
except Exception: self.skipTest("no NV device, test skipped")
x = amd_dev.new_signal()
y = nv_dev.new_signal()
assert type(x) is amd_dev.signal_t
assert type(y) is nv_dev.signal_t
def test_multidevice_p2p(self):
try:
amd_dev = Device["AMD"]
if not issubclass(type(amd_dev.iface), PCIIfaceBase): self.skipTest("Not a pci dev")
except Exception: self.skipTest("no AMD device, test skipped")
try:
nv_dev = Device["NV"]
if not issubclass(type(nv_dev.iface), PCIIfaceBase): self.skipTest("Not a pci dev")
except Exception: self.skipTest("no NV device, test skipped")
def _check_copy(dev1, dev2):
buf1 = Tensor.randn(10, 10, device=dev1).realize()
buf2 = buf1.to(dev2).realize()
np.testing.assert_equal(buf1.numpy(), buf2.numpy(), "p2p failed")
_check_copy("AMD", "NV")
_check_copy("NV", "AMD")
def test_speed_cross_device_rdma_copy_bandwidth(self):
try: d1 = Device[f"{Device.DEFAULT}:7"]
except Exception: self.skipTest("no multidevice, test skipped")
if TestHCQ.d0.peer_group == d1.peer_group: self.skipTest("devices in same peer group, no RDMA path")
SZ = 200_000_000
a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
b = Buffer(f"{Device.DEFAULT}:7", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
# warmup
TestHCQ.d0.allocator._transfer(a._buf, b._buf, SZ, src_dev=d1, dest_dev=TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value - 1)
d1.timeline_signal.wait(d1.timeline_value - 1)
st = time.perf_counter()
TestHCQ.d0.allocator._transfer(a._buf, b._buf, SZ, src_dev=d1, dest_dev=TestHCQ.d0)
TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value - 1)
d1.timeline_signal.wait(d1.timeline_value - 1)
et_ms = (time.perf_counter() - st) * 1e3
gb_s = ((SZ / 1e9) / et_ms) * 1e3
print(f"cross device rdma copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s")
assert 1 <= gb_s <= 100
np.testing.assert_equal(a.numpy(), b.numpy(), "failed")
if __name__ == "__main__":
unittest.main()

61
tinygrad_repo/test/device/test_metal.py Normal file
View File

@@ -0,0 +1,61 @@
import unittest
from tinygrad.device import CompileError, Device, BufferSpec
if Device.DEFAULT=="METAL":
from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
@unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
class TestMetal(unittest.TestCase):
def test_alloc_oom(self):
device = MetalDevice("metal")
with self.assertRaises(MemoryError):
device.allocator.alloc(10000000000000000000)
def test_compile_error(self):
compiler = MetalCompiler()
with self.assertRaises(CompileError):
compiler.compile("this is not valid metal")
def test_compile_success(self):
compiler = MetalCompiler()
ret = compiler.compile("""
#include <metal_stdlib>
using namespace metal;
kernel void E_4n1(device int* data0, const device int* data1, const device int* data2,
uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {
int val0 = *(data1+0);
int val1 = *(data1+1);
int val2 = *(data1+2);
int val3 = *(data1+3);
int val4 = *(data2+0);
int val5 = *(data2+1);
int val6 = *(data2+2);
int val7 = *(data2+3);
*(data0+0) = (val0+val4);
*(data0+1) = (val1+val5);
*(data0+2) = (val2+val6);
*(data0+3) = (val3+val7);
}
""")
assert ret is not None
def test_failed_newLibraryWithData(self):
device = MetalDevice("metal")
compiler = MetalCompiler()
compiled = compiler.compile("""
#include <metal_stdlib>
kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]){
data0[0] = 0;
}
""")
with self.assertRaises(RuntimeError):
compiled = compiled[:40] # corrupt the compiled program
MetalProgram(device, "r_5", compiled)
def test_free(self):
size = 2**16
device = Device['METAL']
before = device.sysdevice.currentAllocatedSize()
buf = device.allocator.alloc(size, BufferSpec(nolru=True))
self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before+size, msg=f"{curr=} - {before=}")
device.allocator.free(buf, buf.size, BufferSpec(nolru=True))
self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before, msg=f"{curr=} - {before=}")

40
tinygrad_repo/test/device/test_ocl.py Normal file
View File

@@ -0,0 +1,40 @@
import unittest
from unittest.mock import patch
from tinygrad import Device
from tinygrad.device import Buffer
from tinygrad.dtype import dtypes
from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler, CLProgram
@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
class TestCLCompileCache(unittest.TestCase):
def test_compile_cached(self):
device = Device[Device.DEFAULT]
src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
CLProgram(device, name="cached_test", lib=src.encode())
with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
CLProgram(device, name="cached_test", lib=src.encode())
@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
class TestCLError(unittest.TestCase):
@unittest.skip("allocates tons of memory")
def test_oom(self):
with self.assertRaises(RuntimeError) as err:
allocator = CLAllocator(CLDevice())
for i in range(1_000_000):
allocator.alloc(1_000_000_000)
assert str(err.exception) == "OpenCL Error -6: CL_OUT_OF_HOST_MEMORY"
def test_invalid_kernel_name(self):
device = Device[Device.DEFAULT]
with self.assertRaises(RuntimeError) as err:
CLProgram(device, name="", lib="__kernel void test(__global int* a) { a[0] = 1; }".encode())
assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
def test_unaligned_copy(self):
data = list(range(65))
unaligned = memoryview(bytearray(data))[1:]
buffer = Buffer("CL", 64, dtypes.uint8).allocate()
buffer.copyin(unaligned)
result = memoryview(bytearray(len(data) - 1))
buffer.copyout(result)
assert unaligned == result, "Unaligned data copied in must be equal to data copied out."

40
tinygrad_repo/test/device/test_validate_with_cpu.py Normal file
View File

@@ -0,0 +1,40 @@
import unittest
from tinygrad import Tensor, Context, Variable, Device
from test.helpers import needs_second_gpu
class TestValidateWithCPU(unittest.TestCase):
def setUp(self):
self.ctx = Context(VALIDATE_WITH_CPU=1)
self.ctx.__enter__()
def tearDown(self): self.ctx.__exit__(None, None, None)
def test_add(self): self.assertListEqual((Tensor([1.,2,3])+Tensor([4.,5,6])).tolist(), [5.0, 7.0, 9.0])
def test_mul(self): self.assertListEqual((Tensor([1.,2,3])*Tensor([4.,5,6])).tolist(), [4.0, 10.0, 18.0])
def test_sum(self): self.assertEqual(Tensor([1.,2,3,4]).sum().item(), 10.0)
def test_reduce_then_op(self): self.assertEqual((Tensor([1.,2,3,4]).sum() * 2).item(), 20.0)
def test_assign(self):
a = Tensor([1.,2,3]).realize()
a.assign(a + 1).realize()
self.assertListEqual(a.tolist(), [2.0, 3.0, 4.0])
def test_buffer_view(self):
self.assertListEqual((Tensor([1.,2,3,4,5,6,7,8])[2:6] + 1).tolist(), [4.0, 5.0, 6.0, 7.0])
def test_symbolic(self):
i = Variable('i', 1, 10)
ones = Tensor.ones(10).contiguous()
self.assertListEqual((ones[:i.bind(5)] + 1).contiguous()[:5].tolist(), [2.0]*5)
def test_multi_kernel(self):
a = (Tensor([1.,2,3]) + 1).contiguous()
b = (a * 2).contiguous()
self.assertListEqual((b - 1).tolist(), [3.0, 5.0, 7.0])
@needs_second_gpu
def test_sharded(self):
t = Tensor([1.,2,3,4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
self.assertListEqual((t + 1).tolist(), [2.0, 3.0, 4.0, 5.0])
if __name__ == "__main__":
unittest.main()

9
tinygrad_repo/test/external/external_benchmark_am.py vendored Normal file
View File

@@ -0,0 +1,9 @@
from tinygrad.helpers import Profiling
from tinygrad import Device
if __name__ == "__main__":
am = Device["AMD"]
# kfd is 0.55ms!
with Profiling("allocation 127.7mb"):
am.allocator.alloc(int(127.7*1024*1024))

18
tinygrad_repo/test/external/external_benchmark_bert_matmuls.py vendored Normal file
View File

@@ -0,0 +1,18 @@
from tinygrad import Tensor, dtypes
dtypes.default_float = dtypes.float16
from tinygrad.dtype import to_dtype
from tinygrad.helpers import getenv
if __name__ == "__main__":
# matmuls in bert layers
BS = getenv("BS", 96//6)
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
tensors = [
(Tensor.empty(BS, 512, 1024), Tensor.empty(1024, 1024).T), # linear to get qkv
(Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3), Tensor.empty(BS, 512, 16, 64).permute(0,2,3,1)), # q@k
(Tensor.empty(BS, 16, 512, 512), Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3)), # qk@v
]
for t0, t1 in tensors:
print(f"{t0.shape=}, {t0.uop.st.is_expanded()=}, {t1.shape=}, {t1.uop.st.is_expanded()=}")
for _ in range(5):
t0.dot(t1, dtype=acc_dtype).realize()

17
tinygrad_repo/test/external/external_benchmark_bert_softmax.py vendored Normal file
View File

@@ -0,0 +1,17 @@
from tinygrad import Tensor, dtypes, GlobalCounters
dtypes.default_float = dtypes.float16
from tinygrad.dtype import to_dtype
from tinygrad.helpers import getenv
from test.backend.test_softmax_fusion import single_kernel_softmax
if __name__ == "__main__":
# softmax in bert layers
BS = getenv("BS", 96//6)
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
t = Tensor.empty(BS, 16, 512, 512)
t.softmax(-1, dtype="half").realize()
# test single kernel softmax
GlobalCounters.reset()
single_kernel_softmax(t, -1, acc_dtype).realize()

8
tinygrad_repo/test/external/external_benchmark_disk_raw.py vendored Normal file
View File

@@ -0,0 +1,8 @@
import pathlib
from tinygrad import Tensor, Device, Context
from tinygrad.helpers import getenv
if __name__ == "__main__":
with Context(DEBUG=2):
disk_llama = Tensor(pathlib.Path(getenv("TESTFILE", "/raid/weights/LLaMA-3/8B/consolidated.00.pth")))
device_llama = disk_llama.to(Device.DEFAULT).realize()

31
tinygrad_repo/test/external/external_benchmark_hip_compile.py vendored Normal file
View File

@@ -0,0 +1,31 @@
import random, os
from tinygrad.helpers import Timing
from tinygrad.runtime.ops_hip import compile_hip, HIPDevice
from tinygrad.runtime.ops_cl import compile_cl, CLDevice
# OMP_NUM_THREADS=1 strace -tt -f -e trace=file python3 test/external/external_benchmark_hip_compile.py
# AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1 python3 test/external/external_benchmark_hip_compile.py
# issue is in https://github.com/ROCm-Developer-Tools/clr/
if __name__ == "__main__":
HIPDevice()
CLDevice()
# warmup
name = "none"+str(random.randint(0, 1000000))
compile_cl.__wrapped__(f"void {name}() {{}}")
print("compile cl warmed up")
compile_hip.__wrapped__(f"void {name}() {{}}")
print("compile hip warmed up")
print("**** benchmark ****")
name = "none"+str(random.randint(0, 1000000))
# this uses AMD_COMGR_ACTION_COMPILE_SOURCE_TO_BC, then it links the lib on the next step
with Timing("compile cl: "): compile_cl.__wrapped__(f"void {name}() {{}}")
# this uses AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, much slower
with Timing("compile hip: "): compile_hip.__wrapped__(f"void {name}() {{}}")
os._exit(0)

20
tinygrad_repo/test/external/external_benchmark_keccak.py vendored Normal file
View File

@@ -0,0 +1,20 @@
from tinygrad import Tensor, dtypes
from tinygrad.engine.jit import TinyJit
from tinygrad.helpers import Timing, getenv
if __name__ == "__main__":
BS = getenv("BS", 2**14)
BLOCKSIZE = getenv("BLOCKSIZE", 4096)
HASHFN = getenv("HASHFN", "shake_128")
NRUNS = getenv("NRUNS", 5)
@TinyJit
def hasher(data: Tensor): return data.keccak(HASHFN)
t = Tensor.randn(BS, BLOCKSIZE, dtype=dtypes.uint8).realize()
ds_mib = t.nbytes() / 1024**2
print(f"--- benchmarking (hash: {HASHFN}, data size: {ds_mib} MiB, block size: {BLOCKSIZE} B, batch size: {BS})")
for i in range(NRUNS):
with Timing(f"run: {i+1}, elapsed time: ", (lambda et: f", throughput: {ds_mib / (et*1e-9):.2f} MiB/s")):
hasher(t).realize()

38
tinygrad_repo/test/external/external_benchmark_kernel_launch.py vendored Normal file
View File

@@ -0,0 +1,38 @@
import time
from tinygrad import Tensor, TinyJit, Device, Context
from tinygrad.helpers import Profiling, Timing, GlobalCounters
# python3 test/speed/external_test_speed_v_torch.py TestSpeed.test_add_a
@TinyJit
def plus(a:Tensor, b:Tensor): return a+b
if __name__ == "__main__":
a = Tensor([1]).realize()
b = Tensor([1]).realize()
for i in range(5):
with Timing(prefix=f"{i}:"):
c = plus(a,b)
Device[c.device].synchronize()
assert c.item() == 2
for i in range(5):
st = time.perf_counter()
c = plus(a,b)
et = time.perf_counter() - st
Device[c.device].synchronize()
print(f"nosync {i}: {et*1e6:.2f} us")
for i in range(5):
st = time.perf_counter()
c = plus(a,b)
Device[c.device].synchronize()
et = time.perf_counter() - st
print(f"precise {i}: {et*1e6:.2f} us")
assert GlobalCounters.time_sum_s == 0
with Context(DEBUG=2):
st = time.perf_counter()
c = plus(a,b)
Device[c.device].synchronize()
et = time.perf_counter() - st
print(f"kernel {GlobalCounters.time_sum_s*1e3:.2f} ms / full {et*1e3:.2f} ms -- {et/(GlobalCounters.time_sum_s+1e-12):.2f} x")
with Profiling():
c = plus(a,b)

35
tinygrad_repo/test/external/external_benchmark_llama_schedule.py vendored Normal file
View File

@@ -0,0 +1,35 @@
from tinygrad import nn, Tensor, dtypes
from tinygrad.helpers import DEV, Timing
from extra.models.llama import Transformer
from examples.llama3 import MODEL_PARAMS
if __name__ == "__main__":
DEV.value = "NULL"
Tensor.training = True
#model_size = "8B"
model_size = "405B"
with Timing("total "):
with Timing("***** create model in "):
model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
max_context=1024, jit=True, disable_kv_cache=True)
with Timing("***** fake state in "):
Tensor.realize(*[p.assign(Tensor.empty(*p.shape, device=p.device, dtype=p.dtype)) for p in nn.state.get_parameters(model)])
with Timing("***** create optim in "):
opt = nn.optim.AdamW(nn.state.get_parameters(model))
with Timing("***** run model in "):
toks = Tensor.empty(1, 1024, dtype=dtypes.int)
out = model(toks, 0, temperature=float('nan'))
with Timing("***** backward in "):
out.mean().backward()
with Timing("***** realize in "):
out.realize()
with Timing("***** step in "):
opt.step()

49
tinygrad_repo/test/external/external_benchmark_multitensor_allreduce.py vendored Normal file
View File

@@ -0,0 +1,49 @@
from tinygrad import Tensor, Device, GlobalCounters, TinyJit, dtypes
from tinygrad.helpers import getenv, Context, DEBUG
def test(devs: list[str], N: int, iters:int = 10, name:str = "allreduce"):
@TinyJit
def f(t: Tensor) -> Tensor: t.sum(0).realize()
secs, gflops, gbs = 0, 0, 0
for i in range(-3, iters):
t = Tensor.empty((len(devs), N))
t = t.shard(devs, 0).realize()
GlobalCounters.reset()
f(t)
for d in devs: Device[d].synchronize()
if i < 0: continue # warm up jit
i_secs = GlobalCounters.time_sum_s
i_gflops = GlobalCounters.global_ops/i_secs/10**9
i_gbs = (N*4)/i_secs/10**9
print(f"{name} iter {i+1}/{iters}: {i_secs:.6f} sec {i_gflops:.2f} GFLOP/s {i_gbs:.2f} GB/s")
secs += i_secs
gflops += i_gflops
gbs += i_gbs
return (gflops/iters, gbs/iters, secs/iters)
def run(sz, n_gpus=6, iters=10, ring=0, all2all=0):
devs = tuple([f"{Device.DEFAULT}:{x}" for x in range(n_gpus)])
N = sz // dtypes.float32.itemsize
name = "all2all" if all2all else ("ring" if ring else "naive")
with Context(RING=(2 if ring else 0), ALL2ALL=(2 if all2all else 0), JIT_BATCH_SIZE=0, DEBUG=max(DEBUG.value, 2)):
return test(devs, N, iters=iters, name=name)
def main():
n_gpus = getenv("GPUS", 6)
iters = getenv("ITERS", 10)
sz = getenv("SZ", 1000) * 10**6 # size of data on each gpu
print(f"Using {sz/10**9:.2f} GB of numbers on each of {n_gpus} GPUs, {n_gpus*sz/10**9:.2f} GB total.")
results = {}
for name, kwargs in [("naive", {}), ("ring", {"ring": 2}), ("all2all", {"all2all": 2})]:
results[name] = run(sz, n_gpus=n_gpus, iters=iters, **kwargs)
print("\n=== RESULTS ===")
for name, (gflops, gbs, secs) in results.items():
print(f"{name.upper()}:\n {secs:.6f} seconds/iter\n {gflops:.2f} GFLOP/s\n {gbs:.2f} GB/s")
if __name__ == "__main__":
main()

102
tinygrad_repo/test/external/external_benchmark_op_conv.py vendored Normal file
View File

@@ -0,0 +1,102 @@
# ruff: noqa: E501 E712 F401
from dataclasses import replace
from tinygrad import dtypes, Device
from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
from tinygrad.helpers import dedup, getenv
from tinygrad.device import Buffer
from tinygrad.dtype import ImageDType, Invalid
# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
def vision_conv_143():
c0 = UOp.param(0, dtypes.imageh((16, 1024, 4)))
c2 = UOp.range(32, 3, AxisType.LOOP)
c5 = UOp.range(128, 4, AxisType.LOOP)
c8 = UOp.range(16, 2, AxisType.LOOP)
c16 = UOp.range(7, 0, AxisType.REDUCE)
c17 = c8*2+c16
c24 = ((c17<3)!=True)&(c17<35)
c26 = UOp.range(7, 1, AxisType.REDUCE)
c27 = c2*2+c26
c32 = ((c27<3)!=True)&(c27<67)
c34 = UOp.param(1, dtypes.imageh((32, 1024, 4)))
c38 = c5//2
c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
c49 = UOp.param(2, dtypes.imageh((64, 49, 4)))
c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
c63 = UOp.param(3, dtypes.float.ptr(128))
c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
c67 = c0.index((c2*128+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
opts = None
# JITBEAM=2
# (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
def vision_conv_153():
c0 = UOp.param(0, dtypes.imageh((8, 1024, 4)))
c2 = UOp.range(16, 3, AxisType.LOOP)
c5 = UOp.range(256, 4, AxisType.LOOP)
c8 = UOp.range(8, 2, AxisType.LOOP)
c16 = UOp.range(7, 0, AxisType.REDUCE)
c17 = c8*2+c16
c24 = ((c17<3)!=True)&(c17<19)
c26 = UOp.range(7, 1, AxisType.REDUCE)
c27 = c2*2+c26
c32 = ((c27<3)!=True)&(c27<35)
c34 = UOp.param(1, dtypes.imageh((16, 1024, 4)))
c38 = c5//2
c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
c49 = UOp.param(2, dtypes.imageh((128, 49, 4)))
c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
c63 = UOp.param(3, dtypes.float.ptr(256))
c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
c67 = c0.index((c2*256+c5+c8*4096), ptr=True).store(c65).end(c8, c2, c5)
opts = None
# JITBEAM=2
# (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
def dm_conv_172():
c0 = UOp.param(0, dtypes.imageh((1, 240, 4)))
c2 = UOp.range(960, 4, AxisType.LOOP)
c5 = UOp.param(1, dtypes.imageh((8, 384, 4)))
c7 = UOp.range(32, 0, AxisType.REDUCE)
c10 = UOp.range(4, 1, AxisType.REDUCE)
c13 = UOp.range(12, 3, AxisType.REDUCE)
c18 = UOp.range(8, 2, AxisType.REDUCE)
c23 = UOp.param(2, dtypes.imageh((240, 128, 4)))
c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
c37 = UOp.param(3, dtypes.float.ptr(960))
c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
c55 = c0.index(c2, ptr=True).store(c53).end(c2)
opts = None
# JITBEAM=2
# (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))
ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()
renderer = Device.default.renderer
allocator = Device.default.allocator
ps = to_program(ast, renderer)
rt = get_runtime(Device.DEFAULT, ps)
gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
# print(len(gs))
# print([g.dtype for g in gs])
bufs = [Buffer(ps.arg.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
gsize, lsize = ps.arg.launch_dims({})
t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
print(f"{t*1e6:.2f} us")

34
tinygrad_repo/test/external/external_benchmark_pyrender.py vendored Normal file
View File

@@ -0,0 +1,34 @@
# benchmark speed of pyrender for all created UOps saved with TRACK_MATCH_STATS=2
import functools, pickle
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import tqdm, temp, time_to_str, cpu_profile
BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
@functools.cache
def create_uop(a:int) -> UOp:
op, dtype, src, arg, *rest = trace.uop_fields[a]
return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
if __name__ == "__main__":
# load rewrite trace
with open(temp("rewrites.pkl", append_user=True), "rb") as f:
trace = pickle.load(f)
# benchmark
result:list[tuple[str, int]] = []
try:
for steps in tqdm(trace.rewrites):
for r in steps:
for _,yn,_,__ in r.matches:
y = create_uop(yn)
if y.op in BENCHMARK_OPS:
with cpu_profile("pyrender") as e:
try: ren = y.render()
except Exception: ren = "PYRENDER_ERR"
result.append((ren, float(e.en-e.st)/1e6))
finally:
N = 10
print(f"Slowst {N} renders from {len(result)} samples:")
for ren,tm in sorted(result, key=lambda x:x[1], reverse=True)[:N]:
print(f"{time_to_str(tm).strip():<10s} {ren}")

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