Files
d4r/engine/rr/rr_templates.py
woutermans 5f520a186d engine/rr: build the Ray Reconstruction models from nvngx_dlssd.dll alone, without a capture
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.
2026-10-14 20:09:08 +02:00

104 lines
4.5 KiB
Python

#!/usr/bin/env python3
"""The two argument-block templates of a preset D / E model (enc0.args, hkpn.args) without a capture.
usage: rr_templates.py NVNGX_DLSSD_DLL D|E OUT_DIR
rr_templates.py NVNGX_DLSSD_DLL D|E --check MODEL_DIR... (compare with templates compile_rr.py took from captures)
engine/runtime_rr.inc fills every field of the two blocks that depends on the frame, the geometry or the bound inputs.
Across captures of different sizes, inputs and frames (2026-10-14: capD, capDT, capD4K, capDX*, capE, capE3) what is
left is the same few constants for both presets, plus two floats of encoder 0 (bytes 540 and 544) that differ between
the presets. NGX sets those two in code: `sub ecx, 4; je D; cmp ecx, 1; jne ...; mov [rbx + X], E0; mov [rbx + X + 4],
E1; jmp ...; D: mov [rbx + X], D0; mov [rbx + X + 4], D1` (ecx = the preset), which find_floats() looks up in the DLL.
The templates are written on the user's machine and are never packaged.
"""
import re
import struct
import sys
from pathlib import Path
ENC0_BYTES, HKPN_BYTES = 1032, 480
# (offset, packed value) of the constants the runtime keeps
ENC0_CONST = [(360, struct.pack('<f', 1.0)), (532, struct.pack('<I', 2)), (988, struct.pack('<I', 1))]
HKPN_CONST = [(148, struct.pack('<f', 1.0)), (436, struct.pack('<f', 0.25)), (456, struct.pack('<I', 1))]
# the byte ranges engine/runtime_rr.inc writes every frame (setFrame), for --check
ENC0_WRITTEN = [(0, 8), (48, 64), (80, 84), (136, 160), (312, 352), (364, 368), (432, 440), (472, 480), (528, 532),
(536, 537), (539, 540), (596, 980), (980, 981), (984, 988), (1024, 1032)]
HKPN_WRITTEN = [(0, 48), (56, 64), (140, 148), (152, 156), (160, 161), (256, 264), (336, 344), (432, 433), (440, 444)]
# mov [rbx + disp32], imm32 is c7 83 disp32 imm32
PATTERN = re.compile(rb'\x83\xe9\x04\x74.\x83\xf9\x01\x75.'
rb'\xc7\x83(....)(....)\xc7\x83(....)(....)\xeb.'
rb'\xc7\x83(....)(....)\xc7\x83(....)(....)', re.S)
def find_floats(dll):
"""{'D': (f0, f1), 'E': (f0, f1)}: encoder 0's bytes 540 / 544 as NGX's code sets them"""
found = set()
for m in PATTERN.finditer(dll):
e0, e0v, e1, e1v, d0, d0v, d1, d1v = m.groups()
x = struct.unpack('<I', e0)[0]
if struct.unpack('<I', e1)[0] == x + 4 and e0 == d0 and e1 == d1:
found.add((d0v + d1v, e0v + e1v))
if len(found) != 1:
raise SystemExit(f'found {len(found)} candidates for the preset constants of encoder 0 (expected one): '
'not a DLL layout this builder knows')
d, e = found.pop()
return {'D': d, 'E': e}
def templates(dll, preset):
enc0, hkpn = bytearray(ENC0_BYTES), bytearray(HKPN_BYTES)
for off, v in ENC0_CONST:
enc0[off:off + len(v)] = v
for off, v in HKPN_CONST:
hkpn[off:off + len(v)] = v
enc0[540:548] = find_floats(dll)[preset]
return bytes(enc0), bytes(hkpn)
def masked(b, written, handles):
b = bytearray(b)
for lo, hi in written:
b[lo:hi] = bytes(hi - lo)
for off in handles:
b[off:off + 8] = bytes(8)
return bytes(b)
def handles(rr_txt, stage):
for line in Path(rr_txt).read_text().splitlines():
f = line.split()
if f[:2] == ['stage', stage]:
tex = [] if f[5] == '-' else [int(x) for x in f[5].split(',')]
surf = [] if f[6] == '-' else [int(x.split(':')[0]) for x in f[6].split(',')]
return tex + surf
return []
def main():
if len(sys.argv) < 4:
raise SystemExit(__doc__)
dll, preset = Path(sys.argv[1]).read_bytes(), sys.argv[2].upper()
if preset not in ('D', 'E'):
raise SystemExit('preset D or E')
enc0, hkpn = templates(dll, preset)
if sys.argv[3] == '--check':
bad = 0
for model in sys.argv[4:]:
for name, ours, written in (('enc0', enc0, ENC0_WRITTEN), ('hkpn', hkpn, HKPN_WRITTEN)):
theirs = (Path(model) / f'{name}.args').read_bytes()
h = handles(Path(model) / 'rr.txt', name)
a, b = masked(ours, written, h), masked(theirs, written, h)
diff = [i for i in range(len(a)) if a[i] != b[i]]
bad += bool(diff)
print(f'{model} {name}: ' + ('same' if not diff else f'{len(diff)} bytes differ, first at {diff[:8]}'))
sys.exit(1 if bad else 0)
out = Path(sys.argv[3])
out.mkdir(parents=True, exist_ok=True)
(out / 'enc0.args').write_bytes(enc0)
(out / 'hkpn.args').write_bytes(hkpn)
if __name__ == '__main__':
main()