woutermans
e0d1549a21
A standalone Vulkan inference engine (engine/) for DLSS presets K and M that runs from weights extracted from the user's own nvngx_dlss.dll, with d4r's own input, network and reconstruction shaders or the native HIP network layers, recorded into the game's command list by the shim. Opt-in ([Engine] Enabled / D4R_ENGINE=1); anything uncovered falls back to the CUDA/ZLUDA path with the reason logged. Byte-identical to NGX's CUDA backend under the phase-capable reference (1706x960 -> 2560x1440, native HIP network, eleven frames, RX 9070 XT): - K render- and display-resolution motion vectors, HDR and LDR - K game-supplied exposure and automatic exposure with any pre-exposure and exposure scale (one residual: effective exposure exactly 2 or 4) - K and M with the DLSS indicator's inverted-axis parameters - M with any exposure settings (NGX ignores them for M) K packages need the current compile_k.py (marker D4RO0003). Also: HIP shared buffers are pooled across feature recreation (ROCm leaks destroyed imports), hipnet.bin must match the preset, M-only games do not load the K model, install_d4r_runtime.sh stages k, k-ldr, m, m-ldr with hip/, the activation table has a generator (gen_activation_m.py), and package_release.sh can bundle the engine's shaders (never its weights). Signed-off-by: woutermans <[email protected]>
151 lines
10 KiB
Python
151 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Run whole preset M frames on the engine's own stages and compare the output with the capture.
|
|
|
|
usage: test_frame_m.py CAPTURE_DIR WORK_DIR KVK EVK [frames]
|
|
Every stage runs as its own process (kvk for the texture stages, evk for the network layers). Only the frame's
|
|
colour, motion and depth, the exposure values and the kernels' parameter blocks come from the capture; all
|
|
intermediate data and all temporal history are the engine's own.
|
|
"""
|
|
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 m_model as mm
|
|
import swin_model as sm
|
|
from test_m import pack_layer, defines, token_codes, report
|
|
cap, work, kvk, evk = sys.argv[1:5]
|
|
frames = int(sys.argv[5]) if len(sys.argv) > 5 else 2
|
|
os.makedirs(work, exist_ok=True)
|
|
LDR = os.environ.get('LDR') == '1' # a capture without the HDR flag: NGX's *_ldr_* kernels, the shaders' LDR paths
|
|
R = 'ldr' if LDR else 'hdr'
|
|
# below a reconstruction/render ratio of 3 the capture holds NGX's rrlite_post_3_2; the engine then needs the
|
|
# larger shared tile (post_m.comp's TILE_W/TILE_H)
|
|
POST = os.environ.get('POST_KERNEL') or next((k for k in (f'rrlite_post_3_1_mvlo_{R}_folded', f'rrlite_post_3_2_mvlo_{R}_folded')
|
|
if glob.glob(f'{cap}/dump/launch-*-{k}-arg*')), f'rrlite_post_3_1_mvlo_{R}_folded')
|
|
POST_TILE = ['-DTILE_W=15', '-DTILE_H=15'] if '3_2' in POST else []
|
|
NAMES = {'enc0': f'rrlite_enc0_4x4_mvlo_{R}_folded', 'dec0': 'rrlite_dec0_4x4_folded', 'post': POST, 'down': f'rrlite_downsample_kernel_static_{R}'}
|
|
STAGE = ['-DLDR=1'] if LDR else []
|
|
|
|
|
|
class Launch:
|
|
def __init__(self, kernel, frame):
|
|
seqs = sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{kernel}-arg*')})
|
|
self.seq = seqs[frame]
|
|
self.res = {int(re.search(r'-arg(\d+)-', f).group(1)): f for f in glob.glob(f'{cap}/dump/launch-{self.seq:03d}-{kernel}-arg*')}
|
|
self.dir = f'{cap}/replay/replay-{self.seq:06d}-{kernel}'
|
|
self.P = open(self.dir + '/args.bin', 'rb').read()
|
|
self.allocs = {}
|
|
for line in open(self.dir + '/manifest.txt'):
|
|
q = line.split()
|
|
if q and q[0] == 'alloc': self.allocs[int(q[1])] = (int(q[2], 16), np.fromfile(f'{self.dir}/alloc-{q[1]}.bin', np.uint8))
|
|
i2 = lambda s, o: struct.unpack_from('<2i', s.P, o)
|
|
f2 = lambda s, o: struct.unpack_from('<2f', s.P, o)
|
|
f1 = lambda s, o: struct.unpack_from('<f', s.P, o)[0]
|
|
u8 = lambda s, o: struct.unpack_from('<Q', s.P, o)[0]
|
|
expo = lambda s, o: float(np.fromfile(s.res[o], np.float16)[0]) if o in s.res else 1.0
|
|
def blob(self, ptr):
|
|
a, b = next((a, b) for a, b in self.allocs.values() if a <= ptr < a + b.size)
|
|
return b[ptr - a:]
|
|
|
|
|
|
def shader(name, flags=()):
|
|
spv = f'{work}/{name}{"".join(flags).replace("-D", "_").replace("=", "")}.spv'
|
|
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, f'{HERE}/{name}.comp', '-o', spv], check=True, stdout=subprocess.DEVNULL)
|
|
return spv
|
|
|
|
|
|
def run(cmd):
|
|
r = subprocess.run(cmd, capture_output=True, text=True)
|
|
if r.returncode:
|
|
print('FAILED', ' '.join(cmd[:3]), r.stderr[-500:]); sys.exit(1)
|
|
|
|
|
|
def zeros(path, n, dtype=np.float16):
|
|
np.zeros(n, dtype).tofile(path); return path
|
|
|
|
|
|
prev = None # the previous frame's files
|
|
for f in range(frames):
|
|
D = f'{work}/f{f}'; os.makedirs(D, exist_ok=True)
|
|
# ---------------------------------------------------------------- input stage
|
|
L = Launch(NAMES['enc0'], f)
|
|
rw, rh = L.i2(56); gw, gh = L.i2(16); tw, th = gw // 2, gh // 2
|
|
hw, hh = L.i2(248)
|
|
wb = L.blob(L.u8(0))
|
|
g = sm.gp(np.arange(64))
|
|
np.ascontiguousarray(wb[sm.woff_table(4096, 512, 2048, 128, 64)][:, g]).tofile(f'{D}/e0.w8')
|
|
np.tile(wb[12288:12288 + 128].view(np.float16)[g][None, :], (16, 1)).tofile(f'{D}/e0.w16')
|
|
wb[:4096].tofile(f'{D}/e0.noise')
|
|
open(f'{D}/e0.params', 'wb').write(struct.pack('<2i2i2i2f2f2f2f2f2f2fff3Iffffff', rw, rh, *L.i2(8), tw, th, *L.f2(168), *L.f2(176), *L.f2(200), *L.f2(208),
|
|
*L.f2(192), *L.f2(184), float(hw), float(hh), L.f1(268), L.f1(272), L.P[216], L.P[316], L.i2(312)[0], L.expo(112), L.expo(120), *L.f2(256), *L.f2(64)))
|
|
colour, motion, depth = L.res[80], L.res[96], L.res[136]
|
|
e1, e2 = L.expo(112), L.expo(120)
|
|
pmv = prev['mvA'] if prev else zeros(f'{D}/zero.mv', rw * rh * 2)
|
|
phalf = prev['half'] if prev else zeros(f'{D}/zero.half', (hw // 2) * (hh // 2) * 4, np.float32)
|
|
pfeat = prev['feat'] if prev else zeros(f'{D}/zero.feat', rw * rh * 4, np.int8)
|
|
phist = prev['hist'] if prev else zeros(f'{D}/zero.hist', hw * hh * 4)
|
|
open(f'{D}/e0.manifest', 'w').write('\n'.join([
|
|
f'param {D}/e0.params', f'grid {(gw + L.i2(8)[0] + 7) // 8} {(gh + L.i2(8)[1] + 7) // 8}',
|
|
f'tex 1 {colour} {rw} {rh} 4 16 nearest', f'tex 2 {motion} {rw} {rh} 2 16 nearest', f'tex 3 {depth} {rw} {rh} 1 32 nearest',
|
|
f'tex 4 {pmv} {rw} {rh} 2 16 linear', f'tex 5 {phalf} {hw // 2} {hh // 2} 4 32 linear', f'tex 6 {pfeat} {rw} {rh} 4 9 linear',
|
|
f'surf 7 {rw} {rh} 2 16 - {D}/mvA', f'surf 8 {rw} {rh} 2 16 - {D}/mvB', f'surf 9 {rw} {rh} 2 16 - {D}/mvC',
|
|
f'lut {D}/e0.w8', f'lut {D}/e0.w16', f'lut {D}/e0.noise', f'outbuf {D}/tokens {tw * th * 64}']) + '\n')
|
|
run([kvk, shader('enc0_m', STAGE), f'{D}/e0.manifest', '1', '32'])
|
|
ref_tokens = token_codes(np.fromfile(L.res[40], np.uint8), L.u8(48) - next(a for a, b in L.allocs.values() if a <= L.u8(48) < a + b.size), tw, th, 64)
|
|
print(f'frame {f}: ' + report('input tokens', np.fromfile(f'{D}/tokens', np.uint8).reshape(th, tw, 64), ref_tokens), flush=True)
|
|
# ---------------------------------------------------------------- network
|
|
Ps, dirs = mm.launches(cap, f)
|
|
own = {Ps[0].inp: f'{D}/tokens'}
|
|
for i, (P, d, name) in enumerate(zip(Ps, dirs, mm.ORDER)):
|
|
C, heads, NPM, CIN = mm.LAYERS[name]
|
|
o, b8, b16 = pack_layer(mm.plain_layer(P, name), name)
|
|
W, H = P.W, P.H
|
|
pre = f'{D}/l{i}'
|
|
open(pre + '.params', 'wb').write(struct.pack('<4i72I', W, H, P.sx, P.sy, *o)); b8.tofile(pre + '.w8'); b16.tofile(pre + '.w16')
|
|
skip = own[P.p32] if CIN else zeros(pre + '.noskip', 64, np.uint8)
|
|
msize = (H // 2) * (W // 2) * NPM if NPM else 64
|
|
run([evk, shader('layer_m', defines(name)), str((W + P.sx + 7) // 8), str((H + P.sy + 7) // 8), '1', f'in:0={pre}.params', f'in:1={pre}.w8', f'in:2={pre}.w16',
|
|
f'in:3={own[P.inp]}', f'in:4={skip}', f'out:5={pre}.out:{H * W * C}', f'out:6={pre}.merged:{msize}'])
|
|
own[P.out] = pre + '.out'
|
|
if NPM: own[P.p48] = pre + '.merged'
|
|
seq = int(os.path.basename(dirs[-1]).split('-')[1])
|
|
after = np.fromfile(glob.glob(f'{cap}/dump/launch-{seq:03d}-*-arg024-*')[0], np.uint8)
|
|
P = Ps[-1]
|
|
print(f'frame {f}: ' + report('network output', np.fromfile(own[P.out], np.uint8).reshape(P.H, P.W, 64), token_codes(after, P.out - P.abase, P.W, P.H, 64)), flush=True)
|
|
# ---------------------------------------------------------------- dec0
|
|
L = Launch(NAMES['dec0'], f)
|
|
wb = L.blob(L.u8(0))
|
|
cols = (np.arange(4)[:, None] * 32 + np.arange(20)[None, :]).ravel()
|
|
np.concatenate([wb[sm.woff_table(0, 512, 1024, 64, 128)][:, cols].ravel(), np.zeros(64, np.uint8)]).tofile(f'{D}/d0.w8')
|
|
np.concatenate([np.tile(wb[8192:8192 + 256].view(np.float16)[cols][None, :], (16, 1)).ravel(), np.zeros(32, np.float16)]).tofile(f'{D}/d0.w16')
|
|
open(f'{D}/d0.params', 'wb').write(struct.pack('<2iiI2f2f2f2f2f2f', rw, rh, tw, L.P[120], *L.f2(72), *L.f2(88), float(L.i2(96)[0]), float(L.i2(96)[1]), *L.f2(80), *L.f2(160), *L.f2(104)))
|
|
open(f'{D}/d0.manifest', 'w').write('\n'.join([
|
|
f'param {D}/d0.params', f'grid {gw // 8} {gh // 8}', f'tex 1 {D}/mvA {rw} {rh} 2 16 nearest', f'tex 2 {pfeat} {rw} {rh} 4 9 linear',
|
|
f'surf 3 {rw} {rh} 4 9 - {D}/feat', f'surf 4 {rw} {rh} 4 8 - {D}/gates', f'surf 5 {rw} {rh} 4 8 - {D}/mix', f'surf 6 {rw} {rh} 4 16 - {D}/cov',
|
|
f'lut {own[Ps[-1].out]}', f'lut {D}/d0.w8', f'lut {D}/d0.w16']) + '\n')
|
|
run([kvk, shader('dec0_m'), f'{D}/d0.manifest', '1', '32'])
|
|
for off, name, dt in ((128, 'feat', np.int8), (136, 'gates', np.uint8), (144, 'mix', np.uint8)):
|
|
a = np.fromfile(L.res[off], dt).astype(np.int32); b = np.fromfile(f'{D}/{name}', dt).astype(np.int32)
|
|
print(f'frame {f}: dec0 {name:5s} identical {100 * (a == b).mean():5.1f}%, within one step {100 * (np.abs(a - b) <= 1).mean():5.1f}%, mean |diff| {np.abs(a - b).mean():.3f} steps', flush=True)
|
|
# ---------------------------------------------------------------- reconstruction and output
|
|
L = Launch(NAMES['post'], f)
|
|
open(f'{D}/post.params', 'wb').write(struct.pack('<2i2i2i2f2f2ffIff', *L.i2(8), hw, hh, rw, rh, *L.f2(124), *L.f2(132), *L.f2(160), L.f1(120), L.P[168], L.expo(16), L.expo(24)))
|
|
open(f'{D}/post.manifest', 'w').write('\n'.join([
|
|
f'param {D}/post.params', f'grid {(hw + 15) // 16} {(hh + 15) // 16}',
|
|
f'tex 1 {colour} {rw} {rh} 4 16 nearest', f'tex 2 {D}/mix {rw} {rh} 4 8 nearest', f'tex 3 {D}/gates {rw} {rh} 4 8 nearest', f'tex 4 {D}/cov {rw} {rh} 4 16 nearest',
|
|
f'tex 5 {phist} {hw} {hh} 4 16 linear', f'tex 6 {D}/mvA {rw} {rh} 2 16 nearest', f'tex 7 {D}/mvB {rw} {rh} 2 16 nearest', f'tex 8 {D}/mvC {rw} {rh} 2 16 linear',
|
|
f'surf 9 {hw} {hh} 4 16 - {D}/hist', f'surf 10 {hw // 2} {hh // 2} 4 32 - {D}/half', f'outbuf {D}/packed {hw * hh * 4}']) + '\n')
|
|
run([kvk, shader('post_m', STAGE + POST_TILE), f'{D}/post.manifest', '1', '64'])
|
|
L = Launch(NAMES['down'], f)
|
|
ow, oh = L.i2(8)
|
|
open(f'{D}/down.params', 'wb').write(struct.pack('<8i4ff', *L.i2(0), *L.i2(40), *L.i2(64), *L.i2(72), *L.f2(80), *L.f2(88), L.expo(48)))
|
|
open(f'{D}/down.manifest', 'w').write('\n'.join([
|
|
f'param {D}/down.params', f'grid {(ow + 7) // 8} {(oh + 63) // 64}', f'tex 1 {colour} {rw} {rh} 4 16 nearest',
|
|
f'surf 2 {ow} {oh} 4 16 - {D}/output', f'lut {D}/packed']) + '\n')
|
|
run([kvk, shader('down_m', STAGE), f'{D}/down.manifest', '1', '64'])
|
|
a = np.fromfile(L.res[32], np.float16).astype(np.float64).reshape(oh, ow, 4)[..., :3]
|
|
b = np.fromfile(f'{D}/output', np.float16).astype(np.float64).reshape(oh, ow, 4)[..., :3]
|
|
d = a - b
|
|
print(f'frame {f}: FINAL OUTPUT {ow}x{oh}: {10 * np.log10(1.0 / np.mean(d * d)):.1f} dB PSNR (peak 1.0), max |diff| {np.abs(d).max():.3f}, non-finite {int((~np.isfinite(b)).sum())}', flush=True)
|
|
prev = {'mvA': f'{D}/mvA', 'half': f'{D}/half', 'feat': f'{D}/feat', 'hist': f'{D}/hist'}
|