woutermans
5f520a186d
Models built from nvngx_dlssd.dll 310.9.1 are byte-identical to the capture-built ones (every file, plus d4r-rr's outputs on capF, capF4K, capF900, capFXparticles and capDT). 310.9.1 contains 310.7's D/E kernels and weights.
264 lines
12 KiB
Python
264 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""What a preset F model takes from a capture, read from nvngx_dlssd.dll 310.9 instead: the argument-block templates
|
|
(<stage>.args), encoder 0's second table (enc0-aux.bin, argument 104) and the stages' plane-arena regions.
|
|
|
|
usage: rr2_templates.py NVNGX_DLSSD_DLL OUT_DIR write the templates and enc0-aux.bin
|
|
rr2_templates.py NVNGX_DLSSD_DLL --check MODEL_DIR compare with a model compile_rr2.py built from a capture
|
|
|
|
engine/runtime_rr2.inc fills every field of the blocks that depends on the frame, the geometry or the bound inputs.
|
|
Across the preset F captures (2026-10-14: capF, capF900, capF1440, capF4K, capFQ and the capFX* input variants) what
|
|
is left is per stage:
|
|
network stages binary16 s0 (byte 92), s1 (94), s2 (96) and, for the decoders, s3 (98) (engine/rr/rr2_ref.py)
|
|
enc0 two float32 (572, 576), two binary16 widened to float32 (580, 584), 1.0 (588, 1072), 1 (1036)
|
|
dec0_1 8.0 (208) and exp(t) (212); post: exp(t') (128); down / downp: 0.25 (204), 1.0 (296, 300)
|
|
NGX's model description is built by one routine of plain stores (mov [rdx + d], imm32 / ax / a pointer); describe()
|
|
replays it. It holds every tensor's pointer and size, s1 / s2 of each stage at 0x220 + 0x238 * n (a decoder level's
|
|
two kernels side by side, its s3 behind them), encoder 0's float32 pair at 0x230 and the binary16 pair at 0x185a, and
|
|
three one-value tensors (t at 0x1838, t' at 0x1848). s0 sits in the weights, before each encoder's relative position
|
|
tables (encoder 0: 16 bytes before; a decoder level uses its encoder's).
|
|
|
|
enc0-aux.bin is preset E's noise table (engine/rr/dll_weights.py) and then encoder 0's embedding, both branches: Wd
|
|
(32 x 64 FP8 codes, 43 inputs used), bd, Wu (32 x 32 FP8), bu (once, after the first branch), the FP8 matrices in the
|
|
kernel's fragment order: a fixed permutation of the index bits of the natural [in][out] order (FRAG_WD, FRAG_WU).
|
|
Everything written is derived from the user's DLL and stays on their machine.
|
|
"""
|
|
import math
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from dll_weights import noise_table # noqa: E402
|
|
from enc1_ref import inter # noqa: E402
|
|
from rr2_ref import ORDER, dll_weights, find_tables # noqa: E402
|
|
|
|
H = np.float16
|
|
NETWORK = ORDER # enc0 ... enc5, dec4_0, dec4_1, ... dec0_1
|
|
ARG_BYTES = dict({s: 104 for s in NETWORK}, enc0=1088, dec0_1=216, post=320, down=416, downp=416, pooled=32)
|
|
# the plane-arena regions of the network stages: argument offset -> region start in units of 1/3152 of NGX's arena
|
|
# (W0 * H0 / 32 bytes; the same on every capture)
|
|
ARENA = {
|
|
'enc0': '32:0,40:1024,64:2560,72:2240', 'enc1': '8:0,32:1536,40:2048,48:2240,64:3008,72:2304',
|
|
'enc2': '8:2048,32:1024,40:1216,48:2304,64:1472,72:1364', 'enc3': '8:1216,32:1280,40:1344,48:1364,64:1368,72:1372',
|
|
'enc4': '8:1344,32:1236,40:1216,48:1372,64:1261,72:1221', 'enc5': '8:1216,32:1256,48:1221,64:1262',
|
|
'dec4_0': '8:1256,16:1236,32:1216', 'dec4_1': '8:1216,32:1344,48:1262,56:1261,64:1364',
|
|
'dec3_0': '8:1344,16:1280,32:1216', 'dec3_1': '8:1216,32:1408,48:1364,56:1368,64:1488',
|
|
'dec2_0': '8:1408,16:1024,32:1216', 'dec2_1': '8:1216,32:2816,48:1488,56:1472,64:3136',
|
|
'dec1_0': '8:2816,16:1536,32:1024', 'dec1_1': '8:1024,32:2048,48:3136,56:3008,64:3072',
|
|
'dec0_0': '8:2048,16:0,32:1024', 'dec0_1': '8:1024,48:3072,56:2560,64:2816',
|
|
}
|
|
# fragment order of encoder 0's FP8 matrices in enc0-aux.bin: bit j of the table index is bit FRAG[j] of the natural
|
|
# index in * 32 + out (found on 310.9.1 and checked on every entry of both branches)
|
|
FRAG_WD = [5, 9, 8, 3, 6, 7, 0, 1, 2, 10, 4]
|
|
FRAG_WU = [5, 9, 8, 3, 6, 7, 0, 1, 2, 4]
|
|
|
|
|
|
def sections(dll):
|
|
pe = struct.unpack_from('<I', dll, 0x3c)[0]
|
|
count, opt = struct.unpack_from('<HH', dll, pe + 6)[0], struct.unpack_from('<H', dll, pe + 20)[0]
|
|
base = struct.unpack_from('<Q', dll, pe + 24 + 24)[0]
|
|
out = []
|
|
for i in range(count):
|
|
at = pe + 24 + opt + 40 * i
|
|
vsize, va, rsize, raw = struct.unpack_from('<IIII', dll, at + 8)
|
|
out.append((base + va, max(vsize, rsize), raw))
|
|
return out
|
|
|
|
|
|
def file_offset(dll, va):
|
|
for start, size, raw in sections(dll):
|
|
if start <= va < start + size:
|
|
return va - start + raw
|
|
raise ValueError(f'address {va:#x} is in no section')
|
|
|
|
|
|
def modrm(b, i):
|
|
m = b[i]
|
|
mod, reg, rm = m >> 6, (m >> 3) & 7, m & 7
|
|
if rm == 4 or (mod == 0 and rm == 5) or mod == 3:
|
|
raise ValueError
|
|
base = {2: 'rdx', 1: 'rcx'}.get(rm)
|
|
if mod == 0:
|
|
return base, reg, 0, 1
|
|
if mod == 1:
|
|
return base, reg, struct.unpack_from('<b', b, i + 1)[0], 2
|
|
return base, reg, struct.unpack_from('<i', b, i + 1)[0], 5
|
|
|
|
|
|
def replay(dll, i, image_base):
|
|
"""the stores of the routine starting at file offset i, up to its ret: ({offset: bytes}, {offset: pointer})"""
|
|
stores, pointers = {}, {}
|
|
eax, r8, rax_ptr = 0, 0, None
|
|
text = next((s, sz, raw) for s, sz, raw in sections(dll) if raw <= i < raw + sz)
|
|
for _ in range(4000):
|
|
b = dll[i:i + 16]
|
|
p66 = b[0] == 0x66
|
|
j = 1 if p66 else 0
|
|
rex = b[j] if 0x40 <= b[j] <= 0x4f else 0
|
|
j += 1 if rex else 0
|
|
op, w, rexr = b[j], bool(rex & 8), bool(rex & 4)
|
|
size = 8 if w else 2 if p66 else 4
|
|
if op == 0xc3:
|
|
return stores, pointers
|
|
if op == 0x33 and rex == 0x45 and b[j + 1] == 0xc0: # xor r8d, r8d
|
|
r8 = 0; i += j + 2; continue
|
|
if op == 0xb8 and not rex: # mov eax, imm32
|
|
eax = struct.unpack_from('<I', b, j + 1)[0]; i += j + 5; continue
|
|
if op == 0x8d and w and b[j + 1] == 0x05: # lea rax, [rip + d]
|
|
rel = struct.unpack_from('<i', b, j + 2)[0]
|
|
rax_ptr = text[0] + (i + j + 6 - text[2]) + rel; i += j + 6; continue
|
|
if op == 0x8b: # mov rax, [rip + d] / mov rax, rcx
|
|
rax_ptr = None; i += j + (2 if b[j + 1] >> 6 == 3 else 6); continue
|
|
if op == 0xc7: # mov [rdx + d], imm32
|
|
base, _, d, n = modrm(b, j + 1)
|
|
if base == 'rdx':
|
|
stores[d] = struct.pack('<q' if w else '<i', struct.unpack_from('<i', b, j + 1 + n)[0])
|
|
i += j + 1 + n + 4; continue
|
|
if op == 0x89: # mov [rdx + d], eax / ax / rax / r8
|
|
base, reg, d, n = modrm(b, j + 1)
|
|
if base == 'rdx' and reg == 0:
|
|
if w and not rexr:
|
|
pointers[d] = rax_ptr
|
|
else:
|
|
v = r8 if rexr else eax
|
|
stores[d] = (v & ((1 << (8 * size)) - 1)).to_bytes(size, 'little')
|
|
i += j + 1 + n; continue
|
|
raise ValueError(f'instruction at file offset {i:#x} ({b[:8].hex()})')
|
|
raise ValueError('no ret')
|
|
|
|
|
|
def describe(dll):
|
|
"""replay the routine that builds the model description: the one that stores encoder 0's tensor sizes 2752
|
|
(Wd, 2 x 32 x 43) at offset 0 and 64 (bd) at 0x10"""
|
|
marker = b'\xc7\x02\xc0\x0a\x00\x00\xc7\x42\x10\x40\x00\x00\x00'
|
|
found = []
|
|
at = dll.find(marker)
|
|
while at >= 0:
|
|
start = dll.rfind(b'\xcc\xcc', 0, at) + 2
|
|
try:
|
|
stores, pointers = replay(dll, start, None)
|
|
if 0x220 in stores and 0x1858 in stores:
|
|
found.append((stores, pointers))
|
|
except ValueError:
|
|
pass
|
|
at = dll.find(marker, at + 1)
|
|
if len(found) != 1:
|
|
raise SystemExit(f'found {len(found)} preset F model descriptions in the DLL (expected one): not a layout this '
|
|
'builder knows')
|
|
return found[0]
|
|
|
|
|
|
def h16(b):
|
|
return np.frombuffer(b, H)[0]
|
|
|
|
|
|
def values(dll):
|
|
"""{stage: {argument offset: bytes}} of every constant the runtime keeps"""
|
|
stores, pointers = describe(dll)
|
|
image = bytearray(max(d + len(v) for d, v in stores.items()))
|
|
for d in sorted(stores):
|
|
image[d:d + len(stores[d])] = stores[d]
|
|
word = lambda d: bytes(image[d:d + 2])
|
|
tables = find_tables(dll)
|
|
s0 = {}
|
|
for n, stage in enumerate(['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5']):
|
|
at = tables[stage] - (16 if stage == 'enc0' else 8)
|
|
s0[stage] = dll[at:at + 2]
|
|
out = {}
|
|
for n, stage in enumerate(NETWORK):
|
|
v = {}
|
|
if stage.startswith('enc'):
|
|
k = int(stage[3])
|
|
v[92] = s0[stage]
|
|
d = 0x220 + 0x238 * k
|
|
v[94] = bytes(image[d:d + 4])
|
|
else:
|
|
level, half = int(stage[3]), int(stage[5])
|
|
v[92] = s0[f'enc{level + 1}']
|
|
d = 0x220 + 0x238 * (10 - level)
|
|
v[94] = bytes(image[d + 4 * half:d + 4 * half + 4])
|
|
v[98] = word(d + 8)
|
|
out[stage] = v
|
|
f16f32 = lambda b: struct.pack('<f', float(h16(b)))
|
|
one = struct.pack('<f', 1.0)
|
|
out['enc0'].update({572: bytes(image[0x230:0x234]), 576: bytes(image[0x234:0x238]), 580: f16f32(word(0x185c)), 584: f16f32(word(0x185a)),
|
|
588: one, 1036: struct.pack('<I', 1), 1072: one})
|
|
tensor = lambda d: dll[file_offset(dll, pointers[d]):file_offset(dll, pointers[d]) + 2]
|
|
expf = lambda b: struct.pack('<f', np.float32(math.exp(float(h16(b)))))
|
|
out['dec0_1'].update({208: struct.pack('<f', 8.0), 212: expf(tensor(0x1838))})
|
|
out['post'] = {128: expf(tensor(0x1848))}
|
|
for s in ('down', 'downp'):
|
|
out[s] = {204: struct.pack('<f', 0.25), 296: one, 300: one}
|
|
out['pooled'] = {}
|
|
return out
|
|
|
|
|
|
def templates(dll):
|
|
blocks = {}
|
|
for stage, v in values(dll).items():
|
|
b = bytearray(ARG_BYTES[stage])
|
|
for off, x in v.items():
|
|
b[off:off + len(x)] = x
|
|
blocks[stage] = bytes(b)
|
|
return blocks
|
|
|
|
|
|
def frag(natural, bits):
|
|
"""natural [in][out] FP8 codes -> the table's order: bit j of the table index is bit bits[j] of the natural index"""
|
|
flat = natural.ravel()
|
|
idx = np.arange(flat.size)
|
|
src = np.zeros_like(idx)
|
|
for j, b in enumerate(bits):
|
|
src |= ((idx >> j) & 1) << b
|
|
return flat[src].tobytes()
|
|
|
|
|
|
def aux_table(dll):
|
|
noise = noise_table(dll)
|
|
if noise is None:
|
|
raise SystemExit('no noise table found in the DLL')
|
|
w = dll_weights(dll, 'enc0')
|
|
|
|
def natural(b_operand): # B = q8(mat).T[inter(k)] (k x n) -> [in][out]
|
|
n = np.empty_like(b_operand)
|
|
n[np.asarray(inter(b_operand.shape[0]))] = b_operand
|
|
return n
|
|
out = [noise]
|
|
for b in range(2):
|
|
out += [frag(natural(np.asarray(w['ewd'][b])), FRAG_WD), np.asarray(w['ebd'][b], H).tobytes(),
|
|
frag(natural(np.asarray(w['ewu'][b])), FRAG_WU)]
|
|
if b == 0:
|
|
out.append(np.asarray(w['ebu'], H).tobytes())
|
|
return b''.join(out)
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) < 3:
|
|
raise SystemExit(__doc__)
|
|
dll = Path(sys.argv[1]).read_bytes()
|
|
blocks, aux = templates(dll), aux_table(dll)
|
|
if sys.argv[2] == '--check':
|
|
model, bad = Path(sys.argv[3]), 0
|
|
theirs = (model / 'enc0-aux.bin').read_bytes()
|
|
print('enc0-aux.bin:', 'same' if theirs == aux else 'DIFFERENT'); bad += theirs != aux
|
|
for stage, ours in blocks.items():
|
|
path = model / f'{stage}.args'
|
|
if not path.exists():
|
|
continue
|
|
t = path.read_bytes()[:len(ours)]
|
|
diff = [i for i in range(len(ours)) if ours[i] != t[i]]
|
|
words = sorted({i // 4 * 4 for i in diff})
|
|
print(f'{stage}: ' + ('same' if not diff else f'{len(words)} words differ: {words[:40]}'))
|
|
sys.exit(1 if bad else 0)
|
|
out = Path(sys.argv[2])
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
for stage, b in blocks.items():
|
|
(out / f'{stage}.args').write_bytes(b)
|
|
(out / 'enc0-aux.bin').write_bytes(aux)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|