ayo
7c5d3c9157
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.
178 lines
9.1 KiB
Python
178 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the engine's own preset M layers (layer_m.comp) on a captured frame and compare with the captured outputs and
|
|
with the plain numpy model (m_model.py).
|
|
|
|
usage: test_m.py CAPTURE_DIR WORK_DIR EVK [--frame N] [--time N] [--model] [--chain] [layer index ...]
|
|
CAPTURE_DIR/replay: layer launches before they ran; CAPTURE_DIR/dump: buffers after each launch.
|
|
--model also runs the numpy model (slow) and reports the shader against it.
|
|
--chain feeds each layer the engine's own previous outputs instead of the captured inputs.
|
|
|
|
The layers' FP8 matrices are exported in the shader's blocked fixed-16 tile layout (pack8 below); the front-end
|
|
embedding/expansion weights that compile_m.py packages for enc0_m/dec0_m stay in NGX's own plain layout.
|
|
"""
|
|
import glob, os, struct, subprocess, sys
|
|
import numpy as np
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, HERE)
|
|
import m_model as mm
|
|
import model_enc3 as me
|
|
|
|
F16 = np.float16
|
|
|
|
|
|
def codes(values):
|
|
"""e4m3 codes of exact e4m3 values"""
|
|
v = np.asarray(values, np.float64); a = np.abs(v)
|
|
c = np.searchsorted(me._POS, a).astype(np.uint8)
|
|
assert np.array_equal(me._POS[c], a), 'not an FP8 value'
|
|
return (c | np.where(np.signbit(v) & (a > 0), 128, 0)).astype(np.uint8)
|
|
|
|
|
|
class Pack:
|
|
def __init__(self, dtype, align):
|
|
self.parts, self.n, self.dtype, self.align = [], 0, dtype, align
|
|
|
|
def add(self, a):
|
|
a = np.ascontiguousarray(a, self.dtype).ravel()
|
|
pad = (-self.n) % self.align
|
|
if pad: self.parts.append(np.zeros(pad, self.dtype)); self.n += pad
|
|
off = self.n; self.parts.append(a); self.n += a.size
|
|
return off
|
|
|
|
def data(self):
|
|
return np.concatenate(self.parts + [np.zeros((-self.n) % 64 + 64, self.dtype)])
|
|
|
|
|
|
def pack8(values):
|
|
"""blocked FP8 bytes of a logical weight matrix W[K][N] (K, N multiples of 16), the layout layer_m.comp loads:
|
|
[K/16][N/16][n16][k16], i.e. reshape(K//16, 16, N//16, 16).transpose(0, 2, 3, 1). contiguous (no arithmetic).
|
|
A fixed 16x16 tile is 256 bytes and starts at 256*(kblock*(N//16) + nblock); the shader's A operand is the tile
|
|
[N16][K16] read row-major with stride 16 (A[n][k] = W[kblock*16 + k][nblock*16 + n]), so a K16 row block of
|
|
every output channel tile stays contiguous and the per-head slice of wo is still base + 32*h*C."""
|
|
W = codes(values)
|
|
K, N = W.shape
|
|
assert K % 16 == 0 and N % 16 == 0, (K, N)
|
|
return np.ascontiguousarray(W.reshape(K // 16, 16, N // 16, 16).transpose(0, 2, 3, 1)).ravel()
|
|
|
|
|
|
def pack_layer(t, name):
|
|
"""(o[72], FP8 bytes, f16 words) of one layer's tensors; the FP8 matrices are blocked by pack8()"""
|
|
C, heads, NPM, CIN = mm.LAYERS[name]
|
|
b8, b16 = Pack(np.uint8, 16), Pack(F16, 16)
|
|
tile = lambda v: np.tile(np.asarray(v, F16)[None, :], (16, 1))
|
|
o = [0] * 72
|
|
o[0], o[1], o[2], o[3] = b16.add(t['g1']), b16.add(tile(t['bo'])), b16.add(t['g2']), b16.add(tile(t['b2']))
|
|
assert t['wo'].shape == (32 * heads, C), t['wo'].shape
|
|
wo = b8.add(pack8(t['wo']))
|
|
for h in range(heads):
|
|
assert t['q'][h].shape == (C, 32) and t['v'][h].shape == (C, 32), (t['q'][h].shape, t['v'][h].shape)
|
|
o[4 + 4 * h], o[5 + 4 * h], o[6 + 4 * h], o[7 + 4 * h] = b8.add(pack8(t['q'][h])), b8.add(pack8(t['v'][h])), b16.add(t['bias'][h]), wo + 32 * h * C
|
|
for c in range(C // 8):
|
|
assert t['w1'][c].shape == (C, 32) and t['w2'][c].shape == (32, C), (t['w1'][c].shape, t['w2'][c].shape)
|
|
o[20 + 3 * c], o[21 + 3 * c], o[22 + 3 * c] = b8.add(pack8(t['w1'][c])), b16.add(tile(t['b1'][c])), b8.add(pack8(t['w2'][c]))
|
|
if NPM:
|
|
assert t['merge_w'].shape == (4 * C, NPM), t['merge_w'].shape
|
|
o[68], o[69] = b8.add(pack8(t['merge_w'])), b16.add(tile(t['merge_b']))
|
|
if CIN:
|
|
assert t['expand_w'].shape == (CIN, 4 * C), t['expand_w'].shape
|
|
o[68], o[69] = b8.add(pack8(t['expand_w'])), b16.add(tile(t['expand_b']))
|
|
if 'emb_w1' in t:
|
|
# preset L's input embedding (layer_m.comp KIND 3): 32 features -> 32 hidden -> C
|
|
assert t['emb_w1'].shape == (32, 32) and t['emb_w2'].shape == (32, C), (t['emb_w1'].shape, t['emb_w2'].shape)
|
|
o[60], o[61] = b8.add(pack8(t['emb_w1'])), b16.add(tile(t['emb_b1']))
|
|
o[62], o[63] = b8.add(pack8(t['emb_w2'])), b16.add(tile(t['emb_b2']))
|
|
return o, b8.data(), b16.data()
|
|
|
|
|
|
def defines(name):
|
|
C, heads, NPM, CIN = mm.LAYERS[name]
|
|
kind = {'enc0l': 3, 'dec0l': 4}.get(name, 1 if NPM else 2 if CIN else 0)
|
|
# compile_m.py's arithmetic (ZLUDA's lowering); M_TEXL=0 gives the native accuracy kernels' for the network layers.
|
|
# M_DEFS: extra defines for every layer_m.comp stage; M_DEFS_ENC1, M_DEFS_ENC3_TUBE, ...: for one layer
|
|
texl = ['-DTEXL=1'] if kind >= 3 or os.environ.get('M_TEXL', '1') != '0' else []
|
|
# M_WPW=2: two 4x4-token windows per wave in the network layers (layer_m.comp WPW; not L's input / expansion)
|
|
wpw = [f"-DWPW={os.environ['M_WPW']}"] if kind < 3 and os.environ.get('M_WPW') else []
|
|
return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={kind}', f'-DXC={NPM or CIN or 16}'] + texl + wpw + os.environ.get('M_DEFS', '').split() + os.environ.get('M_DEFS_' + name.upper(), '').split()
|
|
|
|
|
|
def token_codes(buf, off, W, H, C):
|
|
a = buf[off:off + (C // 32) * H * W * 32].reshape(C // 32, H, W, 32)
|
|
return np.ascontiguousarray(a.transpose(1, 2, 0, 3).reshape(H, W, C))
|
|
|
|
|
|
def report(tag, got, ref):
|
|
rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v)
|
|
g, r = me.E4M3[got].astype(np.float64), me.E4M3[ref].astype(np.float64)
|
|
steps = np.abs(rank(g) - rank(r))
|
|
return (f'{tag} {mm.psnr(g, r):5.1f} dB, identical {100 * (got == ref).mean():5.1f}%, within one step {100 * (steps <= 1).mean():5.1f}%')
|
|
|
|
|
|
def main():
|
|
a = sys.argv[1:]
|
|
def opt(flag, default):
|
|
if flag in a:
|
|
i = a.index(flag); v = a[i + 1]; del a[i:i + 2]; return v
|
|
return default
|
|
frame, reps = int(opt('--frame', 0)), opt('--time', '1')
|
|
model = '--model' in a
|
|
if model: a.remove('--model')
|
|
chain = '--chain' in a
|
|
if chain: a.remove('--chain')
|
|
own = {} # arena address -> the engine's own tokens there
|
|
cap, work, evk = a[:3]
|
|
only = [int(x) for x in a[3:]]
|
|
os.makedirs(work, exist_ok=True)
|
|
Ps, dirs = mm.launches(cap, frame)
|
|
total = 0.0
|
|
for i, (P, d, name) in enumerate(zip(Ps, dirs, mm.ORDER)):
|
|
if only and i not in only: continue
|
|
C, heads, NPM, CIN = mm.LAYERS[name]
|
|
seq = int(os.path.basename(d).split('-')[1])
|
|
after = np.fromfile(glob.glob(f'{cap}/dump/launch-{seq:03d}-*-arg024-*')[0], np.uint8)
|
|
t = mm.plain_layer(P, name)
|
|
o, b8, b16 = pack_layer(t, name)
|
|
W, H = P.W, P.H
|
|
pre = f'{work}/{i}-{name}'
|
|
open(pre + '.params', 'wb').write(struct.pack('<4i72I', W, H, P.sx, P.sy, *o))
|
|
b8.tofile(pre + '.w8'); b16.tofile(pre + '.w16')
|
|
def source(ptr, w, h_, c):
|
|
return own[ptr] if chain and ptr in own else token_codes(P.arena, ptr - P.abase, w, h_, c)
|
|
if CIN:
|
|
source(P.inp, W // 2, H // 2, CIN).tofile(pre + '.in')
|
|
source(P.p32, W, H, C).tofile(pre + '.skip')
|
|
else:
|
|
source(P.inp, W, H, C).tofile(pre + '.in')
|
|
np.zeros(64, np.uint8).tofile(pre + '.skip')
|
|
spv = pre + '.spv'
|
|
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(name), f'{HERE}/layer_m.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL)
|
|
gx, gy = (W + P.sx + 7) // 8, (H + P.sy + 7) // 8
|
|
msize = (H // 2) * (W // 2) * NPM if NPM else 64
|
|
r = subprocess.run([evk, spv, str(gx), str(gy), reps, f'in:0={pre}.params', f'in:1={pre}.w8', f'in:2={pre}.w16', f'in:3={pre}.in',
|
|
f'in:4={pre}.skip', f'out:5={pre}.out:{H * W * C}', f'out:6={pre}.merged:{msize}'], capture_output=True, text=True)
|
|
if r.returncode:
|
|
print(name, 'FAILED', r.stderr[-600:]); sys.exit(1)
|
|
us = [float(l.split(':')[1].split()[0]) for l in r.stdout.split('\n') if 'us per dispatch' in l]
|
|
total += us[0] if us else 0
|
|
got = np.fromfile(pre + '.out', np.uint8).reshape(H, W, C)
|
|
own[P.out] = got
|
|
ref = token_codes(after, P.out - P.abase, W, H, C)
|
|
line = f'{i} {name:9s} {W}x{H}x{C} {us[0] if us else 0:7.1f} us: ' + report('vs capture', got, ref)
|
|
if NPM:
|
|
gm = np.fromfile(pre + '.merged', np.uint8).reshape(H // 2, W // 2, NPM)
|
|
own[P.p48] = gm
|
|
line += '; ' + report('merged', gm, token_codes(after, P.p48 - P.abase, W // 2, H // 2, NPM))
|
|
print(line, flush=True)
|
|
if model:
|
|
if CIN:
|
|
out, merged = mm.forward(t, name, mm.tokens(P.arena, P.inp - P.abase, W // 2, H // 2, CIN), mm.tokens(P.arena, P.p32 - P.abase, W, H, C), P.sx, P.sy)
|
|
else:
|
|
out, merged = mm.forward(t, name, mm.tokens(P.arena, P.inp - P.abase, W, H, C), None, P.sx, P.sy)
|
|
line = ' ' + report('shader vs numpy model', got, codes(out))
|
|
if NPM: line += '; ' + report('merged', gm, codes(merged))
|
|
print(line, flush=True)
|
|
print(f'total {total:.1f} us')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|