mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-10 14:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
1451 lines
74 KiB
Python
1451 lines
74 KiB
Python
# Tokenizer-based expression parser for AMD pcode
|
|
import ast, itertools, operator, re
|
|
from typing import Any, Callable
|
|
|
|
_BINOPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.FloorDiv: operator.floordiv,
|
|
ast.Mod: operator.mod, ast.LShift: operator.lshift, ast.RShift: operator.rshift,
|
|
ast.BitAnd: operator.and_, ast.BitOr: operator.or_, ast.BitXor: operator.xor}
|
|
def _const_int(expr: str) -> int:
|
|
"""Evaluate a compile-time integer expression (integer literals and basic arithmetic only)."""
|
|
def ev(node: ast.AST) -> int:
|
|
if isinstance(node, ast.Expression): return ev(node.body)
|
|
if isinstance(node, ast.Constant) and isinstance(node.value, int): return node.value
|
|
if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
|
|
return (-1 if isinstance(node.op, ast.USub) else 1) * ev(node.operand)
|
|
if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS: return _BINOPS[type(node.op)](ev(node.left), ev(node.right))
|
|
raise ValueError(f"not a constant integer expression: {expr!r}")
|
|
return ev(ast.parse(expr.strip(), mode='eval'))
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.uop.ops import Ops, UOp
|
|
from tinygrad.codegen.decomp.dtype import f2f
|
|
|
|
# Type alias for vars dict: stores UOps and tuples for lambda definitions
|
|
VarVal = UOp | tuple[str, list[str], str]
|
|
|
|
def _const(dt, v): return UOp.const(v, dt)
|
|
def _single_value(v: UOp): return v.vmin if v.vmin == v.vmax else None
|
|
def _u32(v): return _const(dtypes.uint32, v)
|
|
def _u64(v): return _const(dtypes.uint64, v)
|
|
def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
|
|
def _to_bool(v): return v if v.dtype == dtypes.bool else v.ne(_const(v.dtype, 0))
|
|
def _cast_to(v, dt):
|
|
if v.dtype == dt: return v
|
|
if dt == dtypes.half: return v.cast(dtypes.uint16).bitcast(dtypes.half)
|
|
return v.cast(dt) if dt.itemsize != v.dtype.itemsize else v.bitcast(dt)
|
|
|
|
# Float bit extraction - returns (bits, exp_mask, mant_mask, quiet_bit, exp_shift) based on float type
|
|
def _float_info(v: UOp) -> tuple[UOp, UOp, UOp, UOp, int]:
|
|
if v.dtype in (dtypes.float64, dtypes.uint64):
|
|
bits = v.bitcast(dtypes.uint64) if v.dtype == dtypes.float64 else v.cast(dtypes.uint64)
|
|
return bits, _u64(0x7FF0000000000000), _u64(0x000FFFFFFFFFFFFF), _u64(0x0008000000000000), 52
|
|
if v.dtype in (dtypes.half, dtypes.uint16):
|
|
bits = (v.bitcast(dtypes.uint16) if v.dtype == dtypes.half else (v & _u32(0xFFFF)).cast(dtypes.uint16)).cast(dtypes.uint32)
|
|
return bits, _u32(0x7C00), _u32(0x03FF), _u32(0x0200), 10
|
|
bits = v.bitcast(dtypes.uint32) if v.dtype == dtypes.float32 else v.cast(dtypes.uint32)
|
|
return bits, _u32(0x7F800000), _u32(0x007FFFFF), _u32(0x00400000), 23
|
|
|
|
def _isnan(v: UOp) -> UOp:
|
|
bits, exp_m, mant_m, _, _ = _float_info(v.cast(dtypes.float32) if v.dtype == dtypes.half else v)
|
|
return (bits & exp_m).eq(exp_m) & (bits & mant_m).ne(_const(bits.dtype, 0))
|
|
|
|
def _bitreverse(v: UOp, bits: int) -> UOp:
|
|
dt, masks = (dtypes.uint64, [(0x5555555555555555,1),(0x3333333333333333,2),(0x0F0F0F0F0F0F0F0F,4),(0x00FF00FF00FF00FF,8),(0x0000FFFF0000FFFF,16)]) \
|
|
if bits == 64 else (dtypes.uint32, [(0x55555555,1),(0x33333333,2),(0x0F0F0F0F,4),(0x00FF00FF,8)])
|
|
v = v.cast(dt) if v.dtype != dt else v
|
|
for m, s in masks: v = ((v >> _const(dt, s)) & _const(dt, m)) | ((v & _const(dt, m)) << _const(dt, s))
|
|
return (v >> _const(dt, 32 if bits == 64 else 16)) | (v << _const(dt, 32 if bits == 64 else 16))
|
|
|
|
def _extract_bits(val: UOp, hi: int, lo: int) -> UOp:
|
|
dt = dtypes.uint64 if val.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
|
|
width = hi - lo + 1
|
|
# Cast to dt first to ensure shift operands have matching types
|
|
val_cast = val.cast(dt) if val.dtype != dt else val
|
|
result = ((val_cast >> _const(dt, lo)) if lo > 0 else val_cast) & _const(dt, (1 << width) - 1)
|
|
# Downcast to match extracted bit width so brace-concat { hi, lo } computes correct output dtype
|
|
target_dt = _BITS_DT.get(width) or (dtypes.uint32 if width <= 32 else dtypes.uint64 if width <= 64 else dt)
|
|
if result.dtype != target_dt: result = result.cast(target_dt)
|
|
return result
|
|
|
|
def _expr_bits(v: UOp) -> int:
|
|
if v.dtype == dtypes.bool: return 1
|
|
if v.op in (Ops.AND, Ops.XOR):
|
|
widths: list[int] = []
|
|
for src in v.src:
|
|
if isinstance(sv:=_single_value(src), int) and sv > 0 and (sv & (sv + 1)) == 0:
|
|
widths.append(sv.bit_length())
|
|
if widths: return max(widths)
|
|
return v.dtype.bitsize
|
|
|
|
def _countbits(v: UOp) -> UOp:
|
|
dt = dtypes.uint64 if _expr_bits(v) > 32 or v.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
|
|
vv, out = v.cast(dt), _u32(0)
|
|
for i in range(_expr_bits(v)): out = out + ((vv >> _const(dt, i)) & _const(dt, 1)).cast(dtypes.uint32)
|
|
return out
|
|
|
|
def _set_bit(old, pos, val):
|
|
mask = _u32(1) << pos
|
|
return (old & (mask ^ _u32(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & _u32(1)) << pos)
|
|
|
|
def _val_to_bits(val):
|
|
if val.dtype == dtypes.half: return val.bitcast(dtypes.uint16).cast(dtypes.uint32)
|
|
if val.dtype == dtypes.float32: return val.bitcast(dtypes.uint32)
|
|
if val.dtype == dtypes.float64: return val.bitcast(dtypes.uint64)
|
|
return val if val.dtype == dtypes.uint32 else val.cast(dtypes.uint32)
|
|
|
|
def _floor(x):
|
|
t = UOp(Ops.TRUNC, src=(x,))
|
|
return ((x < _const(x.dtype, 0)) & x.ne(t)).where(t - _const(x.dtype, 1), t)
|
|
def _f16_extract(v): return (v & _u32(0xFFFF)).cast(dtypes.uint16).bitcast(dtypes.half) if v.dtype == dtypes.uint32 else v
|
|
|
|
# ═════ FP8 (E4M3) and BF8 (E5M2) conversion helpers ═════
|
|
# f32→fp8/bf8 uses f2f decomposition directly. fp8/bf8→f32 wraps f2f with subnormal handling
|
|
# (f2f flushes denormals to zero, but AMD V_CVT_F32_FP8/BF8 preserves subnormals).
|
|
def _fp8_to_f32(v: UOp) -> UOp:
|
|
b = (v.cast(dtypes.uint32) & _u32(0xFF)).cast(dtypes.uint8)
|
|
# E4M3 subnormal: exp==0, mant!=0 -> (-1)^sign * 2^(1-7) * (mant/8) = (-1)^sign * mant * 2^(-9)
|
|
bu = b.cast(dtypes.uint32)
|
|
sign, exp, mant = (bu >> _u32(7)) << _u32(31), (bu >> _u32(3)) & _u32(0xF), bu & _u32(0x7)
|
|
is_sub = exp.eq(_u32(0)) & mant.ne(_u32(0))
|
|
sub_f32 = (mant.cast(dtypes.float32) * _const(dtypes.float32, 1.0/512.0)).bitcast(dtypes.uint32) | sign
|
|
normal = f2f(b, dtypes.fp8e4m3, dtypes.float32)
|
|
return is_sub.where(sub_f32.bitcast(dtypes.float32), normal)
|
|
|
|
def _bf8_to_f32(v: UOp) -> UOp:
|
|
b = (v.cast(dtypes.uint32) & _u32(0xFF)).cast(dtypes.uint8)
|
|
# E5M2 subnormal: exp==0, mant!=0 -> (-1)^sign * 2^(1-15) * (mant/4) = (-1)^sign * mant * 2^(-16)
|
|
bu = b.cast(dtypes.uint32)
|
|
sign, exp, mant = (bu >> _u32(7)) << _u32(31), (bu >> _u32(2)) & _u32(0x1F), bu & _u32(0x3)
|
|
is_sub = exp.eq(_u32(0)) & mant.ne(_u32(0))
|
|
sub_f32 = (mant.cast(dtypes.float32) * _const(dtypes.float32, 1.0/65536.0)).bitcast(dtypes.uint32) | sign
|
|
normal = f2f(b, dtypes.fp8e5m2, dtypes.float32)
|
|
return is_sub.where(sub_f32.bitcast(dtypes.float32), normal)
|
|
|
|
def _f32_to_fp8(v: UOp) -> UOp:
|
|
return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e4m3, sat=False)
|
|
def _f32_to_bf8(v: UOp) -> UOp:
|
|
return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e5m2, sat=False)
|
|
def _f32_to_bf16(v: UOp) -> UOp:
|
|
"""Convert f32 to bf16 with round-to-nearest-even. BF16 is the upper 16 bits of F32 with rounding."""
|
|
bits = (v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32)
|
|
# Round-to-nearest-even: add rounding bias. If the bit just below the truncation point is 1 and the rest are 0, round to even.
|
|
round_bit = (bits >> _u32(16)) & _u32(1) # bit 16 (LSB of kept part)
|
|
rounding = _u32(0x7FFF) + round_bit # 0x7FFF + bit16: rounds to even
|
|
rounded = bits + rounding
|
|
return (rounded >> _u32(16)).cast(dtypes.uint16)
|
|
def _f32_to_bf16_sr(v: UOp, stoch: UOp) -> UOp:
|
|
"""Convert f32 to bf16 with stochastic rounding."""
|
|
bits = (v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32)
|
|
# Stochastic rounding: add lower 16 bits of stochastic value to lower 16 bits of f32
|
|
rounded = bits + (stoch & _u32(0xFFFF))
|
|
return (rounded >> _u32(16)).cast(dtypes.uint16)
|
|
|
|
def _check_nan(v: UOp, quiet: bool) -> UOp:
|
|
if v.op == Ops.CAST and v.dtype == dtypes.float64: v = v.src[0]
|
|
bits, exp_m, mant_m, qb, _ = _float_info(v)
|
|
is_nan_exp, has_mant, is_q = (bits & exp_m).eq(exp_m), (bits & mant_m).ne(_const(bits.dtype, 0)), (bits & qb).ne(_const(bits.dtype, 0))
|
|
return (is_nan_exp & is_q) if quiet else (is_nan_exp & has_mant & is_q.logical_not())
|
|
|
|
def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
|
|
def cast(v: UOp) -> UOp: return v.bitcast(dt) if dt == dtypes.float32 and v.dtype == dtypes.uint32 else v.cast(dt)
|
|
def minmax(a: UOp, b: UOp) -> UOp:
|
|
if dt in (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64): return (a > b).where(a, b) if is_max else (a < b).where(a, b)
|
|
return a.maximum(b) if is_max else a.minimum(b)
|
|
result = cast(args[0])
|
|
for a in args[1:]:
|
|
b = cast(a)
|
|
if dt == dtypes.float32: result = _isnan(result).where(b, _isnan(b).where(result, minmax(result, b)))
|
|
else: result = minmax(result, b)
|
|
return result
|
|
|
|
def _find_two_pi_mul(x):
|
|
if x.op != Ops.MUL or len(x.src) != 2: return None
|
|
for i, s in enumerate(x.src):
|
|
if (sv:=_single_value(s)) is not None and abs(sv - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
|
|
if s.op == Ops.MUL and len(s.src) == 2:
|
|
vals = [sv for ss in s.src if (sv:=_single_value(ss)) is not None]
|
|
if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
|
|
return None
|
|
|
|
def _trig_reduce(x, phase=0.0):
|
|
match = _find_two_pi_mul(x)
|
|
if match is not None:
|
|
turns, two_pi = match
|
|
if phase: turns = turns + _const(turns.dtype, phase)
|
|
n = _floor(turns + _const(turns.dtype, 0.5))
|
|
return UOp(Ops.SIN, src=((turns - n) * _const(turns.dtype, two_pi),))
|
|
if phase: x = x + _const(x.dtype, phase * 6.283185307179586)
|
|
n = _floor(x * _const(x.dtype, 0.15915494309189535) + _const(x.dtype, 0.5))
|
|
return UOp(Ops.SIN, src=(x - n * _const(x.dtype, 6.283185307179586),))
|
|
|
|
def _signext(val: UOp) -> UOp:
|
|
for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
|
|
if (val.op == Ops.AND and len(val.src) == 2 and _single_value(val.src[1]) == mask) or val.dtype.itemsize == bits // 8:
|
|
v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
|
|
sb = (v32 >> _u32(bits - 1)) & _u32(1)
|
|
return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
|
|
return val.cast(dtypes.int64) if val.dtype in (dtypes.int, dtypes.int32) else val
|
|
|
|
def _signext_4bit(val: UOp) -> UOp:
|
|
"""Sign extend a 4-bit value to 32-bit signed integer."""
|
|
v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
|
|
sb = (v32 >> _u32(3)) & _u32(1) # sign bit at position 3
|
|
return sb.ne(_u32(0)).where(v32 | _u32(0xFFFFFFF0), v32).bitcast(dtypes.int)
|
|
|
|
def _abs(val: UOp) -> UOp:
|
|
if val.dtype not in (dtypes.float32, dtypes.float64, dtypes.half): return val
|
|
_, _, _, _, shift = _float_info(val)
|
|
sign_mask = {10: 0x7FFF, 23: 0x7FFFFFFF, 52: 0x7FFFFFFFFFFFFFFF}[shift]
|
|
bt, ft = {10: (dtypes.uint16, dtypes.half), 23: (dtypes.uint32, dtypes.float32), 52: (dtypes.uint64, dtypes.float64)}[shift]
|
|
return (val.bitcast(bt) & _const(bt, sign_mask)).bitcast(ft)
|
|
|
|
def _f_to_u(f, dt):
|
|
clamped = (f < _const(f.dtype, 0.0)).where(_const(f.dtype, 0.0), f)
|
|
truncated = UOp(Ops.TRUNC, src=(clamped,))
|
|
res = (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
|
|
return _isnan(f).where(_const(dt, 0), res) # float->uint conversion of NaN is 0 on hardware
|
|
|
|
def _f_to_i32(a: UOp) -> UOp:
|
|
"""v_cvt_i32_f32/f64: truncate toward zero, saturate to [INT_MIN, INT_MAX], NaN -> 0.
|
|
(x86 cvttss2si returns 0x80000000 for all of these, which matches hardware only for negative overflow.)"""
|
|
res = (a >= _const(a.dtype, 2147483648.0)).where(_const(dtypes.int, 0x7FFFFFFF), UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int))
|
|
return _isnan(a).where(_const(dtypes.int, 0), res)
|
|
|
|
def _ftz_f32(v: UOp) -> UOp:
|
|
"""Flush f32 denormals to signed zero (RDNA default float mode flushes denormal f32 inputs on select-style ops)."""
|
|
bits = v.bitcast(dtypes.uint32) if v.dtype == dtypes.float32 else v
|
|
return ((bits & _u32(0x7FFFFFFF)) < _u32(0x00800000)).where((bits & _u32(0x80000000)).bitcast(dtypes.float32),
|
|
v if v.dtype == dtypes.float32 else v.bitcast(dtypes.float32))
|
|
|
|
def _cvt_quiet(val: UOp) -> UOp:
|
|
bits, _, _, qb, _ = _float_info(val)
|
|
bt, ft = (dtypes.uint64, dtypes.float64) if val.dtype == dtypes.float64 else \
|
|
(dtypes.uint16, dtypes.half) if val.dtype == dtypes.half else (dtypes.uint32, dtypes.float32)
|
|
return (val.bitcast(bt) | qb).bitcast(ft)
|
|
|
|
def _is_denorm(val: UOp) -> UOp:
|
|
bits, exp_m, mant_m, _, _ = _float_info(val)
|
|
return (bits & exp_m).eq(_const(bits.dtype, 0)) & (bits & mant_m).ne(_const(bits.dtype, 0))
|
|
|
|
_EXP_BITS = {10: 0x1F, 23: 0xFF, 52: 0x7FF}
|
|
def _get_exp(bits: UOp, shift: int) -> UOp: return ((bits >> _const(bits.dtype, shift)) & _const(bits.dtype, _EXP_BITS[shift])).cast(dtypes.int)
|
|
|
|
def _exponent(val: UOp) -> UOp:
|
|
bits, _, _, _, shift = _float_info(val)
|
|
return _get_exp(bits, shift)
|
|
|
|
def _div_would_be_denorm(a: UOp, b: UOp) -> UOp:
|
|
bits_n, _, _, _, shift = _float_info(a)
|
|
bits_d, _, _, _, _ = _float_info(b)
|
|
min_exp = {10: -14, 23: -126, 52: -1022}[shift]
|
|
return (_get_exp(bits_n, shift) - _get_exp(bits_d, shift)) < _const(dtypes.int, min_exp)
|
|
|
|
def _sign(val: UOp) -> UOp:
|
|
bits, _, _, _, shift = _float_info(val)
|
|
sign_shift = {10: 15, 23: 31, 52: 63}[shift]
|
|
return ((bits >> _const(bits.dtype, sign_shift)) & _const(bits.dtype, 1)).cast(dtypes.uint32)
|
|
|
|
def _signext_from_bit(val: UOp, w: UOp) -> UOp:
|
|
is_64bit = val.dtype in (dtypes.uint64, dtypes.int64)
|
|
dt = dtypes.uint64 if is_64bit else dtypes.uint32
|
|
mask_all = _const(dt, 0xFFFFFFFFFFFFFFFF if is_64bit else 0xFFFFFFFF)
|
|
one = _const(dt, 1)
|
|
val_u = val.cast(dt) if val.dtype != dt else val
|
|
w_val = w.cast(dt) if w.dtype != dt else w
|
|
sign_bit = (val_u >> (w_val - one)) & one
|
|
ext_mask = ((one << w_val) - one) ^ mask_all
|
|
return sign_bit.ne(_const(dt, 0)).where(val_u | ext_mask, val_u)
|
|
|
|
def _ldexp(val: UOp, exp: UOp) -> UOp:
|
|
if val.dtype == dtypes.uint32: val = val.bitcast(dtypes.float32)
|
|
elif val.dtype == dtypes.uint64: val = val.bitcast(dtypes.float64)
|
|
if exp.dtype in (dtypes.uint32, dtypes.uint64): exp = exp.cast(dtypes.int if exp.dtype == dtypes.uint32 else dtypes.int64)
|
|
bits = val.bitcast(dtypes.uint32) if val.dtype == dtypes.float32 else val.bitcast(dtypes.uint64)
|
|
abs_max = _const(bits.dtype, 0x7F800000 if val.dtype == dtypes.float32 else 0x7FF0000000000000)
|
|
sign_mask = _const(bits.dtype, 0x80000000 if val.dtype == dtypes.float32 else 0x8000000000000000)
|
|
# hardware flushes denormal inputs to signed zero
|
|
magn_mask = _const(bits.dtype, 0x7FFFFFFF if val.dtype == dtypes.float32 else 0x7FFFFFFFFFFFFFFF)
|
|
is_denorm = ((bits & abs_max).eq(_const(bits.dtype, 0))) & ((bits & magn_mask).ne(_const(bits.dtype, 0)))
|
|
val = is_denorm.where((bits & sign_mask).bitcast(val.dtype), val)
|
|
# hardware propagates 0/+-inf/NaN unchanged (avoids 0*inf = NaN on the host)
|
|
res = val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
|
|
is_special = (bits & abs_max).eq(_const(bits.dtype, 0)) | ((bits & abs_max) >= abs_max)
|
|
return is_special.where(val, res)
|
|
|
|
def _frexp_mant(val: UOp) -> UOp:
|
|
val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
|
|
if val.dtype == dtypes.float32:
|
|
bits = val.bitcast(dtypes.uint32)
|
|
# denormal/zero inputs (exponent field == 0) return signed zero on hardware
|
|
return ((bits & _u32(0x7F800000)).ne(_u32(0))).where(((bits & _u32(0x807FFFFF)) | _u32(0x3F000000)).bitcast(dtypes.float32),
|
|
(bits & _u32(0x80000000)).bitcast(dtypes.float32))
|
|
bits = val.bitcast(dtypes.uint64)
|
|
return ((bits & _const(dtypes.uint64, 0x7FF0000000000000)).ne(_const(dtypes.uint64, 0))).where(
|
|
((bits & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) | _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64),
|
|
(bits & _const(dtypes.uint64, 0x8000000000000000)).bitcast(dtypes.float64))
|
|
|
|
def _msb(val: UOp, bits: int) -> UOp:
|
|
"""Index of the highest set bit, or -1 if val == 0."""
|
|
dt = dtypes.uint64 if bits > 32 else dtypes.uint32
|
|
val = val.cast(dt) if val.dtype != dt else val
|
|
result = _const(dtypes.int, -1)
|
|
for i in range(bits - 1, -1, -1):
|
|
cond = ((val >> _const(dt, i)) & _const(dt, 1)).ne(_const(dt, 0)) & result.eq(_const(dtypes.int, -1))
|
|
result = cond.where(_const(dtypes.int, i), result)
|
|
return result
|
|
|
|
def _frexp_exp(val: UOp) -> UOp:
|
|
val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
|
|
if val.dtype == dtypes.float32:
|
|
e = (val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)
|
|
return e.ne(_u32(0)).where(e.cast(dtypes.int) - _const(dtypes.int, 126), _const(dtypes.int, 0)) # f32 denormals -> 0 (hardware verified)
|
|
bits = val.bitcast(dtypes.uint64)
|
|
e = (bits >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)
|
|
mant = bits & _const(dtypes.uint64, 0xFFFFFFFFFFFFF)
|
|
# f64 denormals: normalized exponent = highest set mantissa bit - 1073, zero -> 0 (hardware verified)
|
|
denorm = mant.ne(_const(dtypes.uint64, 0)).where(_msb(mant, 52) - _const(dtypes.int, 1073), _const(dtypes.int, 0))
|
|
return e.ne(_const(dtypes.uint64, 0)).where(e.cast(dtypes.int) - _const(dtypes.int, 1022), denorm)
|
|
|
|
TWO_OVER_PI = int(
|
|
"0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd"
|
|
"63962534e7dd1046bea5d768909d338e04d68befc827323ac7306a673e93908bf177bf250763ff12fffbc0b301fde5e2316b414"
|
|
"da3eda6cfd9e4f96136e9e8c7ecd3cbfd45aea4f758fd7cbe2f67a0e73ef14a525d4d7f6bf623f1aba10ac06608df8f6", 16)
|
|
# TWO_OVER_PI as 19 u64 words for trig_preop_result (word[0] = bits 0-63, word[18] = bits 1152-1200)
|
|
_PREOP_WORDS = tuple((TWO_OVER_PI >> (64 * i)) & 0xFFFFFFFFFFFFFFFF for i in range(19))
|
|
def _trig_preop(val: UOp) -> UOp:
|
|
# Extract 53 bits from position (1148 - shift) in the 1201-bit 2/PI constant
|
|
# Using word-based selection: 19 conditions instead of 1149
|
|
shift = val.cast(dtypes.uint32)
|
|
bit_pos = _u32(1148) - shift # starting bit position from LSB
|
|
word_idx = bit_pos >> _u32(6) # // 64
|
|
bit_off = bit_pos & _u32(63) # % 64
|
|
# Select lo_word and hi_word using shared conditions
|
|
lo_word, hi_word = _u64(_PREOP_WORDS[18]), _u64(0)
|
|
for i in range(17, -1, -1):
|
|
cond = word_idx.eq(_u32(i))
|
|
lo_word = cond.where(_u64(_PREOP_WORDS[i]), lo_word)
|
|
hi_word = cond.where(_u64(_PREOP_WORDS[i + 1]), hi_word)
|
|
# Combine and extract 53 bits: ((lo >> bit_off) | (hi << (64 - bit_off))) & mask
|
|
bit_off_64 = bit_off.cast(dtypes.uint64)
|
|
result = ((lo_word >> bit_off_64) | (hi_word << (_u64(64) - bit_off_64))) & _u64(0x1fffffffffffff)
|
|
return result.cast(dtypes.float64)
|
|
|
|
def _ff1(val: UOp, bits: int) -> UOp:
|
|
dt = dtypes.uint64 if bits == 64 else dtypes.uint32
|
|
val = val.cast(dt) if val.dtype != dt else val
|
|
result = _const(dtypes.int, -1)
|
|
for i in range(bits):
|
|
cond = ((val >> _const(dt, i)) & _const(dt, 1)).ne(_const(dt, 0)) & result.eq(_const(dtypes.int, -1))
|
|
result = cond.where(_const(dtypes.int, i), result)
|
|
return result
|
|
|
|
def _sad_u8(a: UOp, b: UOp, acc: UOp, masked: bool = False) -> UOp:
|
|
"""Sum of absolute differences of 4 unsigned bytes + accumulator. If masked, skips bytes where a == 0."""
|
|
a, b, acc = a.cast(dtypes.uint32), b.cast(dtypes.uint32), acc.cast(dtypes.uint32)
|
|
result = acc
|
|
for i in range(4):
|
|
a_byte = (a >> _u32(i * 8)) & _u32(0xFF)
|
|
b_byte = (b >> _u32(i * 8)) & _u32(0xFF)
|
|
diff = (a_byte > b_byte).where(a_byte - b_byte, b_byte - a_byte)
|
|
result = result + (a_byte.ne(_u32(0)).where(diff, _u32(0)) if masked else diff)
|
|
return result
|
|
|
|
_FUNCS: dict[str, Callable[..., UOp]] = {
|
|
'sqrt': lambda a: UOp(Ops.SQRT, src=(a,)), 'trunc': lambda a: UOp(Ops.TRUNC, src=(a,)),
|
|
'log2': lambda a: UOp(Ops.LOG2, src=(a,)), 'sin': lambda a: _trig_reduce(a),
|
|
'cos': lambda a: _trig_reduce(a, 0.25), 'floor': _floor, 'fract': lambda a: a - _floor(a),
|
|
'signext': _signext, 'abs': _abs,
|
|
'isEven': lambda a: (UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int) & _const(dtypes.int, 1)).eq(_const(dtypes.int, 0)),
|
|
'max': lambda a, b: UOp(Ops.MAX, src=(a, b)),
|
|
'min': lambda a, b: UOp(Ops.MAX, src=(a.neg(), b.neg())).neg(),
|
|
'pow': lambda a, b: UOp(Ops.EXP2, src=(b.bitcast(dtypes.float32),)),
|
|
'fma': lambda a, b, c: a * b + c,
|
|
'i32_to_f32': lambda a: a.cast(dtypes.int).cast(dtypes.float32),
|
|
'u32_to_f32': lambda a: a.cast(dtypes.uint32).cast(dtypes.float32),
|
|
'f32_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float32)),
|
|
'f32_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint32),
|
|
'f64_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float64)),
|
|
'f64_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float64), dtypes.uint32),
|
|
'f16_to_f32': lambda a: _f16_extract(a).cast(dtypes.float32),
|
|
'f32_to_f16': lambda a: a.cast(dtypes.half),
|
|
'f32_to_f64': lambda a: a.bitcast(dtypes.float32).cast(dtypes.float64),
|
|
'f64_to_f32': lambda a: a.bitcast(dtypes.float64).cast(dtypes.float32),
|
|
'i32_to_f64': lambda a: a.cast(dtypes.int).cast(dtypes.float64),
|
|
'u32_to_f64': lambda a: a.cast(dtypes.uint32).cast(dtypes.float64),
|
|
'f16_to_i16': lambda a: UOp(Ops.TRUNC, src=(_f16_extract(a),)).cast(dtypes.int16),
|
|
'f16_to_u16': lambda a: UOp(Ops.TRUNC, src=(_f16_extract(a),)).cast(dtypes.uint16),
|
|
'i16_to_f16': lambda a: a.cast(dtypes.int16).cast(dtypes.half),
|
|
'u16_to_f16': lambda a: a.cast(dtypes.uint16).cast(dtypes.half),
|
|
'bf16_to_f32': lambda a: (((a.cast(dtypes.uint32) if a.dtype != dtypes.uint32 else a) & _u32(0xFFFF)) << _u32(16)).bitcast(dtypes.float32),
|
|
'isNAN': _isnan, 'isSignalNAN': lambda a: _check_nan(a, False),
|
|
'isQuietNAN': lambda a: _check_nan(a, True), 'cvtToQuietNAN': _cvt_quiet,
|
|
'isDENORM': _is_denorm, 'exponent': _exponent, 'divWouldBeDenorm': _div_would_be_denorm, 'sign': _sign,
|
|
'signext_from_bit': _signext_from_bit, 'ldexp': _ldexp, 'frexp_mant': _frexp_mant, 'mantissa': _frexp_mant,
|
|
'frexp_exp': _frexp_exp, 'trig_preop_result': _trig_preop,
|
|
's_ff1_i32_b32': lambda a: _ff1(a, 32), 's_ff1_i32_b64': lambda a: _ff1(a, 64),
|
|
# Normalization conversions: map [-1,1] or [0,1] to integer range
|
|
# Use floor(x + 0.5) for round-to-nearest
|
|
# SNORM: round(value * 32767), range is [-32767, 32767] (hardware behavior)
|
|
'f16_to_snorm': lambda a: _floor(
|
|
_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
|
|
'f16_to_unorm': lambda a: _floor(
|
|
_f16_extract(a).cast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
|
|
'f32_to_snorm': lambda a: _floor(
|
|
a.bitcast(dtypes.float32) * _const(dtypes.float32, 32767) + _const(dtypes.float32, 0.5)).cast(dtypes.int).cast(dtypes.int16),
|
|
'f32_to_unorm': lambda a: _floor(
|
|
a.bitcast(dtypes.float32) * _const(dtypes.float32, 65535) + _const(dtypes.float32, 0.5)).cast(dtypes.uint16),
|
|
'f32_to_u8': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint8),
|
|
# Integer truncation conversions
|
|
'i32_to_i16': lambda a: a.cast(dtypes.int).cast(dtypes.int16),
|
|
'u32_to_u16': lambda a: a.cast(dtypes.uint32).cast(dtypes.uint16),
|
|
'u16_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xFFFF)),
|
|
'u8_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xFF)),
|
|
'u4_to_u32': lambda a: (a.cast(dtypes.uint32) & _u32(0xF)),
|
|
# Signed extraction with sign extension for dot products
|
|
'i16_to_i32': lambda a: _signext(a.cast(dtypes.uint32) & _u32(0xFFFF)),
|
|
'i8_to_i32': lambda a: _signext(a.cast(dtypes.uint32) & _u32(0xFF)),
|
|
'i4_to_i32': lambda a: _signext_4bit(a.cast(dtypes.uint32) & _u32(0xF)),
|
|
# Float to int16 conversions
|
|
'v_cvt_i16_f32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float32),)).cast(dtypes.int16),
|
|
'v_cvt_u16_f32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint16),
|
|
# SAD (Sum of Absolute Differences) - sum |a_i - b_i| for 4 bytes + accumulator
|
|
'v_sad_u8': lambda a, b, c: _sad_u8(a, b, c),
|
|
'v_msad_u8': lambda a, b, c: _sad_u8(a, b, c, masked=True),
|
|
# System NOPs - these are scheduling hints, no effect on emulation
|
|
'MIN': lambda a, b: (a < b).where(a, b),
|
|
's_nop': lambda a: _u32(0),
|
|
'countbits': _countbits,
|
|
# Address calculation for memory operations
|
|
'CalcDsAddr': lambda a, o, *r: a.cast(dtypes.uint32) + o.cast(dtypes.uint32),
|
|
'CalcGlobalAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
|
|
'CalcScratchAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
|
|
# FP8/BF8/BF16 conversion functions
|
|
'fp8_to_f32': _fp8_to_f32, 'bf8_to_f32': _bf8_to_f32, 'f32_to_fp8': _f32_to_fp8, 'f32_to_bf8': _f32_to_bf8,
|
|
'f32_to_bf16': _f32_to_bf16, 'f32_to_bf16_SR': _f32_to_bf16_sr, 'f32_to_bf16_sr': _f32_to_bf16_sr,
|
|
}
|
|
# min/max family: min/max + 3-input (x3), IEEE num variants (f16/f32 only), and long names minimum/maximum (f16/f32 only)
|
|
for is_max, name, full in [(False, 'min', 'minimum'), (True, 'max', 'maximum')]:
|
|
for dt, sfx, pre in [(dtypes.float32, 'f32', None), (dtypes.int, 'i32', None), (dtypes.uint32, 'u32', None),
|
|
(dtypes.int16, 'i16', None), (dtypes.uint16, 'u16', None), (dtypes.half, 'f16', _f16_extract)]:
|
|
def mm(*a, im=is_max, d=dt, p=pre): return _minmax_reduce(im, d, *(a if p is None else [p(x) for x in a]))
|
|
extra = (f'v_{name}_num_{sfx}', f'v_{name}3_num_{sfx}', f'v_{full}_{sfx}', f'v_{full}3_{sfx}') if dt in (dtypes.float32, dtypes.half) else ()
|
|
for fn in (f'v_{name}_{sfx}', f'v_{name}3_{sfx}', *extra): _FUNCS[fn] = mm
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════════
|
|
# TOKENIZER/PARSER
|
|
# ═══════════════════════════════════════════════════════════════════════════════
|
|
|
|
DTYPES = {'u32': dtypes.uint32, 'i32': dtypes.int, 'f32': dtypes.float32, 'b32': dtypes.uint32, 'u64': dtypes.uint64, 'i64': dtypes.int64,
|
|
'f64': dtypes.float64, 'b64': dtypes.uint64, 'u16': dtypes.uint16, 'i16': dtypes.short, 'f16': dtypes.half, 'b16': dtypes.uint16,
|
|
'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u4': dtypes.uint8, 'i4': dtypes.int8, 'u1': dtypes.uint32,
|
|
'fp8': dtypes.uint8, 'bf8': dtypes.uint8, 'b3': dtypes.uint8, 'b2': dtypes.uint8}
|
|
_BITS_DT = {8: dtypes.uint8, 16: dtypes.uint16, 32: dtypes.uint32, 64: dtypes.uint64}
|
|
_NUM_SUFFIXES = ('ULL', 'LL', 'UL', 'U', 'L', 'F', 'f')
|
|
def _strip_suffix(num: str) -> tuple[str, str]:
|
|
for sfx in _NUM_SUFFIXES:
|
|
if num.endswith(sfx): return sfx, num[:-len(sfx)]
|
|
return '', num
|
|
_SINGLE_CHAR = {'(': 'LPAREN', ')': 'RPAREN', '[': 'LBRACKET', ']': 'RBRACKET', '{': 'LBRACE', '}': 'RBRACE',
|
|
':': 'COLON', ',': 'COMMA', '?': 'QUESTION', '.': 'DOT', '=': 'EQUALS', "'": 'QUOTE'}
|
|
|
|
class Token:
|
|
__slots__ = ('type', 'val')
|
|
def __init__(self, kind: str, val: str): self.type, self.val = kind, val
|
|
def __repr__(self): return f'{self.type}:{self.val}'
|
|
|
|
def tokenize(s: str) -> list[Token]:
|
|
tokens, i, n = [], 0, len(s)
|
|
while i < n:
|
|
c = s[i]
|
|
if c.isspace():
|
|
i += 1
|
|
continue
|
|
if i + 1 < n and s[i:i+2] in ('+=', '-=', '|=', '&=', '^='):
|
|
tokens.append(Token('ASSIGN_OP', s[i:i+2]))
|
|
i += 2
|
|
continue
|
|
if i + 1 < n and s[i:i+2] in ('||', '&&', '>=', '<=', '==', '!=', '<>', '>>', '<<', '**', '+:', '-:'):
|
|
tokens.append(Token('OP', s[i:i+2]))
|
|
i += 2
|
|
continue
|
|
if c in '|^&><+-*/~!%':
|
|
tokens.append(Token('OP', c))
|
|
i += 1
|
|
continue
|
|
if (t := _SINGLE_CHAR.get(c)):
|
|
tokens.append(Token(t, c))
|
|
i += 1
|
|
continue
|
|
if c == ';':
|
|
i += 1
|
|
continue
|
|
if c.isdigit() or (c == '-' and i + 1 < n and s[i+1].isdigit()):
|
|
start = i
|
|
if c == '-': i += 1
|
|
if i + 1 < n and s[i] == '0' and s[i+1] in 'xX':
|
|
i += 2
|
|
while i < n and s[i] in '0123456789abcdefABCDEF': i += 1
|
|
else:
|
|
while i < n and s[i].isdigit(): i += 1
|
|
if i < n and s[i] == '.' and i + 1 < n and s[i+1].isdigit():
|
|
i += 1
|
|
while i < n and s[i].isdigit(): i += 1
|
|
for sfx in ('ULL', 'LL', 'UL', 'U', 'L', 'F', 'f'):
|
|
if s[i:i+len(sfx)] == sfx:
|
|
i += len(sfx)
|
|
break
|
|
tokens.append(Token('NUM', s[start:i]))
|
|
continue
|
|
if c.isalpha() or c == '_':
|
|
start = i
|
|
while i < n and (s[i].isalnum() or s[i] == '_'): i += 1
|
|
tokens.append(Token('IDENT', s[start:i]))
|
|
continue
|
|
raise RuntimeError(f"unexpected char '{c}' at pos {i} in: {s}")
|
|
tokens.append(Token('EOF', ''))
|
|
return tokens
|
|
|
|
class Parser:
|
|
def __init__(self, tokens: list[Token], env: dict, funcs: dict | None = None):
|
|
self.tokens, self.vars, self.funcs, self.pos = tokens, env, funcs if funcs is not None else _FUNCS, 0
|
|
|
|
def peek(self, offset=0) -> Token: return self.tokens[min(self.pos + offset, len(self.tokens) - 1)]
|
|
def at(self, *types) -> bool: return self.peek().type in types
|
|
def _advance(self) -> Token:
|
|
tok = self.tokens[self.pos]
|
|
self.pos += 1
|
|
return tok
|
|
def eat(self, kind: str) -> Token:
|
|
if self.peek().type != kind: raise RuntimeError(f"expected {kind}, got {self.peek()}")
|
|
return self._advance()
|
|
def try_eat(self, kind: str) -> Token | None: return self._advance() if self.peek().type == kind else None
|
|
def try_eat_val(self, val: str, kind: str) -> Token | None:
|
|
return self._advance() if self.peek().type == kind and self.peek().val == val else None
|
|
def eat_val(self, val: str, kind: str) -> Token:
|
|
if self.peek().type != kind or self.peek().val != val: raise RuntimeError(f"expected {kind}:{val}, got {self.peek()}")
|
|
return self._advance()
|
|
|
|
def parse(self) -> UOp:
|
|
cond = self.binop(0)
|
|
if self.try_eat('QUESTION'):
|
|
then_val = self.parse()
|
|
self.eat('COLON')
|
|
return _to_bool(cond).where(then_val, self.parse())
|
|
return cond
|
|
|
|
def _apply_binop(self, left, right, op):
|
|
if op in ('||', '&&', '|', '^', '&'): left, right = self._coerce_bitwise(left, right)
|
|
elif op in ('>=', '<=', '>', '<', '==', '!=', '<>', '>>', '<<'): left, right = self._coerce_cmp(left, right)
|
|
elif left.dtype != right.dtype: right = right.cast(left.dtype)
|
|
match op:
|
|
case '||' | '|': return left | right
|
|
case '&&' | '&': return left & right
|
|
case '^': return left ^ right
|
|
case '==': return left.eq(right)
|
|
case '!=': return left.ne(right)
|
|
case '>=' | '<=' | '>' | '<' | '<>':
|
|
ops = {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b),'<>':(lambda a,b:a.ne(b))}
|
|
return self._cmp_nan(left, right, ops[op])
|
|
case '>>' | '<<':
|
|
if not dtypes.is_int(left.dtype): left = left.cast(dtypes.uint32)
|
|
if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
|
|
return (left >> right) if op == '>>' else (left << right)
|
|
case '+' | '-':
|
|
if op == '-' and (lv:=_single_value(left)) is not None and (rv:=_single_value(right)) is not None: return _const(left.dtype, lv - rv)
|
|
return (left + right) if op == '+' else (left - right)
|
|
case '*' | '/':
|
|
# Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
|
|
# (e.g. SOPP branch offset: SIMM16.i16 * 16'4 can exceed int16 range)
|
|
if op == '*' and left.dtype.itemsize == 2 and left.dtype in (dtypes.int16, dtypes.short, dtypes.uint16, dtypes.ushort):
|
|
pdt = dtypes.int if left.dtype in (dtypes.int16, dtypes.short) else dtypes.uint
|
|
left, right = left.cast(pdt), right.cast(pdt)
|
|
if op == '*': return left * right
|
|
return (left // right) if dtypes.is_int(left.dtype) else (left / right)
|
|
case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _single_value(left) == 2.0 else left
|
|
|
|
_PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
|
|
|
|
def binop(self, prec: int) -> UOp:
|
|
if prec >= len(self._PREC): return self.unary()
|
|
left = self.binop(prec + 1)
|
|
ops = self._PREC[prec]
|
|
while self.at('OP') and self.peek().val in ops:
|
|
op = self.eat('OP').val
|
|
left = self._apply_binop(left, self.binop(prec + 1), op)
|
|
return left
|
|
|
|
def unary(self) -> UOp:
|
|
if self.try_eat_val('~', 'OP'):
|
|
inner = self.unary()
|
|
return inner ^ _const(inner.dtype, (1 << _expr_bits(inner)) - 1)
|
|
if self.try_eat_val('!', 'OP'):
|
|
inner = self.unary()
|
|
return inner.eq(_const(inner.dtype, 0))
|
|
if self.try_eat_val('-', 'OP'):
|
|
inner = self.unary()
|
|
if (v:=_single_value(inner)) is not None:
|
|
return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -v)
|
|
return inner.neg()
|
|
if self.try_eat_val('+', 'OP'): return self.unary()
|
|
return self.postfix()
|
|
|
|
def postfix(self) -> UOp:
|
|
base = self.primary()
|
|
while True:
|
|
if self.try_eat('DOT'):
|
|
field = self.eat('IDENT').val
|
|
base = self._handle_dot(base, field)
|
|
elif self.at('LBRACKET'):
|
|
base = self._handle_bracket(base)
|
|
elif self.at('LBRACE'):
|
|
base = self._handle_brace_index(base)
|
|
else:
|
|
break
|
|
return base
|
|
|
|
def primary(self) -> UOp:
|
|
if self.try_eat('LPAREN'):
|
|
e = self.parse()
|
|
self.eat('RPAREN')
|
|
return e
|
|
if self.try_eat('LBRACE'):
|
|
hi = self.parse()
|
|
self.eat('COMMA')
|
|
lo = self.parse()
|
|
self.eat('RBRACE')
|
|
lo_bits, hi_bits = _expr_bits(lo), _expr_bits(hi)
|
|
total_bits = lo_bits + hi_bits
|
|
dt = _BITS_DT.get(total_bits, dtypes.uint32 if total_bits <= 32 else dtypes.uint64)
|
|
return (hi.cast(dt) << _const(dt, lo_bits)) | lo.cast(dt)
|
|
if self.at('NUM'):
|
|
num = self.eat('NUM').val
|
|
if self.try_eat('QUOTE'):
|
|
return self._sized_literal(int(num.rstrip('ULlf')))
|
|
return self._parse_number(num)
|
|
if self.at('IDENT'):
|
|
name = self.eat('IDENT').val
|
|
if name == 'MEM':
|
|
self.eat('LBRACKET')
|
|
addr = self.parse()
|
|
self.eat('RBRACKET')
|
|
self.eat('DOT')
|
|
dt_name = self.eat('IDENT').val
|
|
return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
|
|
if name in self.funcs and self.try_eat('LBRACKET'):
|
|
index = self.parse()
|
|
self.eat('RBRACKET')
|
|
return self.funcs[name](index)
|
|
if name == 'VGPR' and self.at('LBRACKET'):
|
|
self.eat('LBRACKET')
|
|
lane = self.parse()
|
|
self.eat('RBRACKET')
|
|
self.eat('LBRACKET')
|
|
reg = self.parse()
|
|
self.eat('RBRACKET')
|
|
vgpr = self.vars.get('_vgpr')
|
|
if vgpr is None: return _u32(0)
|
|
ws = self.vars.get('_wave_size', 32)
|
|
return vgpr.index(_to_u32(reg) * _u32(ws) + _to_u32(lane)).load()
|
|
if self.try_eat('LPAREN'):
|
|
args = self._parse_args()
|
|
self.eat('RPAREN')
|
|
return self._call_func(name, args)
|
|
if name == 'PI': return _const(dtypes.float32, 3.141592653589793)
|
|
if name == 'INF': return _const(dtypes.float64, float('inf'))
|
|
if name == 'NAN': return _const(dtypes.uint32, 0x7FC00000).bitcast(dtypes.float32)
|
|
if name == 'UNDERFLOW_F32': return _const(dtypes.uint32, 1).bitcast(dtypes.float32)
|
|
if name == 'OVERFLOW_F32': return _const(dtypes.uint32, 0x7F7FFFFF).bitcast(dtypes.float32)
|
|
if name == 'UNDERFLOW_F64': return _const(dtypes.uint64, 1).bitcast(dtypes.float64)
|
|
if name == 'OVERFLOW_F64': return _const(dtypes.uint64, 0x7FEFFFFFFFFFFFFF).bitcast(dtypes.float64)
|
|
if name.lower() == 'wave32': return _const(dtypes.bool, self.vars.get('_wave_size', 32) <= 32)
|
|
if name.lower() == 'wave64': return _const(dtypes.bool, self.vars.get('_wave_size', 32) > 32)
|
|
if name == 'WAVE_MODE' and self.try_eat('DOT') and self.try_eat_val('IEEE', 'IDENT'): return _u32(1)
|
|
if self.try_eat('LBRACE'):
|
|
idx = self.eat('NUM').val
|
|
self.eat('RBRACE')
|
|
# Handle VGPR{lane}[reg] - 2D array access after loop unrolling
|
|
if name == 'VGPR' and self.at('LBRACKET'):
|
|
self.eat('LBRACKET')
|
|
reg = self.parse()
|
|
self.eat('RBRACKET')
|
|
vgpr = self.vars.get('_vgpr')
|
|
if vgpr is None: return _u32(0)
|
|
ws = self.vars.get('_wave_size', 32)
|
|
return vgpr.index(_to_u32(reg) * _u32(ws) + _u32(int(idx))).load()
|
|
elem = self.vars.get(f'{name}@{idx}', self.vars.get(f'{name}{idx}'))
|
|
if elem is None:
|
|
# Extract bit idx from base variable (like var[idx])
|
|
base = self.vars.get(name)
|
|
assert isinstance(base, UOp), f"unknown variable: {name}{idx}"
|
|
dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
|
|
elem = (base.cast(dt) >> _const(dt, int(idx))) & _const(dt, 1)
|
|
if self.try_eat('DOT'):
|
|
dt_name = self.eat('IDENT').val
|
|
return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
|
|
if self.at('LBRACKET'):
|
|
return self._handle_bracket(elem, name + idx)
|
|
return elem
|
|
if self.at('LBRACKET') and name not in self.vars:
|
|
self.eat('LBRACKET')
|
|
first = self.parse()
|
|
return self._handle_bracket_rest(first, _u32(0), name)
|
|
if name in self.vars:
|
|
v = self.vars[name]
|
|
assert isinstance(v, UOp), f"expected UOp for {name}, got {type(v)}"
|
|
return v
|
|
raise RuntimeError(f"unknown variable: {name}")
|
|
raise RuntimeError(f"unexpected token in primary: {self.peek()}")
|
|
|
|
def _handle_dot(self, base, field: str) -> UOp:
|
|
assert isinstance(base, UOp), f"expected UOp for dot access, got {type(base)}"
|
|
if field == 'u64' and self.at('LBRACKET') and self.peek(1).type == 'IDENT' and self.peek(1).val == 'laneId':
|
|
self.eat('LBRACKET')
|
|
self.eat_val('laneId', 'IDENT')
|
|
self.eat('RBRACKET')
|
|
lane = self.vars['laneId']
|
|
shift = lane.cast(base.dtype) if base.dtype != dtypes.uint32 else _to_u32(lane)
|
|
result = (base >> shift) & _const(base.dtype, 1)
|
|
if self.try_eat('DOT'):
|
|
dt_name = self.eat('IDENT').val
|
|
return result.cast(DTYPES.get(dt_name, dtypes.uint32))
|
|
return result
|
|
dt = DTYPES.get(field)
|
|
if dt is None: return base
|
|
if dt == base.dtype: return base
|
|
if dt.itemsize == 2 and base.dtype.itemsize == 4:
|
|
if dt == dtypes.uint16: return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16)
|
|
return (base & _const(base.dtype, 0xFFFF)).cast(dtypes.uint16).bitcast(dt)
|
|
if field == 'i4': return _signext_4bit(base)
|
|
return _cast_to(base, dt)
|
|
|
|
def _handle_bracket(self, base, var_name: str | None = None) -> UOp:
|
|
self.eat('LBRACKET')
|
|
return self._handle_bracket_rest(self.parse(), base, var_name)
|
|
|
|
def _handle_bracket_rest(self, first: UOp, base: UOp, var_name: str | None = None) -> UOp:
|
|
if self.at('OP') and self.peek().val in ('+:', '-:'):
|
|
self.eat('OP')
|
|
width = self.parse()
|
|
self.eat('RBRACKET')
|
|
if isinstance(w:=_single_value(width), int):
|
|
return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
|
|
return base
|
|
if self.try_eat('COLON'):
|
|
second = self.parse()
|
|
self.eat('RBRACKET')
|
|
if isinstance(a:=_single_value(first), int) and isinstance(b:=_single_value(second), int):
|
|
if a < b: return _bitreverse(base, b - a + 1)
|
|
hi, lo = a, b
|
|
if lo >= base.dtype.itemsize * 8:
|
|
vn = var_name or self._find_var_name(base)
|
|
if vn and f'{vn}{lo // 32}' in self.vars:
|
|
base = self.vars[f'{vn}{lo // 32}']
|
|
lo, hi = lo % 32, (hi % 32) + (lo % 32)
|
|
return _extract_bits(base, hi, lo)
|
|
# Dynamic bit slice: (base >> lo) & ((1 << (hi - lo + 1)) - 1)
|
|
dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
|
|
hi_u, lo_u = first.cast(dt), second.cast(dt)
|
|
width = hi_u - lo_u + _const(dt, 1)
|
|
mask = (_const(dt, 1) << width) - _const(dt, 1)
|
|
return (base.cast(dt) >> lo_u) & mask
|
|
self.eat('RBRACKET')
|
|
dt_suffix = None
|
|
if self.try_eat('DOT'):
|
|
dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32)
|
|
if var_name is None:
|
|
var_name = self._find_var_name(base)
|
|
if isinstance(idx:=_single_value(first), int):
|
|
# Check for array element (var@idx)
|
|
if var_name and f'{var_name}@{idx}' in self.vars:
|
|
v = self.vars[f'{var_name}@{idx}']
|
|
return _cast_to(v, dt_suffix) if dt_suffix else v
|
|
# Bit extraction
|
|
dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
|
|
base_cast = base.cast(dt) if base.dtype != dt else base
|
|
result = ((base_cast >> _const(dt, idx)) & _const(dt, 1))
|
|
return _cast_to(result, dt_suffix) if dt_suffix else result.cast(dtypes.bool)
|
|
if var_name:
|
|
idx_u32 = _to_u32(first)
|
|
elems = [(i, self.vars[f'{var_name}@{i}']) for i in range(256) if f'{var_name}@{i}' in self.vars]
|
|
if elems:
|
|
result = elems[-1][1]
|
|
for ei, ev in reversed(elems[:-1]):
|
|
if ev.dtype != result.dtype and ev.dtype.itemsize == result.dtype.itemsize: result = result.cast(ev.dtype)
|
|
elif ev.dtype != result.dtype: ev = ev.cast(result.dtype)
|
|
result = idx_u32.eq(_u32(ei)).where(ev, result)
|
|
return result
|
|
dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
|
|
base_cast = base.cast(dt) if base.dtype != dt else base
|
|
result = (base_cast >> first.cast(dt)) & _const(dt, 1)
|
|
return _cast_to(result, dt_suffix) if dt_suffix else result.cast(dtypes.bool)
|
|
|
|
def _handle_brace_index(self, base) -> UOp:
|
|
self.eat('LBRACE')
|
|
idx = self.eat('NUM').val
|
|
self.eat('RBRACE')
|
|
var_name = self._find_var_name(base)
|
|
if var_name:
|
|
elem = self.vars.get(f'{var_name}@{idx}', _u32(0)) # use @ to avoid collision with temps like A4
|
|
if self.try_eat('DOT'):
|
|
dt_name = self.eat('IDENT').val
|
|
return _cast_to(elem, DTYPES.get(dt_name, dtypes.uint32))
|
|
if self.at('LBRACKET'):
|
|
return self._handle_bracket(elem)
|
|
return elem
|
|
return _u32(0)
|
|
|
|
def _find_var_name(self, base: UOp) -> str | None:
|
|
if base.op == Ops.PARAM and base.arg.name is not None: return base.arg.name
|
|
for name, v in self.vars.items():
|
|
if isinstance(v, UOp) and v is base: return name
|
|
return None
|
|
|
|
def _sized_literal(self, bits: int) -> UOp:
|
|
if self.at('IDENT') and self.peek().val in ('U', 'I', 'F', 'B', 'BF'):
|
|
type_char = self.eat('IDENT').val
|
|
self.eat('LPAREN')
|
|
inner = self.parse()
|
|
self.eat('RPAREN')
|
|
dt = {('U',32): dtypes.uint32, ('U',64): dtypes.uint64, ('I',32): dtypes.int, ('I',64): dtypes.int64,
|
|
('F',16): dtypes.half, ('F',32): dtypes.float32, ('F',64): dtypes.float64,
|
|
('BF',16): dtypes.bfloat16,
|
|
('B',32): dtypes.uint32, ('B',64): dtypes.uint64}.get((type_char, bits), dtypes.uint64 if bits > 32 else dtypes.uint32)
|
|
if type_char == 'F' and inner.dtype in (dtypes.uint32, dtypes.uint64, dtypes.ulong, dtypes.int, dtypes.int64):
|
|
if inner.dtype.itemsize != dt.itemsize: inner = inner.cast(dtypes.uint32 if dt.itemsize == 4 else dtypes.uint64)
|
|
return inner.bitcast(dt)
|
|
return inner.cast(dt)
|
|
if self.at('IDENT'):
|
|
ident = self.peek().val
|
|
fmt = ident[0].lower()
|
|
if fmt in ('h', 'b', 'd'):
|
|
self.eat('IDENT')
|
|
if len(ident) > 1: num = ident[1:]
|
|
elif self.at('NUM'): num = self.eat('NUM').val
|
|
elif self.at('IDENT'): num = self.eat('IDENT').val
|
|
else: raise RuntimeError(f"expected number after {bits}'{fmt}")
|
|
if fmt == 'h': val = int(num, 16)
|
|
elif fmt == 'b': val = int(num, 2)
|
|
else: val = int(num)
|
|
return _const(_BITS_DT.get(bits, dtypes.uint32), val)
|
|
if self.at('NUM') and self.peek().val.startswith('0x'):
|
|
num = self.eat('NUM').val
|
|
return _const(_BITS_DT.get(bits, dtypes.uint32), int(num, 16))
|
|
if self.at('NUM') or (self.at('OP') and self.peek().val == '-'):
|
|
neg = self.try_eat_val('-', 'OP') is not None
|
|
suffix, num = _strip_suffix(self.eat('NUM').val)
|
|
if num.startswith('0x'):
|
|
val = int(num, 16)
|
|
if neg: val = -val
|
|
elif '.' in num:
|
|
fval = float(num)
|
|
if neg: fval = -fval
|
|
return _const({16: dtypes.half, 32: dtypes.float32, 64: dtypes.float64}.get(bits, dtypes.float32), fval)
|
|
else:
|
|
val = int(num)
|
|
if neg: val = -val
|
|
dt = {1: dtypes.uint32, 8: dtypes.uint8, 16: dtypes.int16 if 'U' not in suffix else dtypes.uint16,
|
|
32: dtypes.int if 'U' not in suffix else dtypes.uint32, 64: dtypes.int64 if 'U' not in suffix else dtypes.uint64}.get(bits, dtypes.uint32)
|
|
return _const(dt, val)
|
|
raise RuntimeError(f"unexpected token after {bits}': {self.peek()}")
|
|
|
|
def _parse_number(self, num: str) -> UOp:
|
|
if num.startswith('0x') or num.startswith('0X'):
|
|
is_u64 = num.upper().endswith('ULL') or num.upper().endswith('LL') or num.upper().endswith('UL')
|
|
return _const(dtypes.uint64 if is_u64 else dtypes.uint32, int(num.rstrip('ULul'), 16))
|
|
suffix, num_str = _strip_suffix(num)
|
|
if '.' in num_str or suffix in ('F', 'f'):
|
|
return _const(dtypes.float32 if suffix in ('F', 'f') else dtypes.float64, float(num_str))
|
|
val = int(num_str)
|
|
if 'ULL' in suffix or 'LL' in suffix or 'L' in suffix: return _const(dtypes.uint64, val)
|
|
if 'U' in suffix: return _const(dtypes.uint32, val)
|
|
return _const(dtypes.int if val < 0 else dtypes.uint32, val)
|
|
|
|
def _parse_args(self) -> list[UOp]:
|
|
if self.at('RPAREN'): return []
|
|
args = [self.parse()]
|
|
while self.try_eat('COMMA'):
|
|
args.append(self.parse())
|
|
return args
|
|
|
|
def _call_func(self, name: str, args: list[UOp]) -> UOp:
|
|
if name in self.vars and isinstance(self.vars[name], tuple) and self.vars[name][0] == 'lambda':
|
|
_, params, body = self.vars[name]
|
|
lv = {**self.vars, **dict(zip(params, args))}
|
|
if ';' in body or '\n' in body or 'return' in body.lower():
|
|
lines = [l.strip() for l in body.replace(';', '\n').split('\n') if l.strip() and not l.strip().startswith('//')]
|
|
_, _, result = parse_block(lines, 0, lv, self.funcs)
|
|
assert result is not None, f"lambda {name} must return a value"
|
|
return result
|
|
return parse_expr(body, lv, self.funcs)
|
|
if name in self.funcs:
|
|
return self.funcs[name](*args)
|
|
raise RuntimeError(f"unknown function: {name}")
|
|
|
|
def _handle_mem_load(self, addr: UOp, dt) -> UOp:
|
|
mem = self.vars.get('_vmem') if '_vmem' in self.vars else self.vars.get('_lds')
|
|
assert mem is not None, "memory load requires _vmem or _lds"
|
|
adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
|
|
active = self.vars.get('_active')
|
|
def mindex(idx:UOp): return mem.index(idx.valid(active) if active is not None else idx)
|
|
byte_mem = mem.dtype == dtypes.uint8
|
|
if byte_mem:
|
|
idx = addr
|
|
if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
|
|
val = _u32(0).cast(dtypes.uint64)
|
|
for i in range(8): val = val | (mindex(idx + _const(adt, i)).load().cast(dtypes.uint64) << _u64(i * 8))
|
|
elif dt in (dtypes.uint8, dtypes.int8):
|
|
val = mindex(idx).load().cast(dt)
|
|
elif dt in (dtypes.uint16, dtypes.int16, dtypes.short):
|
|
lo = mindex(idx).load().cast(dtypes.uint32)
|
|
hi = mindex(idx + _const(adt, 1)).load().cast(dtypes.uint32)
|
|
val = (lo | (hi << _u32(8))).cast(dt)
|
|
else:
|
|
val = _u32(0)
|
|
for i in range(4): val = val | (mindex(idx + _const(adt, i)).load().cast(dtypes.uint32) << _u32(i * 8))
|
|
else:
|
|
idx = addr >> _const(addr.dtype, 2)
|
|
val = mindex(idx)
|
|
if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
|
|
idx2 = (addr + _const(adt, 4)) >> _const(adt, 2)
|
|
val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32))
|
|
elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
|
|
else:
|
|
# Handle unaligned 16/32-bit loads: combine two consecutive dwords and shift.
|
|
# The next dword is only read when the value straddles into it, so a load at the end of a buffer stays in bounds.
|
|
# Use int64 for the WHERE to avoid 32-bit int overflow in C pointer arithmetic (addr can be >8GB).
|
|
byte_off = (addr & _const(adt, 3)).cast(dtypes.uint32)
|
|
idx_native = (addr >> _const(adt, 2)).cast(dtypes.int64)
|
|
idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
|
|
hi = mindex((byte_off > _u32(4 - dt.itemsize)).where(idx_hi_native, idx_native))
|
|
combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
|
|
val = (combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32)
|
|
return _cast_to(val, dt)
|
|
|
|
def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
|
|
if l.dtype != r.dtype:
|
|
if r.dtype == dtypes.int and isinstance(rv:=_single_value(r), int) and rv < 0: l = l.cast(dtypes.int)
|
|
else: r = r.cast(l.dtype)
|
|
return l, r
|
|
|
|
def _coerce_bitwise(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
|
|
if l.dtype != r.dtype:
|
|
if l.dtype.itemsize == r.dtype.itemsize:
|
|
t = dtypes.uint32 if l.dtype.itemsize == 4 else dtypes.uint64 if l.dtype.itemsize == 8 else l.dtype
|
|
l, r = l.bitcast(t), r.bitcast(t)
|
|
else: r = r.cast(l.dtype)
|
|
return l, r
|
|
|
|
def _cmp_nan(self, l: UOp, r: UOp, fn) -> UOp:
|
|
result = fn(l, r)
|
|
if l.dtype in (dtypes.float32, dtypes.float64, dtypes.half):
|
|
return result & _isnan(l).logical_not() & _isnan(r).logical_not()
|
|
return result
|
|
|
|
_break_var_ids = itertools.count() # unique names for per-loop break-tracking variables
|
|
|
|
def _match_bracket(toks: list[Token], start: int) -> tuple[int, list[Token]]:
|
|
"""Match brackets from start, return (end_idx, inner_tokens)."""
|
|
j, depth = start + 1, 1
|
|
while j < len(toks) and depth > 0:
|
|
if toks[j].type == 'LBRACKET': depth += 1
|
|
elif toks[j].type == 'RBRACKET': depth -= 1
|
|
j += 1
|
|
return j, [t for t in toks[start+1:j-1] if t.type != 'EOF']
|
|
|
|
def _tok_str(toks: list[Token]) -> str: return ' '.join(t.val for t in toks if t.type != 'EOF')
|
|
def parse_tokens(toks: list[Token], env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
|
|
return Parser(toks, env, funcs).parse()
|
|
|
|
# Unified block parser for pcode
|
|
def _subst_loop_var(line: str, loop_var: str, val: int) -> str:
|
|
"""Substitute loop variable with its value."""
|
|
toks = tokenize(line)
|
|
return ' '.join(str(val) if t.type == 'IDENT' and t.val == loop_var else t.val for t in toks if t.type != 'EOF')
|
|
|
|
def _set_bits(old: UOp, val: UOp, width: int, offset: int) -> UOp:
|
|
"""Set bits [offset:offset+width) in old to val, masking and shifting appropriately."""
|
|
if old.dtype in (dtypes.half, dtypes.float32): old = _val_to_bits(old)
|
|
is64 = old.dtype in (dtypes.uint64, dtypes.int64) or offset + width > 32
|
|
if is64:
|
|
old = old.cast(dtypes.uint64) if old.dtype != dtypes.uint64 else old
|
|
mask = _u64(((1 << width) - 1) << offset)
|
|
v = (val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val) & _u64((1 << width) - 1)
|
|
return (old & (mask ^ _u64(0xFFFFFFFFFFFFFFFF))) | (v << _u64(offset))
|
|
mask = _u32(((1 << width) - 1) << offset)
|
|
v = (val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val) & _u32((1 << width) - 1)
|
|
return (old & (mask ^ _u32(0xFFFFFFFF))) | (v << _u32(offset))
|
|
|
|
def _find_paren_end(s: str, start: int = 0, open_ch: str = '(', close_ch: str = ')') -> int:
|
|
"""Find index of matching close paren, starting after the open paren at start."""
|
|
depth = 0
|
|
for j, ch in enumerate(s[start:], start):
|
|
if ch == open_ch: depth += 1
|
|
elif ch == close_ch:
|
|
depth -= 1
|
|
if depth == 0: return j
|
|
return len(s)
|
|
|
|
def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dict | None = None,
|
|
assigns: list | None = None) -> tuple[int, dict[str, VarVal], UOp | None]:
|
|
"""Parse a block of pcode. Returns (next_line, block_assigns, return_value).
|
|
If assigns list is provided, side effects (MEM/VGPR writes) are appended to it."""
|
|
if funcs is None: funcs = _FUNCS
|
|
block_assigns: dict[str, VarVal] = {}
|
|
i = start
|
|
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
toks = tokenize(line)
|
|
if toks[0].type != 'IDENT' and toks[0].type != 'LBRACE':
|
|
i += 1
|
|
continue
|
|
first = toks[0].val.lower() if toks[0].type == 'IDENT' else '{'
|
|
|
|
# Block terminators
|
|
if first in ('elsif', 'else', 'endif', 'endfor'): break
|
|
|
|
# return expr (lambda bodies)
|
|
if first == 'return':
|
|
rest = line[line.lower().find('return') + 6:].strip()
|
|
return i + 1, block_assigns, parse_expr(rest, env, funcs)
|
|
|
|
# for loop
|
|
if first == 'for':
|
|
# C-style loops use an exclusive bound; for/in loops use an inclusive bound.
|
|
if m := re.fullmatch(r'for\s*\(\s*(\w+)\s*=\s*(\d+);\s*\1\s*<\s*(\d+);\s*\1\s*(\+\+|\+=\s*\d+)\s*\)', line):
|
|
loop_var, start_val, end_val = m[1], int(m[2]), int(m[3]) - 1
|
|
step = 1 if m[4] == '++' else int(m[4][2:])
|
|
else:
|
|
p = Parser(toks, env, funcs)
|
|
p.eat_val('for', 'IDENT')
|
|
loop_var = p.eat('IDENT').val
|
|
p.eat_val('in', 'IDENT')
|
|
def parse_bound():
|
|
if p.at('NUM') and p.peek(1).type == 'QUOTE':
|
|
p.eat('NUM')
|
|
p.eat('QUOTE')
|
|
if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
|
|
return int(p.parse())
|
|
start_val = parse_bound()
|
|
p.eat('COLON')
|
|
end_val, step = parse_bound(), 1
|
|
# Collect body
|
|
i += 1
|
|
body_lines: list[str] = []
|
|
depth = 1
|
|
while i < len(lines) and depth > 0:
|
|
btoks = tokenize(lines[i])
|
|
if btoks[0].type == 'IDENT':
|
|
if btoks[0].val.lower() == 'for': depth += 1
|
|
elif btoks[0].val.lower() == 'endfor': depth -= 1
|
|
if depth > 0: body_lines.append(lines[i])
|
|
i += 1
|
|
# Execute loop with break support
|
|
has_break = any('break' in bl.lower() for bl in body_lines)
|
|
found_var = f'_found_{next(_break_var_ids)}' if has_break else None
|
|
if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
|
|
for loop_i in range(start_val, end_val + 1, step):
|
|
subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
|
|
_, iter_assigns, _ = parse_block(subst_lines, 0, {**env, **block_assigns}, funcs, assigns)
|
|
if has_break:
|
|
assert found_var is not None
|
|
found = block_assigns.get(found_var, env.get(found_var))
|
|
assert isinstance(found, UOp)
|
|
not_found = found.eq(_const(dtypes.bool, False))
|
|
for var, val in iter_assigns.items():
|
|
if var != found_var and isinstance(val, UOp):
|
|
old = block_assigns.get(var, env.get(var, _u32(0)))
|
|
if isinstance(old, UOp):
|
|
block_assigns[var] = env[var] = not_found.where(
|
|
val, old.cast(val.dtype) if val.dtype != old.dtype and val.dtype.itemsize == old.dtype.itemsize else old)
|
|
for j, bl in enumerate(body_lines):
|
|
bl_l = bl.strip().lower()
|
|
if bl_l.startswith('if ') and bl_l.endswith(' then'):
|
|
if any(body_lines[k].strip().lower() == 'break' for k in range(j+1, len(body_lines))):
|
|
cond_str = _subst_loop_var(bl.strip()[3:-5].strip(), loop_var, loop_i)
|
|
cond = _to_bool(parse_expr(cond_str, env, funcs))
|
|
block_assigns[found_var] = env[found_var] = not_found.where(cond, found)
|
|
break
|
|
else:
|
|
block_assigns.update(iter_assigns)
|
|
env.update(iter_assigns)
|
|
continue
|
|
|
|
# declare
|
|
if first == 'declare':
|
|
# Initialize scalar declarations (skip arrays and env already passed as srcs)
|
|
if '[' not in line and len(toks) >= 2 and toks[1].type == 'IDENT':
|
|
env.setdefault(toks[1].val, _u32(0))
|
|
i += 1
|
|
continue
|
|
|
|
# lambda definition
|
|
if first != '{' and '=' in line and 'lambda' in line and any(t.type == 'IDENT' and t.val == 'lambda' for t in toks):
|
|
name = toks[0].val
|
|
body_start = line[line.find('(', line.find('lambda')):]
|
|
params_end = _find_paren_end(body_start) + 1
|
|
params = [p.strip() for p in body_start[1:params_end-1].split(',') if p.strip()]
|
|
rest = body_start[params_end:].strip()
|
|
if rest.startswith('('):
|
|
body_end = _find_paren_end(rest)
|
|
if body_end < len(rest): # found matching paren on same line
|
|
body = rest[1:body_end].strip()
|
|
i += 1
|
|
else: # multiline body
|
|
body_lines_lst, depth = [rest[1:]], 1
|
|
i += 1
|
|
while i < len(lines) and depth > 0:
|
|
for j, ch in enumerate(lines[i]):
|
|
if ch == '(': depth += 1
|
|
elif ch == ')':
|
|
depth -= 1
|
|
if depth == 0:
|
|
body_lines_lst.append(lines[i][:j])
|
|
break
|
|
else: body_lines_lst.append(lines[i])
|
|
i += 1
|
|
body = '\n'.join(body_lines_lst).strip()
|
|
env[name] = ('lambda', params, body)
|
|
continue
|
|
|
|
# MEM assignment: MEM[addr].type (+|-)?= value
|
|
if first == 'mem' and toks[1].type == 'LBRACKET':
|
|
j, addr_toks = _match_bracket(toks, 1)
|
|
addr = parse_tokens(addr_toks, env, funcs)
|
|
if j < len(toks) and toks[j].type == 'DOT': j += 1
|
|
dt_name = toks[j].val if j < len(toks) and toks[j].type == 'IDENT' else 'u32'
|
|
dt, j = DTYPES.get(dt_name, dtypes.uint32), j + 1
|
|
compound_op = None
|
|
if j < len(toks) and toks[j].type == 'ASSIGN_OP':
|
|
compound_op = toks[j].val
|
|
j += 1
|
|
elif j < len(toks) and toks[j].type == 'EQUALS': j += 1
|
|
rhs = parse_tokens(toks[j:], env, funcs)
|
|
if compound_op:
|
|
old = Parser([Token('EOF', '')], env, funcs)._handle_mem_load(addr, dt)
|
|
if compound_op == '+=': rhs = old + rhs
|
|
elif compound_op == '-=': rhs = old - rhs
|
|
elif compound_op == '|=': rhs = old | rhs
|
|
elif compound_op == '&=': rhs = old & rhs
|
|
elif compound_op == '^=': rhs = old ^ rhs
|
|
if assigns is not None: assigns.append((f'MEM[{_tok_str(addr_toks)}].{dt_name}', (addr, rhs)))
|
|
i += 1
|
|
continue
|
|
|
|
# VGPR assignment: VGPR[lane][reg] = value or VGPR[lane][reg][hi:lo].type = { ... }
|
|
if first == 'vgpr' and toks[1].type == 'LBRACKET':
|
|
j, lane_toks = _match_bracket(toks, 1)
|
|
if j < len(toks) and toks[j].type == 'LBRACKET':
|
|
j, reg_toks = _match_bracket(toks, j)
|
|
# Check for bit-slice: VGPR[lane][reg][hi:lo].type = value (read-modify-write)
|
|
if j < len(toks) and toks[j].type == 'LBRACKET':
|
|
j, slice_toks = _match_bracket(toks, j)
|
|
slice_str = _tok_str(slice_toks)
|
|
hi_str, lo_str = slice_str.split(':')
|
|
hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
|
|
if j < len(toks) and toks[j].type == 'DOT': j += 2 # skip .type suffix
|
|
if j < len(toks) and toks[j].type == 'EQUALS': j += 1
|
|
ln = parse_tokens(lane_toks, env, funcs)
|
|
rg, val = parse_tokens(reg_toks, env, funcs), parse_tokens(toks[j:], env, funcs)
|
|
ws = env.get('_wave_size', 32)
|
|
vgpr_idx = _to_u32(rg) * _u32(ws) + _to_u32(ln)
|
|
if assigns is not None:
|
|
assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}][{hi_val}:{lo_val}]', (vgpr_idx, val, _u32(hi_val), _u32(lo_val))))
|
|
i += 1
|
|
continue
|
|
if j < len(toks) and toks[j].type == 'DOT': j += 2 # skip .type suffix
|
|
if j < len(toks) and toks[j].type == 'EQUALS': j += 1
|
|
ln = parse_tokens(lane_toks, env, funcs)
|
|
rg, val = parse_tokens(reg_toks, env, funcs), parse_tokens(toks[j:], env, funcs)
|
|
if assigns is not None:
|
|
ws = env.get('_wave_size', 32)
|
|
assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(ws) + _to_u32(ln), val)))
|
|
i += 1
|
|
continue
|
|
|
|
# Compound destination: {hi.type, lo.type} = value
|
|
if first == '{':
|
|
j = 1
|
|
if j+2 < len(toks) and toks[j].type == 'IDENT' and toks[j+1].type == 'DOT':
|
|
hi_var, hi_type = toks[j].val, toks[j+2].val
|
|
j += 3
|
|
if j < len(toks) and toks[j].type == 'COMMA': j += 1
|
|
if j+2 < len(toks) and toks[j].type == 'IDENT' and toks[j+1].type == 'DOT':
|
|
lo_var, lo_type = toks[j].val, toks[j+2].val
|
|
j += 3
|
|
if j < len(toks) and toks[j].type == 'RBRACE': j += 1
|
|
if j < len(toks) and toks[j].type == 'EQUALS': j += 1
|
|
val = parse_tokens(toks[j:], env, funcs)
|
|
lo_dt, hi_dt = DTYPES.get(lo_type, dtypes.uint64), DTYPES.get(hi_type, dtypes.uint32)
|
|
lo_bits = 64 if lo_dt in (dtypes.uint64, dtypes.int64) else 32
|
|
lo_u = val.cast(lo_dt) if val.dtype.itemsize * 8 <= lo_bits else (val & _const(val.dtype, (1 << lo_bits) - 1)).cast(lo_dt)
|
|
hi_u = (val >> _const(val.dtype, lo_bits)).cast(hi_dt)
|
|
block_assigns[lo_var] = env[lo_var] = lo_u
|
|
block_assigns[hi_var] = env[hi_var] = hi_u
|
|
if assigns is not None: assigns.extend([(f'{lo_var}.{lo_type}', lo_u), (f'{hi_var}.{hi_type}', hi_u)])
|
|
i += 1
|
|
continue
|
|
|
|
# Bit slice/index: var[hi:lo] = value, var.type[hi:lo] = value, or var[expr] = value
|
|
if len(toks) >= 5 and toks[0].type == 'IDENT' and \
|
|
(toks[1].type == 'LBRACKET' or (toks[1].type == 'DOT' and toks[3].type == 'LBRACKET')):
|
|
bracket_start = 2 if toks[1].type == 'LBRACKET' else 4
|
|
j = bracket_start
|
|
colon_pos = None
|
|
while j < len(toks) and toks[j].type != 'RBRACKET':
|
|
if toks[j].type == 'COLON': colon_pos = j
|
|
j += 1
|
|
var = toks[0].val
|
|
if colon_pos is not None: # bit slice: var[hi:lo]
|
|
hi_str = ' '.join(t.val for t in toks[bracket_start:colon_pos] if t.type != 'EOF')
|
|
lo_str = ' '.join(t.val for t in toks[colon_pos+1:j] if t.type != 'EOF')
|
|
try:
|
|
hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
|
|
hi, lo = max(hi_val, lo_val), min(hi_val, lo_val)
|
|
j += 1
|
|
if j < len(toks) and toks[j].type == 'DOT': j += 2
|
|
if j < len(toks) and toks[j].type == 'EQUALS': j += 1
|
|
val = parse_tokens(toks[j:], env, funcs)
|
|
dt_suffix = toks[2].val if toks[1].type == 'DOT' else None
|
|
if assigns is not None: assigns.append((f'{var}[{hi}:{lo}]' + (f'.{dt_suffix}' if dt_suffix else ''), val))
|
|
if var not in env: env[var] = _const(dtypes.uint64 if hi >= 32 else dtypes.uint32, 0)
|
|
old = block_assigns.get(var, env.get(var))
|
|
assert isinstance(old, UOp)
|
|
block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
|
|
i += 1
|
|
continue
|
|
except (ValueError, SyntaxError): pass # non-constant slice bounds - fall through to other statement forms
|
|
elif toks[1].type == 'LBRACKET': # bit index: var[expr] (only for var[...], not var.type[...])
|
|
existing = block_assigns.get(var, env.get(var))
|
|
if existing is not None and isinstance(existing, UOp) and \
|
|
not any(f'{var}{k}' in env or f'{var}{k}' in block_assigns for k in range(8)):
|
|
bit_toks = toks[2:j]
|
|
j += 1
|
|
while j < len(toks) and toks[j].type != 'EQUALS': j += 1
|
|
if j < len(toks):
|
|
block_assigns[var] = env[var] = _set_bit(
|
|
existing, _to_u32(parse_tokens(bit_toks, env, funcs)), parse_tokens(toks[j+1:], env, funcs))
|
|
i += 1
|
|
continue
|
|
|
|
# Array element: var[idx] = value (static index) or var[expr] = value (dynamic)
|
|
if len(toks) >= 4 and toks[0].type == 'IDENT' and toks[1].type == 'LBRACKET':
|
|
var = toks[0].val
|
|
j, idx_toks = _match_bracket(toks, 1)
|
|
if j < len(toks) and toks[j].type == 'EQUALS':
|
|
idx_expr = parse_tokens(idx_toks, env, funcs)
|
|
# Static index: var[NUM] = value
|
|
if isinstance(idx := _single_value(idx_expr), int):
|
|
val = parse_tokens(toks[j+1:], env, funcs)
|
|
existing = block_assigns.get(var, env.get(var))
|
|
if existing is not None and isinstance(existing, UOp):
|
|
block_assigns[var] = env[var] = _set_bit(existing, _u32(idx), val)
|
|
else:
|
|
block_assigns[f'{var}@{idx}'] = env[f'{var}@{idx}'] = val
|
|
i += 1
|
|
continue
|
|
# Dynamic index: var[expr] = value where var has @-elements
|
|
elems = [(k.split('@')[1], v) for k, v in {**env, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
|
|
if elems:
|
|
val = parse_tokens(toks[j+1:], env, funcs)
|
|
for elem_idx_str, old_elem in elems:
|
|
elem_idx = int(elem_idx_str)
|
|
cond = _to_u32(idx_expr).eq(_u32(elem_idx))
|
|
new_val = cond.where(val.cast(old_elem.dtype) if val.dtype != old_elem.dtype else val, old_elem)
|
|
block_assigns[f'{var}@{elem_idx}'] = env[f'{var}@{elem_idx}'] = new_val
|
|
i += 1
|
|
continue
|
|
|
|
# Compound assignment: var += or var -=
|
|
assign_op = next((j for j, t in enumerate(toks) if t.type == 'ASSIGN_OP'), None)
|
|
if assign_op is not None:
|
|
var = toks[0].val
|
|
old = block_assigns.get(var, env.get(var, _u32(0)))
|
|
rhs = parse_tokens(toks[assign_op+1:], env, funcs)
|
|
if rhs.dtype != old.dtype: rhs = rhs.cast(old.dtype)
|
|
if toks[assign_op].val == '+=': block_assigns[var] = env[var] = old + rhs
|
|
elif toks[assign_op].val == '-=': block_assigns[var] = env[var] = old - rhs
|
|
elif toks[assign_op].val == '|=': block_assigns[var] = env[var] = old | rhs
|
|
elif toks[assign_op].val == '&=': block_assigns[var] = env[var] = old & rhs
|
|
elif toks[assign_op].val == '^=': block_assigns[var] = env[var] = old ^ rhs
|
|
i += 1
|
|
continue
|
|
|
|
# Typed element: var.type[idx] = value
|
|
if len(toks) >= 7 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and \
|
|
toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET' and toks[4].type == 'NUM':
|
|
var, dt_name, idx = toks[0].val, toks[2].val, int(toks[4].val)
|
|
dt = DTYPES.get(dt_name, dtypes.uint32)
|
|
j = 6
|
|
while j < len(toks) and toks[j].type != 'EQUALS': j += 1
|
|
if j < len(toks):
|
|
val, old = parse_tokens(toks[j+1:], env, funcs), block_assigns.get(var, env.get(var, _u32(0)))
|
|
bw = dt.itemsize * 8
|
|
block_assigns[var] = env[var] = _set_bits(old, val, bw, idx * bw)
|
|
if assigns is not None: assigns.append((f'{var}.{dt_name}[{idx}]', val))
|
|
i += 1
|
|
continue
|
|
|
|
# Dynamic bit: var.type[expr_with_brackets] = value
|
|
if len(toks) >= 5 and toks[0].type == 'IDENT' and toks[1].type == 'DOT' and \
|
|
toks[2].type == 'IDENT' and toks[3].type == 'LBRACKET':
|
|
j, depth, has_inner = 4, 1, False
|
|
while j < len(toks) and depth > 0:
|
|
if toks[j].type == 'LBRACKET':
|
|
depth += 1
|
|
has_inner = True
|
|
elif toks[j].type == 'RBRACKET': depth -= 1
|
|
j += 1
|
|
if has_inner:
|
|
var = toks[0].val
|
|
bit_pos = _to_u32(parse_tokens(toks[4:j-1], env, funcs))
|
|
while j < len(toks) and toks[j].type != 'EQUALS': j += 1
|
|
if j < len(toks):
|
|
val = parse_tokens(toks[j+1:], env, funcs)
|
|
old = block_assigns.get(var, env.get(var, _u32(0)))
|
|
block_assigns[var] = env[var] = _set_bit(old, bit_pos, val)
|
|
i += 1
|
|
continue
|
|
|
|
# If/elsif/else - skip branches with statically false conditions (WAVE32/WAVE64)
|
|
if first == 'if':
|
|
def parse_cond(s, kw):
|
|
ll = s.lower()
|
|
return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs))
|
|
def is_const(c, v): return c.op == Ops.CONST and c.val is v
|
|
cond = parse_cond(line, 'if')
|
|
conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
|
|
branch_assigns: list[tuple[UOp, list]] = [] # (cond, assigns_list) for side-effect merging
|
|
else_branch: tuple[UOp | None, dict[str, VarVal]] = (None, {})
|
|
else_side_effects: list = []
|
|
env_snap = dict(env)
|
|
static_true = is_const(cond, True) # track if any condition is statically true
|
|
i += 1
|
|
if_side: list = [] if assigns is not None and not is_const(cond, False) else []
|
|
i, branch, ret = parse_block(lines, i, env, funcs, if_side if assigns is not None and not is_const(cond, False) else None)
|
|
if conditions: conditions[0] = (cond, ret if ret is not None else branch)
|
|
if assigns is not None and not is_const(cond, False): branch_assigns.append((cond, if_side))
|
|
env.clear()
|
|
env.update(env_snap)
|
|
while i < len(lines):
|
|
ltoks = tokenize(lines[i])
|
|
if ltoks[0].type != 'IDENT': break
|
|
lf = ltoks[0].val.lower()
|
|
if lf == 'elsif':
|
|
c = parse_cond(lines[i], 'elsif')
|
|
take = not static_true and not is_const(c, False)
|
|
i += 1
|
|
br_side: list = [] if assigns is not None and take else []
|
|
i, branch, ret = parse_block(lines, i, env, funcs, br_side if assigns is not None and take else None)
|
|
if take:
|
|
conditions.append((c, ret if ret is not None else branch))
|
|
if is_const(c, True): static_true = True
|
|
if assigns is not None: branch_assigns.append((c, br_side))
|
|
env.clear()
|
|
env.update(env_snap)
|
|
elif lf == 'else':
|
|
i += 1
|
|
el_side: list = [] if assigns is not None and not static_true else []
|
|
i, branch, ret = parse_block(lines, i, env, funcs, el_side if assigns is not None and not static_true else None)
|
|
if not static_true:
|
|
else_branch = (ret, branch)
|
|
if assigns is not None: else_side_effects = el_side
|
|
env.clear()
|
|
env.update(env_snap)
|
|
elif lf == 'endif':
|
|
i += 1
|
|
break
|
|
else: break
|
|
# Check if any branch returned a value (lambda-style)
|
|
if any(isinstance(br, UOp) for _, br in conditions):
|
|
result = else_branch[0]
|
|
for c, rv in reversed(conditions):
|
|
if isinstance(rv, UOp) and isinstance(result, UOp):
|
|
if rv.dtype != result.dtype: result = result.cast(rv.dtype)
|
|
result = c.where(rv, result)
|
|
return i, block_assigns, result
|
|
# If statically true, use that branch directly; otherwise merge with WHERE
|
|
if static_true:
|
|
ba = next((b for c, b in conditions if is_const(c, True) and isinstance(b, dict)), {})
|
|
block_assigns.update(ba)
|
|
env.update(ba)
|
|
# For static true, forward side effects unconditionally
|
|
if assigns is not None:
|
|
for bc, bse in branch_assigns:
|
|
if is_const(bc, True): assigns.extend(bse)
|
|
else:
|
|
else_assigns = else_branch[1]
|
|
all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
|
|
for var in all_vars:
|
|
res: Any = else_assigns.get(var, block_assigns.get(var, env.get(var, _u32(0))))
|
|
for cond, ba in reversed(conditions): # type: ignore[assignment]
|
|
if isinstance(ba, dict) and var in ba:
|
|
tv = ba[var]
|
|
if isinstance(tv, UOp) and isinstance(res, UOp):
|
|
res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype else res)
|
|
block_assigns[var] = env[var] = res
|
|
# Merge side effects from branches with conditions
|
|
if assigns is not None:
|
|
def _cond_side_effect(cnd, dest, val):
|
|
if isinstance(val, tuple) and len(val) == 4: # VGPR bit-slice: (idx, rhs, hi, lo) -> add condition
|
|
return (dest, (val[0], val[1], val[2], val[3], cnd))
|
|
if isinstance(val, tuple) and len(val) == 2: # VGPR/MEM write: (addr, rhs) -> condition rhs
|
|
return (dest, (val[0], cnd.where(val[1], val[1])))
|
|
return (dest, val)
|
|
# Build combined condition: each branch fires when its cond is true AND no earlier cond was true
|
|
remaining = UOp.const(True)
|
|
for bc, bse in branch_assigns:
|
|
effective = remaining & bc if remaining.op != Ops.CONST else bc
|
|
for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val))
|
|
remaining = remaining & bc.logical_not() if remaining.op != Ops.CONST else bc.logical_not()
|
|
for dest, val in else_side_effects: assigns.append(_cond_side_effect(remaining, dest, val))
|
|
continue
|
|
|
|
# Regular assignment: var = value
|
|
for j, t in enumerate(toks):
|
|
if t.type == 'EQUALS':
|
|
if any(toks[k].type == 'OP' and toks[k].val in ('<', '>', '!', '=') for k in range(j)): break
|
|
base_var = toks[0].val
|
|
block_assigns[base_var] = env[base_var] = parse_tokens(toks[j+1:], env, funcs)
|
|
i += 1
|
|
break
|
|
else: i += 1
|
|
return i, block_assigns, None
|
|
|
|
def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
|
|
return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
|
|
|
|
def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None, funcs: dict | None = None) -> tuple[dict, list]:
|
|
env: dict = srcs.copy() if srcs else {}
|
|
assigns: list[tuple[str, UOp]] = []
|
|
lines: list[str] = []
|
|
blocks: list[str] = []
|
|
for raw in pcode.splitlines():
|
|
line = raw.split('//')[0].strip().rstrip(';')
|
|
if not line: continue
|
|
# Both block syntaxes share the same parser; braces supply the implicit end markers.
|
|
if line.startswith('}') and blocks:
|
|
end = blocks.pop()
|
|
line = line[1:].strip()
|
|
if not line.startswith(('elsif', 'else')): lines.append(end)
|
|
if m := re.match(r'(if|elsif|else|for)\b.*\{$', line):
|
|
blocks.append('endfor' if m[1] == 'for' else 'endif')
|
|
line = line[:-1].rstrip()
|
|
if m[1] in ('if', 'elsif'): line += ' then'
|
|
if not line: continue
|
|
line = re.sub(r'=\s*(\w+):(\w+)$', r'= {\1, \2}', line)
|
|
if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] += ' ' + line
|
|
else: lines.append(line)
|
|
assert not blocks, "unclosed pcode block"
|
|
_, final, _ = parse_block(lines, 0, env, {**_FUNCS, **funcs} if funcs else None, assigns=assigns)
|
|
sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
|
|
for var, val in final.items():
|
|
if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):
|
|
if var in sliced and not any(re.match(rf'{var}\.\w+\s*=', l) for l in lines): continue
|
|
for l in lines:
|
|
if (m := re.match(rf'{var}\.(\w+(?:\[\w+\])?)', l)):
|
|
assigns.append((f'{var}.{m.group(1)}', val))
|
|
break
|
|
else: assigns.append((var, val))
|
|
return env, assigns
|
|
|