feat: add d4r inference engines and Windows parity tooling
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
This commit is contained in:
ayo committed
1 parent
25493b3198
commit
7c5d3c9157
430 files changed
+76191
No files matched your search
@@ -0,0 +1,883 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Minimal vectorised PTX interpreter for DLSS rrlite kernels.
|
||||
|
||||
Runs one thread block of a fully unrolled kernel on the CPU (all threads of the
|
||||
block as numpy lanes), against memory restored from a bridge replay dump
|
||||
(D4R_CUDA_REPLAY_DUMP_DIR). It is a reference and an instrument: every
|
||||
statement can be observed (hooks), which is how the logical computation
|
||||
(GEMMs, weight layouts, normalisations) of a kernel is recovered.
|
||||
|
||||
Only the instruction forms the rrlite kernels use are implemented; anything
|
||||
else raises NotImplementedError naming the statement.
|
||||
"""
|
||||
import re
|
||||
import struct
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
F16, F32, F64 = np.float16, np.float32, np.float64
|
||||
U64 = np.uint64
|
||||
PARAM_BASE = 1 << 60 # virtual address of the kernel parameter block
|
||||
MASK = {8: 0xff, 16: 0xffff, 32: 0xffffffff, 64: 0xffffffffffffffff}
|
||||
|
||||
|
||||
def _e4m3_decode_table():
|
||||
t = np.zeros(256, F64)
|
||||
for c in range(256):
|
||||
s = -1.0 if c & 0x80 else 1.0
|
||||
e, m = (c >> 3) & 0xf, c & 7
|
||||
if (c & 0x7f) == 0x7f:
|
||||
t[c] = np.nan
|
||||
elif e == 0:
|
||||
t[c] = s * m * 2.0 ** -9
|
||||
else:
|
||||
t[c] = s * (1 + m / 8) * 2.0 ** (e - 7)
|
||||
return t
|
||||
|
||||
|
||||
E4M3 = _e4m3_decode_table()
|
||||
|
||||
|
||||
def _f16_to_e4m3_table():
|
||||
"""Exact RNE + satfinite f16 -> e4m3 for all 65536 f16 codes (ZLUDA's helper)."""
|
||||
h = np.arange(65536, dtype=np.uint32)
|
||||
mag = h & 0x7fff
|
||||
normal = (mag - 0x2000 + 0x3f + ((mag >> 7) & 1)) >> 7
|
||||
exp = mag >> 10
|
||||
sig = (mag & 0x3ff) | np.where(exp != 0, 0x400, 0)
|
||||
cl = np.clip(exp, 1, 8)
|
||||
sh = 16 - cl
|
||||
sub = (sig + (np.left_shift(1, sh - 1)) - 1 + ((sig >> sh) & 1)) >> sh
|
||||
code = np.where(mag >= 0x2400, normal, sub)
|
||||
code = np.where(mag > 0x5f00, 0x7e, code)
|
||||
code = np.where(mag > 0x7c00, 0x7f, code)
|
||||
return (((h >> 8) & 0x80) | code).astype(np.uint32)
|
||||
|
||||
|
||||
F16_TO_E4M3 = _f16_to_e4m3_table()
|
||||
|
||||
|
||||
def bits_to_f16(x):
|
||||
return (x & 0xffff).astype(np.uint16).view(F16)
|
||||
|
||||
|
||||
def f16_to_bits(v):
|
||||
return np.asarray(v, F16).view(np.uint16).astype(U64)
|
||||
|
||||
|
||||
def bits_to_f32(x):
|
||||
return (x & 0xffffffff).astype(np.uint32).view(F32)
|
||||
|
||||
|
||||
def f32_to_bits(v):
|
||||
return np.asarray(v, F32).view(np.uint32).astype(U64)
|
||||
|
||||
|
||||
class Stmt:
|
||||
__slots__ = ('labels', 'pred', 'neg', 'op', 'args', 'text', 'index')
|
||||
|
||||
|
||||
def split_args(s):
|
||||
out, depth, cur = [], 0, ''
|
||||
for ch in s:
|
||||
if ch in '{[':
|
||||
depth += 1
|
||||
elif ch in '}]':
|
||||
depth -= 1
|
||||
if ch == ',' and depth == 0:
|
||||
out.append(cur.strip())
|
||||
cur = ''
|
||||
else:
|
||||
cur += ch
|
||||
if cur.strip():
|
||||
out.append(cur.strip())
|
||||
return out
|
||||
|
||||
|
||||
def parse_kernel(path, name):
|
||||
with open(path) as file:
|
||||
src = file.read()
|
||||
i = src.index('.entry ' + name + '(')
|
||||
j = src.find('\n.visible', i + 10)
|
||||
j = len(src) if j < 0 else j
|
||||
body = src[i:j]
|
||||
body = body[body.index(')') + 1:]
|
||||
# module-scope .shared arrays (declared before the entry) are visible to it too
|
||||
module_shared = {m.group(1): int(m.group(2)) for m in re.finditer(r'^\.shared[^\n;]*?(\w+)\[(\d+)\]', src[:i], re.M)}
|
||||
body = re.sub(r'//[^\n]*', '', body)
|
||||
body = re.sub(r'\.(maxntid|reqntid|minnctapersm|maxnreg)[^\n{]*', '', body)
|
||||
stmts, types, shared = [], {}, {}
|
||||
pending_labels = []
|
||||
for raw in re.sub(r'\s*\n\s*', ' ', body).split(';'):
|
||||
s = raw.strip()
|
||||
while True:
|
||||
s = s.strip()
|
||||
if s.startswith('{') and not s.startswith('{%'):
|
||||
s = s[1:]
|
||||
continue
|
||||
if s.startswith('}'):
|
||||
s = s[1:]
|
||||
continue
|
||||
m = re.match(r'(\$[\w]+):\s*', s)
|
||||
if m:
|
||||
pending_labels.append(m.group(1))
|
||||
s = s[m.end():]
|
||||
continue
|
||||
break
|
||||
if not s:
|
||||
continue
|
||||
if s.startswith('.reg'):
|
||||
m = re.match(r'\.reg\s*\.(\w+)\s+(.*)', s)
|
||||
ty = m.group(1)
|
||||
for decl in m.group(2).split(','):
|
||||
decl = decl.strip()
|
||||
m2 = re.match(r'(%[\w]+)<(\d+)>', decl)
|
||||
if m2:
|
||||
types[('prefix', m2.group(1))] = ty
|
||||
else:
|
||||
types[decl] = ty
|
||||
continue
|
||||
if s.startswith('.shared'):
|
||||
m = re.search(r'(\w+)\[(\d+)\]', s)
|
||||
shared[m.group(1)] = int(m.group(2))
|
||||
continue
|
||||
if s.startswith('.'):
|
||||
continue
|
||||
st = Stmt()
|
||||
st.labels, pending_labels = pending_labels, []
|
||||
st.pred, st.neg = None, False
|
||||
m = re.match(r'@(!?)(%\w+)\s+', s)
|
||||
if m:
|
||||
st.neg, st.pred = m.group(1) == '!', m.group(2)
|
||||
s = s[m.end():]
|
||||
parts = s.split(None, 1)
|
||||
st.op = parts[0]
|
||||
st.args = split_args(parts[1]) if len(parts) > 1 else []
|
||||
st.text = s
|
||||
st.index = len(stmts)
|
||||
stmts.append(st)
|
||||
if pending_labels:
|
||||
st = Stmt()
|
||||
st.labels, st.pred, st.neg, st.op, st.args, st.text, st.index = pending_labels, None, False, 'nop', [], 'nop', len(stmts)
|
||||
stmts.append(st)
|
||||
for k, v in module_shared.items():
|
||||
shared.setdefault(k, v)
|
||||
return stmts, types, shared
|
||||
|
||||
|
||||
def reg_width(types, name):
|
||||
ty = types.get(name)
|
||||
if ty is None:
|
||||
m = re.match(r'(%[a-z]+)', name)
|
||||
ty = types.get(('prefix', m.group(1)))
|
||||
return {'pred': 1, 'b16': 16, 'u16': 16, 's16': 16, 'f16': 16, 'b32': 32, 'u32': 32, 's32': 32, 'f32': 32,
|
||||
'b64': 64, 'u64': 64, 's64': 64, 'f64': 64, 'b8': 8}[ty]
|
||||
|
||||
|
||||
class Memory:
|
||||
"""Global memory: the replay dump's allocations at their captured addresses."""
|
||||
|
||||
def __init__(self, dump_dir):
|
||||
manifest = open(dump_dir + '/manifest.txt').read().split('\n')
|
||||
self.args = open(dump_dir + '/args.bin', 'rb').read()
|
||||
self.bases, self.data, self.ids = [], [], []
|
||||
for line in manifest:
|
||||
p = line.split()
|
||||
if p and p[0] == 'alloc':
|
||||
idx, base, size = int(p[1]), int(p[2], 16), int(p[3])
|
||||
buf = np.fromfile(f'{dump_dir}/alloc-{idx}.bin', np.uint8)
|
||||
assert buf.size == size
|
||||
self.bases.append(base)
|
||||
self.data.append(buf)
|
||||
self.ids.append(idx)
|
||||
order = np.argsort(self.bases)
|
||||
self.bases = [self.bases[i] for i in order]
|
||||
self.data = [self.data[i] for i in order]
|
||||
self.ids = [self.ids[i] for i in order]
|
||||
self.base_arr = np.array(self.bases, U64)
|
||||
self.stores = [] # (addresses, bytes) of every global store
|
||||
|
||||
def locate(self, addr):
|
||||
idx = np.searchsorted(self.base_arr, addr, side='right') - 1
|
||||
if np.any(idx < 0):
|
||||
raise ValueError('address below every allocation')
|
||||
return idx
|
||||
|
||||
def load(self, addr, nbytes, mask):
|
||||
out = np.zeros((addr.size, nbytes), np.uint8)
|
||||
idx = self.locate(np.where(mask, addr, self.base_arr[0]))
|
||||
for a in np.unique(idx[mask]):
|
||||
sel = mask & (idx == a)
|
||||
off = (addr[sel] - self.bases[a]).astype(np.int64)
|
||||
if np.any(off + nbytes > self.data[a].size):
|
||||
raise ValueError('load past the end of an allocation')
|
||||
out[sel] = self.data[a][off[:, None] + np.arange(nbytes)]
|
||||
return out
|
||||
|
||||
def store(self, addr, raw, mask):
|
||||
idx = self.locate(np.where(mask, addr, self.base_arr[0]))
|
||||
n = raw.shape[1]
|
||||
for a in np.unique(idx[mask]):
|
||||
sel = mask & (idx == a)
|
||||
off = (addr[sel] - self.bases[a]).astype(np.int64)
|
||||
self.data[a][off[:, None] + np.arange(n)] = raw[sel]
|
||||
self.stores.append((self.ids[a], off.copy(), raw[sel].copy()))
|
||||
|
||||
|
||||
class Block:
|
||||
def __init__(self, stmts, types, shared, mem, param_name, block_dim, grid_dim, ctaid):
|
||||
self.stmts, self.types, self.mem = stmts, types, mem
|
||||
self.param_name = param_name
|
||||
bx, by, bz = block_dim
|
||||
self.T = bx * by * bz
|
||||
t = np.arange(self.T)
|
||||
self.tid = {'x': t % bx, 'y': (t // bx) % by, 'z': t // (bx * by)}
|
||||
self.lane = t % 32
|
||||
self.warp = t // 32
|
||||
self.ctaid = ctaid
|
||||
self.grid = grid_dim
|
||||
self.ntid = block_dim
|
||||
self.regs = {}
|
||||
self.shared = {k: np.zeros(v, np.uint8) for k, v in shared.items()}
|
||||
self.shared_base = {}
|
||||
base = 0
|
||||
for k, v in shared.items():
|
||||
self.shared_base[k] = base
|
||||
base += (v + 15) & ~15
|
||||
self.shared_mem = np.zeros(max(base, 16), np.uint8)
|
||||
self.labels = {}
|
||||
for st in stmts:
|
||||
for l in st.labels:
|
||||
self.labels[l] = st.index
|
||||
self.hooks = []
|
||||
self.retired = np.zeros(self.T, bool)
|
||||
self.prov = {}
|
||||
|
||||
# --- operands -----------------------------------------------------------
|
||||
def width(self, name):
|
||||
return reg_width(self.types, name)
|
||||
|
||||
def special(self, name):
|
||||
m = re.match(r'%(tid|ntid|ctaid|nctaid)\.([xyz])', name)
|
||||
if m:
|
||||
kind, ax = m.groups()
|
||||
k = 'xyz'.index(ax)
|
||||
if kind == 'tid':
|
||||
return self.tid[ax].astype(U64)
|
||||
if kind == 'ntid':
|
||||
return np.full(self.T, self.ntid[k], U64)
|
||||
if kind == 'ctaid':
|
||||
return np.full(self.T, self.ctaid[k], U64)
|
||||
return np.full(self.T, self.grid[k], U64)
|
||||
if name == 'WARP_SZ':
|
||||
return np.full(self.T, 32, U64)
|
||||
if name == '%laneid':
|
||||
return self.lane.astype(U64)
|
||||
if name == '%warpid':
|
||||
return self.warp.astype(U64)
|
||||
return None
|
||||
|
||||
def imm(self, s, width=64):
|
||||
if s.startswith('0f'):
|
||||
return np.full(self.T, int(s[2:], 16), U64)
|
||||
if s.startswith('0d'):
|
||||
return np.full(self.T, int(s[2:], 16), U64)
|
||||
v = int(s, 0)
|
||||
return np.full(self.T, v & MASK[64], U64)
|
||||
|
||||
def get(self, s, width=None):
|
||||
s = s.strip()
|
||||
if s.startswith('%') or s in self.types or s == 'WARP_SZ':
|
||||
sp = self.special(s) if (s.startswith('%') or s == 'WARP_SZ') else None
|
||||
if sp is not None:
|
||||
return sp
|
||||
if s not in self.regs:
|
||||
w = self.width(s)
|
||||
self.regs[s] = np.zeros(self.T, bool if w == 1 else U64)
|
||||
return self.regs[s]
|
||||
if s in self.shared_base:
|
||||
return np.full(self.T, self.shared_base[s], U64)
|
||||
if s == self.param_name:
|
||||
return np.full(self.T, PARAM_BASE, U64)
|
||||
if s[0].isdigit() or s[0] == '-':
|
||||
v = self.imm(s)
|
||||
if width:
|
||||
v = v & np.uint64(MASK[width])
|
||||
return v
|
||||
raise NotImplementedError('operand ' + s)
|
||||
|
||||
def put(self, name, value, mask, keep_prov=False):
|
||||
name = name.strip()
|
||||
if not keep_prov:
|
||||
self.prov.pop(name, None)
|
||||
w = self.width(name)
|
||||
old = self.get(name)
|
||||
if w == 1:
|
||||
self.regs[name] = np.where(mask, value.astype(bool), old)
|
||||
else:
|
||||
self.regs[name] = np.where(mask, value.astype(U64) & np.uint64(MASK[w]), old)
|
||||
|
||||
def vec(self, s):
|
||||
return [x.strip() for x in s.strip()[1:-1].split(',')]
|
||||
|
||||
def address(self, s):
|
||||
"""[reg+off] or [symbol+off] -> (space, address array)."""
|
||||
inner = s.strip()[1:-1]
|
||||
m = re.match(r'([%\w$.]+)\s*(\+\s*-?\d+)?', inner)
|
||||
base, off = m.group(1), int(m.group(2).replace('+', '')) if m.group(2) else 0
|
||||
if base == self.param_name:
|
||||
return 'param', np.full(self.T, off, U64)
|
||||
if base in self.shared_base:
|
||||
return 'shared', np.full(self.T, self.shared_base[base] + off, U64)
|
||||
return 'reg', (self.get(base) + np.uint64(off & MASK[64])) & np.uint64(MASK[64])
|
||||
|
||||
# --- execution ------------------------------------------------------------
|
||||
def run(self, limit=None):
|
||||
stmts = self.stmts
|
||||
active = np.ones(self.T, bool)
|
||||
pending = {}
|
||||
pc = 0
|
||||
steps = 0
|
||||
while pc < len(stmts):
|
||||
st = stmts[pc]
|
||||
if pc in pending:
|
||||
active |= pending.pop(pc)
|
||||
if not active.any():
|
||||
# Reconverge at either a forward target or a loop fallthrough.
|
||||
targets = sorted(pending)
|
||||
if not targets:
|
||||
break
|
||||
pc = targets[0]
|
||||
continue
|
||||
active &= ~self.retired
|
||||
mask = active
|
||||
if st.pred is not None:
|
||||
p = self.get(st.pred)
|
||||
mask = active & (~p if st.neg else p)
|
||||
if st.op == 'bra' or st.op.startswith('bra.'):
|
||||
target = self.labels[st.args[0]]
|
||||
taking = mask
|
||||
if target <= pc:
|
||||
if (taking == active).all():
|
||||
pc = target
|
||||
continue
|
||||
if not taking.any():
|
||||
pc += 1
|
||||
continue
|
||||
# Lanes finishing a loop wait at its fallthrough while the
|
||||
# other lanes execute their remaining iterations. Their
|
||||
# registers and memory writes stay masked until reconvergence.
|
||||
pending[pc + 1] = pending.get(pc + 1, np.zeros(self.T, bool)) | (active & ~taking)
|
||||
active = taking.copy()
|
||||
pc = target
|
||||
continue
|
||||
active = active & ~taking
|
||||
if taking.any():
|
||||
pending[target] = pending.get(target, np.zeros(self.T, bool)) | taking
|
||||
pc += 1
|
||||
continue
|
||||
if mask.any() or st.op.startswith(('mma', 'shfl', 'movmatrix', 'bar')):
|
||||
self.execute(st, mask)
|
||||
for hook in self.hooks:
|
||||
hook(self, st, mask)
|
||||
pc += 1
|
||||
steps += 1
|
||||
if limit and steps >= limit:
|
||||
break
|
||||
return steps
|
||||
|
||||
def execute(self, st, mask):
|
||||
op, a = st.op, st.args
|
||||
parts = op.split('.')
|
||||
base = parts[0]
|
||||
fn = getattr(self, 'op_' + base, None)
|
||||
if fn is None:
|
||||
raise NotImplementedError(st.text)
|
||||
try:
|
||||
fn(st, parts, a, mask)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f'{type(e).__name__} {e} in: {st.text}') from e
|
||||
|
||||
# --- helpers --------------------------------------------------------------
|
||||
@staticmethod
|
||||
def sext(v, w):
|
||||
v = v.astype(np.int64) & ((1 << w) - 1) if w < 64 else v.view(np.int64)
|
||||
if w < 64:
|
||||
v = np.where(v >= (1 << (w - 1)), v - (1 << w), v)
|
||||
return v
|
||||
|
||||
def typed(self, s, ty):
|
||||
w = int(re.sub(r'\D', '', ty) or 32) if ty not in ('pred',) else 1
|
||||
v = self.get(s, w)
|
||||
if ty.startswith('s'):
|
||||
return self.sext(v, w)
|
||||
if ty.startswith('f'):
|
||||
if ty == 'f16':
|
||||
return bits_to_f16(v).astype(F64)
|
||||
if ty == 'f32':
|
||||
return bits_to_f32(v).astype(F64)
|
||||
return v.view(F64)
|
||||
return (v & np.uint64(MASK[w])).astype(np.int64) if w < 64 else v
|
||||
|
||||
def result(self, ty, value):
|
||||
if ty == 'f16':
|
||||
return f16_to_bits(np.asarray(value).astype(F16))
|
||||
if ty == 'f32':
|
||||
return f32_to_bits(np.asarray(value).astype(F32))
|
||||
if ty == 'f64':
|
||||
return np.asarray(value, F64).view(U64)
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
return (np.asarray(value).astype(np.int64).astype(U64)) & np.uint64(MASK[w])
|
||||
|
||||
# --- instructions -----------------------------------------------------------
|
||||
def op_nop(self, st, parts, a, mask):
|
||||
pass
|
||||
|
||||
def op_bar(self, st, parts, a, mask):
|
||||
pass
|
||||
|
||||
def op_ret(self, st, parts, a, mask):
|
||||
self.retired |= mask
|
||||
|
||||
def op_neg(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty == 'f16x2':
|
||||
self.put(a[0], self.get(a[1]) ^ np.uint64(0x80008000), mask)
|
||||
elif ty == 'f16':
|
||||
self.put(a[0], self.get(a[1]) ^ np.uint64(0x8000), mask)
|
||||
elif ty == 'f32':
|
||||
self.put(a[0], self.get(a[1]) ^ np.uint64(0x80000000), mask)
|
||||
else:
|
||||
self.put(a[0], self.result('u' + ty[1:], -self.typed(a[1], ty)), mask)
|
||||
|
||||
def op_cvta(self, st, parts, a, mask):
|
||||
self.put(a[0], self.get(a[1]), mask)
|
||||
|
||||
def op_mov(self, st, parts, a, mask):
|
||||
d, s = a
|
||||
if d.startswith('{'):
|
||||
names = self.vec(d)
|
||||
v = self.get(s)
|
||||
w = self.width(names[0])
|
||||
for i, n in enumerate(names):
|
||||
self.put(n, (v >> np.uint64(i * w)) & np.uint64(MASK[w]), mask)
|
||||
return
|
||||
if s.startswith('{'):
|
||||
names = self.vec(s)
|
||||
w = self.width(names[0])
|
||||
v = np.zeros(self.T, U64)
|
||||
for i, n in enumerate(names):
|
||||
v |= (self.get(n) & np.uint64(MASK[w])) << np.uint64(i * w)
|
||||
self.put(d, v, mask)
|
||||
return
|
||||
self.put(d, self.get(s), mask)
|
||||
if s in self.prov:
|
||||
self.prov[d] = self.prov[s]
|
||||
|
||||
def op_ld(self, st, parts, a, mask):
|
||||
space = 'global' if 'global' in parts else 'shared' if 'shared' in parts else 'param' if 'param' in parts else None
|
||||
ty = parts[-1]
|
||||
vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
nbytes = w // 8
|
||||
kind, addr = self.address(a[1])
|
||||
total = nbytes * vec
|
||||
if kind == 'param':
|
||||
raw = np.frombuffer(self.mem.args, np.uint8)
|
||||
off = int(addr[0])
|
||||
data = np.tile(raw[off:off + total], (self.T, 1))
|
||||
elif space == 'param':
|
||||
# ld.param through a computed address (mov.b64 %rd, <param>; add; ld.param [%rd])
|
||||
raw = np.frombuffer(self.mem.args, np.uint8)
|
||||
off = np.where(mask, addr - np.uint64(PARAM_BASE), 0).astype(np.int64)
|
||||
data = raw[off[:, None] + np.arange(total)]
|
||||
elif kind == 'shared' or space == 'shared':
|
||||
off = addr.astype(np.int64)
|
||||
data = self.shared_mem[off[:, None] + np.arange(total)]
|
||||
else:
|
||||
data = self.mem.load(addr, total, mask)
|
||||
names = self.vec(a[0]) if a[0].startswith('{') else [a[0]]
|
||||
for i, n in enumerate(names):
|
||||
chunk = data[:, i * nbytes:(i + 1) * nbytes]
|
||||
v = np.zeros(self.T, U64)
|
||||
for b in range(nbytes):
|
||||
v |= chunk[:, b].astype(U64) << np.uint64(8 * b)
|
||||
if ty.startswith('s') and w < self.width(n):
|
||||
v = (self.sext(v, w).astype(np.int64)).astype(U64)
|
||||
self.put(n, v, mask)
|
||||
self.prov[n] = (st.index, kind if kind != 'reg' else space, addr + np.uint64(i * nbytes), nbytes)
|
||||
|
||||
def op_st(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
nbytes = w // 8
|
||||
kind, addr = self.address(a[0])
|
||||
names = self.vec(a[1]) if a[1].startswith('{') else [a[1]]
|
||||
raw = np.zeros((self.T, nbytes * len(names)), np.uint8)
|
||||
for i, n in enumerate(names):
|
||||
v = self.get(n, w)
|
||||
for b in range(nbytes):
|
||||
raw[:, i * nbytes + b] = ((v >> np.uint64(8 * b)) & np.uint64(0xff)).astype(np.uint8)
|
||||
if kind == 'shared' or 'shared' in parts:
|
||||
off = addr.astype(np.int64)[mask]
|
||||
self.shared_mem[off[:, None] + np.arange(raw.shape[1])] = raw[mask]
|
||||
else:
|
||||
self.mem.store(addr, raw, mask)
|
||||
|
||||
def int_binop(self, parts, a, mask, fn):
|
||||
ty = parts[-1]
|
||||
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
|
||||
self.put(a[0], self.result(ty if not ty.startswith('b') else 'u' + ty[1:], fn(x, y)), mask)
|
||||
|
||||
def op_add(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) + self.typed(a[2], ty)), mask)
|
||||
elif ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x + y), mask)
|
||||
else:
|
||||
self.int_binop(parts, a, mask, lambda x, y: x + y)
|
||||
|
||||
def op_sub(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) - self.typed(a[2], ty)), mask)
|
||||
elif ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x - y), mask)
|
||||
else:
|
||||
self.int_binop(parts, a, mask, lambda x, y: x - y)
|
||||
|
||||
def op_mul(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if 'wide' in parts:
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
|
||||
self.put(a[0], (x * y).astype(np.int64).astype(U64), mask)
|
||||
elif ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty)), mask)
|
||||
elif ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x * y), mask)
|
||||
else:
|
||||
self.int_binop(parts, a, mask, lambda x, y: x * y)
|
||||
|
||||
def op_mad(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
x, y, z = self.typed(a[1], ty), self.typed(a[2], ty), self.typed(a[3], ty)
|
||||
self.put(a[0], self.result('u' + ty[1:], x * y + z), mask)
|
||||
|
||||
def h2(self, sa, sb, fn, sc=None):
|
||||
va, vb = self.get(sa), self.get(sb)
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
x = bits_to_f16(va >> np.uint64(16 * i)).astype(F64)
|
||||
y = bits_to_f16(vb >> np.uint64(16 * i)).astype(F64)
|
||||
if sc is not None:
|
||||
z = bits_to_f16(self.get(sc) >> np.uint64(16 * i)).astype(F64)
|
||||
r = fn(x, y, z)
|
||||
else:
|
||||
r = fn(x, y)
|
||||
out |= f16_to_bits(np.asarray(r).astype(F16)) << np.uint64(16 * i)
|
||||
return out
|
||||
|
||||
def op_fma(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y, z: x * y + z, a[3]), mask)
|
||||
elif ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty) + self.typed(a[3], ty)), mask)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
|
||||
def minmax(self, parts, a, mask, is_max):
|
||||
ty = parts[-1]
|
||||
if ty == 'f16x2':
|
||||
f = np.fmax if is_max else np.fmin
|
||||
self.put(a[0], self.h2(a[1], a[2], f), mask)
|
||||
elif ty in ('f16', 'f32'):
|
||||
f = np.fmax if is_max else np.fmin
|
||||
self.put(a[0], self.result(ty, f(self.typed(a[1], ty), self.typed(a[2], ty))), mask)
|
||||
else:
|
||||
f = np.maximum if is_max else np.minimum
|
||||
self.int_binop(parts, a, mask, f)
|
||||
|
||||
def op_min(self, st, parts, a, mask):
|
||||
self.minmax(parts, a, mask, False)
|
||||
|
||||
def op_max(self, st, parts, a, mask):
|
||||
self.minmax(parts, a, mask, True)
|
||||
|
||||
def op_abs(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty in ('f16x2', 'f16', 'f32'):
|
||||
m = {'f16x2': 0x7fff7fff, 'f16': 0x7fff, 'f32': 0x7fffffff}[ty]
|
||||
self.put(a[0], self.get(a[1]) & np.uint64(m), mask)
|
||||
return
|
||||
self.put(a[0], self.result('u' + ty[1:], np.abs(self.typed(a[1], ty))), mask)
|
||||
|
||||
def logic(self, a, mask, fn):
|
||||
if self.width(a[0]) == 1:
|
||||
self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
|
||||
else:
|
||||
self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
|
||||
|
||||
def op_and(self, st, parts, a, mask):
|
||||
self.logic(a, mask, lambda x, y: x & y)
|
||||
|
||||
def op_or(self, st, parts, a, mask):
|
||||
self.logic(a, mask, lambda x, y: x | y)
|
||||
|
||||
def op_xor(self, st, parts, a, mask):
|
||||
self.logic(a, mask, lambda x, y: x ^ y)
|
||||
|
||||
def op_not(self, st, parts, a, mask):
|
||||
self.put(a[0], ~self.get(a[1]), mask)
|
||||
|
||||
def op_shl(self, st, parts, a, mask):
|
||||
w = int(re.sub(r'\D', '', parts[-1]))
|
||||
n = self.get(a[2]) & np.uint64(0xff)
|
||||
v = np.where(n >= w, 0, (self.get(a[1]) << np.minimum(n, np.uint64(63))))
|
||||
self.put(a[0], v.astype(U64), mask)
|
||||
|
||||
def op_shr(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
n = (self.get(a[2]) & np.uint64(0xff)).astype(np.int64)
|
||||
if ty.startswith('s'):
|
||||
v = self.sext(self.get(a[1]), w)
|
||||
r = v >> np.minimum(n, w - 1)
|
||||
self.put(a[0], r.astype(np.int64).astype(U64), mask)
|
||||
else:
|
||||
v = self.get(a[1]) & np.uint64(MASK[w])
|
||||
r = np.where(n >= w, 0, v >> np.minimum(n, 63).astype(U64))
|
||||
self.put(a[0], r.astype(U64), mask)
|
||||
|
||||
def op_bfi(self, st, parts, a, mask):
|
||||
f, b = self.get(a[1]), self.get(a[2])
|
||||
pos = self.get(a[3]).astype(np.int64) & 0xff
|
||||
ln = self.get(a[4]).astype(np.int64) & 0xff
|
||||
w = int(re.sub(r'\D', '', parts[-1]))
|
||||
out = np.array(b, U64)
|
||||
for i in range(self.T):
|
||||
p, l = int(pos[i]), int(ln[i])
|
||||
if l == 0 or p >= w:
|
||||
continue
|
||||
l = min(l, w - p)
|
||||
m = ((1 << l) - 1) << p
|
||||
out[i] = (int(b[i]) & ~m & MASK[w]) | ((int(f[i]) << p) & m)
|
||||
self.put(a[0], out, mask)
|
||||
|
||||
def op_setp(self, st, parts, a, mask):
|
||||
cmp, ty = parts[1], parts[-1]
|
||||
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
|
||||
r = {'lt': x < y, 'le': x <= y, 'gt': x > y, 'ge': x >= y, 'eq': x == y, 'ne': x != y,
|
||||
'lo': x < y, 'ls': x <= y, 'hi': x > y, 'hs': x >= y}[cmp]
|
||||
dests = a[0].split('|')
|
||||
self.put(dests[0], r, mask)
|
||||
|
||||
def op_selp(self, st, parts, a, mask):
|
||||
p = self.get(a[3])
|
||||
self.put(a[0], np.where(p, self.get(a[1]), self.get(a[2])), mask)
|
||||
|
||||
def op_ex2(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
v = self.get(a[1])
|
||||
if ty == 'f16x2':
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
h = bits_to_f16((v >> np.uint64(16 * i)) & np.uint64(0xffff)).astype(F64)
|
||||
out |= f16_to_bits(np.exp2(h).astype(F16)) << np.uint64(16 * i)
|
||||
self.put(a[0], out, mask)
|
||||
elif ty == 'f16':
|
||||
self.put(a[0], f16_to_bits(np.exp2(bits_to_f16(v).astype(F64)).astype(F16)), mask)
|
||||
elif ty == 'f32':
|
||||
self.put(a[0], f32_to_bits(np.exp2(bits_to_f32(v).astype(F64)).astype(F32)), mask)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
|
||||
def op_rsqrt(self, st, parts, a, mask):
|
||||
x = self.typed(a[1], 'f32').astype(F32)
|
||||
with np.errstate(divide='ignore', invalid='ignore'):
|
||||
self.put(a[0], f32_to_bits((1.0 / np.sqrt(x.astype(F64))).astype(F32)), mask)
|
||||
|
||||
def op_rcp(self, st, parts, a, mask):
|
||||
x = self.typed(a[1], 'f32')
|
||||
with np.errstate(divide='ignore'):
|
||||
self.put(a[0], f32_to_bits((1.0 / x).astype(F32)), mask)
|
||||
|
||||
def op_div(self, st, parts, a, mask):
|
||||
x, y = self.typed(a[1], 'f32'), self.typed(a[2], 'f32')
|
||||
with np.errstate(divide='ignore', invalid='ignore'):
|
||||
self.put(a[0], f32_to_bits((x / y).astype(F32)), mask)
|
||||
|
||||
def op_cvt(self, st, parts, a, mask):
|
||||
types = [p for p in parts[1:] if p not in ('rn', 'rz', 'rm', 'rp', 'rni', 'rzi', 'sat', 'satfinite', 'relu', 'ftz')]
|
||||
dt, stp = types[0], types[1]
|
||||
relu = 'relu' in parts
|
||||
v = self.get(a[1])
|
||||
if dt == 'e4m3x2' and stp == 'f16x2':
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
h = (v >> np.uint64(16 * i)) & np.uint64(0xffff)
|
||||
if relu:
|
||||
f = bits_to_f16(h)
|
||||
h = np.where((f < 0) & ~np.isnan(f.astype(F32)), 0, h).astype(U64)
|
||||
out |= F16_TO_E4M3[h.astype(np.int64)].astype(U64) << np.uint64(8 * i)
|
||||
self.put(a[0], out, mask)
|
||||
elif dt == 'f16x2' and stp == 'e4m3x2':
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
c = ((v >> np.uint64(8 * i)) & np.uint64(0xff)).astype(np.int64)
|
||||
out |= f16_to_bits(E4M3[c].astype(F16)) << np.uint64(16 * i)
|
||||
self.put(a[0], out, mask)
|
||||
elif dt == 'f16' and stp == 'f32':
|
||||
self.put(a[0], f16_to_bits(bits_to_f32(v).astype(F16)), mask)
|
||||
elif dt == 'f16' and stp == 'f64':
|
||||
self.put(a[0], f16_to_bits(v.view(F64).astype(F16)), mask)
|
||||
elif dt == 'f32' and stp == 'f16':
|
||||
self.put(a[0], f32_to_bits(bits_to_f16(v).astype(F32)), mask)
|
||||
elif dt in ('f16', 'f32') and stp[0] in 'us':
|
||||
ws = int(re.sub(r'\D', '', stp))
|
||||
iv = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
|
||||
if dt == 'f16':
|
||||
self.put(a[0], f16_to_bits(iv.astype(F64).astype(F16)), mask)
|
||||
else:
|
||||
self.put(a[0], f32_to_bits(iv.astype(F64).astype(F32)), mask)
|
||||
elif dt[0] in 'us' and stp[0] in 'us':
|
||||
ws = int(re.sub(r'\D', '', stp))
|
||||
val = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
|
||||
self.put(a[0], val.astype(np.int64).astype(U64), mask)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
|
||||
def op_shfl(self, st, parts, a, mask):
|
||||
mode = parts[2]
|
||||
v = self.get(a[1])
|
||||
b = self.get(a[2]).astype(np.int64)
|
||||
out = np.array(v)
|
||||
for w in range(self.T // 32):
|
||||
sl = slice(32 * w, 32 * w + 32)
|
||||
lane = np.arange(32)
|
||||
if mode == 'bfly':
|
||||
src = lane ^ b[sl]
|
||||
elif mode == 'down':
|
||||
src = np.where(lane + b[sl] < 32, lane + b[sl], lane)
|
||||
elif mode == 'up':
|
||||
src = np.where(lane - b[sl] >= 0, lane - b[sl], lane)
|
||||
elif mode == 'idx':
|
||||
src = b[sl] & 31
|
||||
out[sl] = v[sl][src]
|
||||
self.put(a[0], out, mask)
|
||||
|
||||
def op_movmatrix(self, st, parts, a, mask):
|
||||
v = self.get(a[1])
|
||||
out = np.zeros(self.T, U64)
|
||||
for w in range(self.T // 32):
|
||||
m = np.zeros((8, 8), U64)
|
||||
for lane in range(32):
|
||||
x = int(v[32 * w + lane])
|
||||
r, c = lane // 4, (lane % 4) * 2
|
||||
m[r, c], m[r, c + 1] = x & 0xffff, (x >> 16) & 0xffff
|
||||
mt = m.T
|
||||
for lane in range(32):
|
||||
r, c = lane // 4, (lane % 4) * 2
|
||||
out[32 * w + lane] = int(mt[r, c]) | (int(mt[r, c + 1]) << 16)
|
||||
self.put(a[0], out, mask)
|
||||
|
||||
def op_mma(self, st, parts, a, mask):
|
||||
shape = parts[3]
|
||||
d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
|
||||
A = [self.get(n) for n in a_names]
|
||||
Bv = [self.get(n) for n in b_names]
|
||||
C = [self.get(n) for n in c_names]
|
||||
dout = [np.array(self.get(n)) for n in d_names]
|
||||
for w in range(self.T // 32):
|
||||
base = 32 * w
|
||||
if shape == 'm16n8k32':
|
||||
Am = np.zeros((16, 32)); Bm = np.zeros((32, 8)); Cm = np.zeros((16, 8))
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for i in range(16):
|
||||
byte = (int(A[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
|
||||
row = g if (i < 4 or 8 <= i < 12) else g + 8
|
||||
col = t * 4 + (i & 3) + (16 if i >= 8 else 0)
|
||||
Am[row, col] = E4M3[byte]
|
||||
for i in range(8):
|
||||
byte = (int(Bv[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
|
||||
Bm[t * 4 + (i & 3) + (16 if i >= 4 else 0), g] = E4M3[byte]
|
||||
elif shape == 'm16n8k16':
|
||||
Am = np.zeros((16, 16)); Bm = np.zeros((16, 8)); Cm = np.zeros((16, 8))
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for i in range(8):
|
||||
h = (int(A[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
|
||||
row = g if (i < 2 or 4 <= i < 6) else g + 8
|
||||
col = t * 2 + (i & 1) + (8 if i >= 4 else 0)
|
||||
Am[row, col] = np.uint16(h).view(F16)
|
||||
for i in range(4):
|
||||
h = (int(Bv[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
|
||||
Bm[t * 2 + (i & 1) + (8 if i >= 2 else 0), g] = np.uint16(h).view(F16)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for i in range(4):
|
||||
h = (int(C[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
|
||||
Cm[g if i < 2 else g + 8, t * 2 + (i & 1)] = np.uint16(h).view(F16)
|
||||
Dm = (Am @ Bm + Cm).astype(F16)
|
||||
if self.mma_hook:
|
||||
self.mma_hook(self, st, w, Am, Bm, Cm, Dm, [self.prov.get(n) for n in a_names], [self.prov.get(n) for n in b_names])
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for r in range(2):
|
||||
lo = int(Dm[g + 8 * r, t * 2].view(np.uint16))
|
||||
hi = int(Dm[g + 8 * r, t * 2 + 1].view(np.uint16))
|
||||
dout[r][base + lane] = lo | (hi << 16)
|
||||
for n, v in zip(d_names, dout):
|
||||
self.put(n, v, mask)
|
||||
|
||||
mma_hook = None
|
||||
|
||||
# wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 with register fragments: two m16n8k16 MMAs
|
||||
# sharing A registers 0..3, B registers 0,1 / 2,3 and C/D registers 0,1 / 2,3 (the upper A and B
|
||||
# registers of the fragment ABI are duplicates; ZLUDA's LLVMZludaSplitWmmaM16n16k16 mapping)
|
||||
def op_wmma(self, st, parts, a, mask):
|
||||
if 'm16n16k16' not in parts or 'mma' not in parts:
|
||||
raise NotImplementedError(st.text)
|
||||
d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
|
||||
for half in range(2):
|
||||
sub = Stmt_like(st, 'mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16',
|
||||
['{%s}' % ', '.join(d_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(a_names[0:4]),
|
||||
'{%s}' % ', '.join(b_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(c_names[2 * half:2 * half + 2])])
|
||||
self.op_mma(sub, sub.op.split('.'), sub.args, mask)
|
||||
|
||||
|
||||
class Stmt_like:
|
||||
def __init__(self, st, op, args):
|
||||
self.op, self.args, self.text = op, args, st.text + f' [as {op}]'
|
||||
for k in Stmt.__slots__:
|
||||
if k not in ('op', 'args', 'text'):
|
||||
setattr(self, k, getattr(st, k, None))
|
||||
|
||||
|
||||
def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None):
|
||||
stmts, types, shared = parse_kernel(ptx, kernel)
|
||||
mem = Memory(dump)
|
||||
launch = [l for l in open(dump + '/manifest.txt') if l.startswith('launch')][0].split()
|
||||
grid = tuple(int(x) for x in launch[1:4])
|
||||
block = tuple(int(x) for x in launch[4:7])
|
||||
b = Block(stmts, types, shared, mem, kernel + '_param_0', block, grid, ctaid)
|
||||
b.hooks = list(hooks)
|
||||
b.mma_hook = mma_hook
|
||||
steps = b.run(limit)
|
||||
return b, mem, steps
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
ptx, kernel, dump = sys.argv[1:4]
|
||||
cx, cy = (int(v) for v in sys.argv[4].split(',')) if len(sys.argv) > 4 else (0, 0)
|
||||
import time
|
||||
t0 = time.time()
|
||||
b, mem, steps = run_block(ptx, kernel, dump, (cx, cy, 0))
|
||||
print(f'block ({cx},{cy}): {steps} statements in {time.time() - t0:.1f}s, {len(mem.stores)} store groups')
|
||||
Reference in new issue
Block a user