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.
146 lines
8.7 KiB
Python
146 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the engine's own preset L stages on a captured launch and compare them with its outputs.
|
|
|
|
usage: test_l.py enc0|dec0 CAPTURE_DIR FRAME WORK_DIR KVK [repeats]
|
|
enc0: layer_m.comp KIND 3 (features, embedding, block, merge) against rrlite_enc0_4x4_mv*_*: the block output
|
|
(dec0's skip input), the merged tokens (the network's input) and the three motion surfaces
|
|
dec0: layer_m.comp KIND 4 (expand, block, head) against rrlite_dec0_4x4: its four surfaces
|
|
CAPTURE_DIR as for test_enc0_m.py (replay/ before each launch, dump/ after it, filter `rrlite_`). KVK is
|
|
vulkan/kvk built with its `bindings` manifest line. The half-resolution previous result (texture 152) is not in the
|
|
dumps: enc0 uses WORK_DIR/half-prev.bin if it exists (zeros otherwise), which only matters after the first frame.
|
|
"""
|
|
import glob, os, re, struct, subprocess, sys
|
|
import numpy as np
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, HERE)
|
|
import l_model as lm
|
|
import model_enc3 as me
|
|
from test_m import pack_layer, defines, token_codes, report
|
|
|
|
|
|
def launch(cap, kernel_re, frame):
|
|
"""(sequence number, replay dir, dumped resources by parameter offset) of the frame's launch of the kernel"""
|
|
names = sorted({re.match(r'launch-(\d+)-(.*)-arg', os.path.basename(f)).groups() for f in glob.glob(f'{cap}/dump/launch-*-arg*')})
|
|
hits = [(int(s), k) for s, k in names if re.fullmatch(kernel_re, k)]
|
|
if len(hits) <= frame:
|
|
raise SystemExit(f'no launch of {kernel_re} for frame {frame}')
|
|
seq, k = hits[frame]
|
|
res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{k}-arg*')}
|
|
return seq, k, f'{cap}/replay/replay-{seq:06d}-{k}', res
|
|
|
|
|
|
class Replay:
|
|
def __init__(self, d):
|
|
self.P = open(d + '/args.bin', 'rb').read()
|
|
self.allocs = {}
|
|
for line in open(d + '/manifest.txt'):
|
|
q = line.split()
|
|
if q and q[0] == 'alloc':
|
|
self.allocs[int(q[1])] = (int(q[2], 16), np.fromfile(f'{d}/alloc-{q[1]}.bin', np.uint8))
|
|
|
|
def at(self, ptr):
|
|
"""the replayed allocation bytes from address ptr on"""
|
|
a, b = next((a, b) for a, b in self.allocs.values() if a <= ptr < a + b.size)
|
|
return b[ptr - a:]
|
|
|
|
def q(self, o): return struct.unpack_from('<Q', self.P, o)[0]
|
|
def i2(self, o): return struct.unpack_from('<2i', self.P, o)
|
|
def f2(self, o): return struct.unpack_from('<2f', self.P, o)
|
|
def f1(self, o): return struct.unpack_from('<f', self.P, o)[0]
|
|
|
|
|
|
def compile_stage(name, out, extra=()):
|
|
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(name), *extra, f'{HERE}/layer_m.comp', '-o', out],
|
|
check=True, stdout=subprocess.DEVNULL)
|
|
|
|
|
|
def run(kvk, spv, manifest, reps):
|
|
r = subprocess.run([kvk, spv, manifest, reps, '32'], capture_output=True, text=True)
|
|
if r.returncode:
|
|
print('FAILED', r.stderr[-800:]); sys.exit(1)
|
|
for l in r.stdout.split('\n'):
|
|
if 'us per' in l: print(l)
|
|
|
|
|
|
def enc0(cap, frame, work, kvk, reps):
|
|
seq, k, d, res = launch(cap, r'rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)', frame)
|
|
R = Replay(d)
|
|
wbuf = R.at(R.q(0))
|
|
t = lm.enc0_layer(wbuf)
|
|
o, b8, b16 = pack_layer(t, 'enc0l')
|
|
rw, rh = R.i2(56); gw, gh = R.i2(16); sx, sy = R.i2(8)
|
|
b8.tofile(f'{work}/w8.bin'); b16.tofile(f'{work}/w16.bin'); wbuf[:4096].tofile(f'{work}/noise.bin')
|
|
np.zeros(64, np.uint8).tofile(f'{work}/dummy.bin')
|
|
open(f'{work}/lp.bin', 'wb').write(struct.pack('<4i72I', gw, gh, sx, sy, *o))
|
|
expo = lambda off: float(np.fromfile(res[off], np.float16)[0])
|
|
hw, hh = R.i2(248)
|
|
open(f'{work}/ep.bin', 'wb').write(struct.pack('<2i2i2i2f2f2f2f2f2f2fff3Iffffff', rw, rh, sx, sy, gw // 2, gh // 2, *R.f2(168), *R.f2(176),
|
|
*R.f2(200), *R.f2(208), *R.f2(192), *R.f2(184), float(hw), float(hh), R.f1(268), R.f1(272),
|
|
R.P[216], R.P[316], R.i2(312)[0], expo(112), expo(120), *R.f2(256), *R.f2(64)))
|
|
half = f'{work}/half-prev.bin'
|
|
if not os.path.exists(half):
|
|
np.zeros((hh // 2) * (hw // 2) * 4, np.float32).tofile(half)
|
|
lines = [f'param {work}/lp.bin', f'grid {(gw + sx + 7) // 8} {(gh + sy + 7) // 8}',
|
|
f'tex 1 {res[80]} {rw} {rh} 4 16 nearest', f'tex 2 {res[96]} {rw} {rh} 2 16 nearest', f'tex 3 {res[136]} {rw} {rh} 1 32 nearest',
|
|
f'tex 4 {res[128]} {rw} {rh} 2 16 linear', f'tex 5 {half} {hw // 2} {hh // 2} 4 32 linear', f'tex 6 {res[160]} {rw} {rh} 4 9 linear']
|
|
for i, off in enumerate((280, 288, 296)): lines.append(f'surf {7 + i} {rw} {rh} 2 16 - {work}/mv-{off}.bin')
|
|
lines += [f'lut {work}/w8.bin', f'lut {work}/w16.bin', f'lut {work}/dummy.bin', f'lut {work}/dummy.bin',
|
|
f'outbuf {work}/skip.bin {gw * gh * 32}', f'outbuf {work}/merged.bin {(gw // 2) * (gh // 2) * 64}',
|
|
f'lut {work}/ep.bin', f'lut {work}/noise.bin',
|
|
'bindings 0 8 9 10 11 12 13 14 15 16 1 2 3 4 5 6 7 17']
|
|
open(f'{work}/manifest', 'w').write('\n'.join(lines) + '\n')
|
|
spv = f'{work}/enc0_l.spv'
|
|
compile_stage('enc0l', spv, ['-DLDR=1'] if k.endswith('_ldr') else [])
|
|
run(kvk, spv, f'{work}/manifest', reps)
|
|
print(f'frame {frame} ({k}, launch {seq}) {rw}x{rh}, pixel grid {gw}x{gh}, shift {sx},{sy}, reset {R.P[316]}')
|
|
for off, name in ((280, 'A'), (288, 'B'), (296, 'C')):
|
|
a = np.fromfile(res[off], np.float16).astype(np.float64); b = np.fromfile(f'{work}/mv-{off}.bin', np.float16).astype(np.float64)
|
|
print(f' motion {name}: identical {100 * (a == b).mean():5.1f}% max |diff| {np.abs(a - b).max():.4g}')
|
|
out, merged = R.q(40), R.q(48)
|
|
after = np.fromfile(res[40], np.uint8); base = out
|
|
ref = token_codes(after, 0, gw, gh, 32)
|
|
got = np.fromfile(f'{work}/skip.bin', np.uint8).reshape(gh, gw, 32)
|
|
print(' ' + report('block output', got, ref))
|
|
print(' ' + report('block output, rows inside the image', got[:rh], ref[:rh]))
|
|
mref = token_codes(after, merged - base, gw // 2, gh // 2, 64)
|
|
mgot = np.fromfile(f'{work}/merged.bin', np.uint8).reshape(gh // 2, gw // 2, 64)
|
|
print(' ' + report('merged', mgot, mref))
|
|
print(' ' + report('merged, rows inside the image', mgot[:rh // 2], mref[:rh // 2]))
|
|
|
|
|
|
def dec0(cap, frame, work, kvk, reps):
|
|
seq, k, d, res = launch(cap, r'rrlite_dec0_4x4', frame)
|
|
R = Replay(d)
|
|
t = lm.dec0_layer(R.at(R.q(0)))
|
|
o, b8, b16 = pack_layer(t, 'dec0l')
|
|
sx, sy = R.i2(8); gw, gh = R.i2(16)
|
|
rw, rh = [int(x) for x in re.search(r'surface-(\d+)x(\d+)x', res[128]).groups()]
|
|
b8.tofile(f'{work}/dw8.bin'); b16.tofile(f'{work}/dw16.bin')
|
|
np.zeros(64, np.uint8).tofile(f'{work}/dummy.bin')
|
|
token_codes(R.at(R.q(24)), 0, gw // 2, gh // 2, 64).tofile(f'{work}/din.bin')
|
|
token_codes(R.at(R.q(32)), 0, gw, gh, 32).tofile(f'{work}/dskip.bin')
|
|
open(f'{work}/dlp.bin', 'wb').write(struct.pack('<4i72I', gw, gh, sx, sy, *o))
|
|
open(f'{work}/dp.bin', 'wb').write(struct.pack('<2iiI2f2f2f2fff2f', rw, rh, gw // 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, R.f1(160), R.f1(164), 0, 0))
|
|
lines = [f'param {work}/dlp.bin', f'grid {(gw + sx + 7) // 8} {(gh + sy + 7) // 8}']
|
|
fmt = {128: (4, 9), 136: (4, 8), 144: (4, 8), 152: (4, 16)}
|
|
for off in (128, 136, 144, 152): lines.append(f'surf {off} {rw} {rh} {fmt[off][0]} {fmt[off][1]} - {work}/s{off}.bin')
|
|
lines += [f'lut {work}/dw8.bin', f'lut {work}/dw16.bin', f'lut {work}/din.bin', f'lut {work}/dskip.bin',
|
|
f'outbuf {work}/dump.bin {gw * gh * 64}', f'lut {work}/dummy.bin', f'lut {work}/dp.bin', 'bindings 0 8 9 10 11 1 2 3 4 5 6 7']
|
|
open(f'{work}/dmanifest', 'w').write('\n'.join(lines) + '\n')
|
|
spv = f'{work}/dec0_l.spv'
|
|
compile_stage('dec0l', spv, [a for a in os.environ.get('L_DEFS', '').split()])
|
|
run(kvk, spv, f'{work}/dmanifest', reps)
|
|
print(f'frame {frame} (launch {seq}) {rw}x{rh}, pixel grid {gw}x{gh}, shift {sx},{sy}, sigma {R.f1(160):.4g} rho {R.f1(164):.4g}')
|
|
for off, name, dt in ((128, 'feature', np.int8), (136, 'gates', np.uint8), (144, 'mix', np.uint8), (152, 'covariance', np.float16)):
|
|
a = np.fromfile(res[off], dt).reshape(rh, rw, 4).astype(np.float64); b = np.fromfile(f'{work}/s{off}.bin', dt).reshape(rh, rw, 4).astype(np.float64)
|
|
same = (a == b) | (np.isnan(a) & np.isnan(b))
|
|
print(f' {name:10s}: identical {100 * same.mean():6.2f}% (per channel ' + ' '.join(f'{100 * same[..., c].mean():6.2f}' for c in range(4)) +
|
|
f'), max |diff| {np.nanmax(np.abs(a - b)):.4g}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
stage, cap, frame, work, kvk = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4], sys.argv[5]
|
|
reps = sys.argv[6] if len(sys.argv) > 6 else '1'
|
|
os.makedirs(work, exist_ok=True)
|
|
{'enc0': enc0, 'dec0': dec0}[stage](cap, frame, work, kvk, reps)
|