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.
This commit is contained in:
woutermans committed 2026-10-14 20:09:08 +02:00
1 parent e3606cedf9
commit 5f520a186d
7 files changed
+508 -45

No files matched your search

+46 -14
View File
@@ -1,7 +1,8 @@
#!/usr/bin/env python3
"""Build the engine's Ray Reconstruction model directory (preset D or E) from a capture of the denoiser.
usage: compile_rr.py CAPTURE_DIR OUT_DIR [--ptx-dir DIR] [--frame N]
usage: compile_rr.py CAPTURE_DIR OUT_DIR [--ptx-dir DIR] [--frame N] [--network-bindings-only]
compile_rr.py --dll NVNGX_DLSSD_DLL --preset D|E OUT_DIR [--ptx-dir DIR] --network-bindings-only
CAPTURE_DIR is a bridge capture of a few denoiser frames (engine/rr/rr_capture.py: replay/ with the launches' weight
allocations, dump/, run/run.log), taken with the preset the model is for. Writes:
@@ -14,6 +15,12 @@ generated from the kernel's PTX in the user's own nvngx_dlssd.dll by engine/rr/p
kernels/tools/extract_dlss_ptx.py wrote, named NNNN-<first entry>.ptx; default ~/.cache/d4r-rr-corpus).
The weights, the templates and any generated shader are derived from NVIDIA's DLL: they stay on the user's machine
and are never packaged (scripts/check_redistributable.py).
With --dll instead of a capture nothing is taken from a capture: the templates come from engine/rr/rr_templates.py,
the output stage's weight image is NGX's 16 zero bytes, and the network stages get no weight image, so the model is
only usable once engine/rr/build_native_model.sh has replaced them with the native stages (weights read from the DLL).
--network-bindings-only translates enc0 ... dec0 for their bindings (and enc0's gather, engine/rr/gen_enc0_m.py) but
does not compile them: build_native_model.sh replaces them, and the translations' pipeline compiles take tens of GB.
"""
import argparse
import json
@@ -63,26 +70,49 @@ def find_module(ptx_dir, kernel):
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument('capture')
ap.add_argument('capture', nargs='?')
ap.add_argument('output', type=Path)
ap.add_argument('--dll', help="build without a capture, from the user's nvngx_dlssd.dll (with --preset)")
ap.add_argument('--preset', choices=['D', 'E'])
ap.add_argument('--network-bindings-only', action='store_true',
help='translate the network stages for their bindings only, without compiling them')
ap.add_argument('--ptx-dir', default=os.path.expanduser('~/.cache/d4r-rr-corpus'))
ap.add_argument('--frame', type=int, default=2, help='which captured frame supplies the templates (1-based)')
ap.add_argument('--reset-only', action='store_true', help='use the own output shader that has no history path')
a = ap.parse_args()
cap = Capture(a.capture)
preset_d = any('diamond_wallaby' in l.kernel for l in cap.launches)
if a.dll:
if a.capture or not a.preset or not a.network_bindings_only:
ap.error('--dll takes --preset and --network-bindings-only, and no capture')
cap = None
preset_d = a.preset == 'D'
from rr_templates import templates
dll_templates = dict(zip(('enc0', 'hkpn'), templates(Path(a.dll).read_bytes(), a.preset)))
else:
if not a.capture:
ap.error('a capture directory, or --dll')
cap = Capture(a.capture)
preset_d = any('diamond_wallaby' in l.kernel for l in cap.launches)
kernels = dict(KERNEL, **(D_KERNEL if preset_d else {}))
a.output.mkdir(parents=True, exist_ok=True)
lines = ['D4RR0001', f'preset {4 if preset_d else 5}']
info = {'preset': 'D' if preset_d else 'E', 'stages': {}}
for stage in STAGES:
launches = [l for l in cap.launches if l.kernel == kernels[stage]]
if len(launches) < a.frame:
raise SystemExit(f'the capture has {len(launches)} launches of {kernels[stage]}; frame {a.frame} needed')
launch = launches[a.frame - 1]
(a.output / f'weights-{stage}.bin').write_bytes(launch.weights)
if stage in ('enc0', 'hkpn'):
(a.output / f'{stage}.args').write_bytes(stage_params(launch))
network = stage != 'hkpn'
if cap is None:
# the output kernel's weight argument: NGX uploads 16 zero bytes for it (every capture of 310.7, D and E)
weights = bytes(16) if not network else None
if stage in dll_templates:
(a.output / f'{stage}.args').write_bytes(dll_templates[stage])
else:
launches = [l for l in cap.launches if l.kernel == kernels[stage]]
if len(launches) < a.frame:
raise SystemExit(f'the capture has {len(launches)} launches of {kernels[stage]}; frame {a.frame} needed')
launch = launches[a.frame - 1]
weights = launch.weights
if stage in ('enc0', 'hkpn'):
(a.output / f'{stage}.args').write_bytes(stage_params(launch))
if weights is not None:
(a.output / f'weights-{stage}.bin').write_bytes(weights)
own = HERE / 'rr' / (f'{stage}_d.comp' if preset_d and stage in D_KERNEL else f'{stage}.comp')
spv = a.output / f'{stage}.spv'
if own.exists() and (stage not in RESET_ONLY or a.reset_only):
@@ -94,15 +124,17 @@ def main():
formats = [f'--surface={off}={fmt}' for off, fmt in SURFACE_FORMATS.get(stage, {}).items()]
subprocess.run([sys.executable, str(HERE / 'rr' / 'ptx2rr.py'), str(find_module(a.ptx_dir, kernels[stage])),
kernels[stage], str(generated), *formats], check=True)
glsl(generated, spv)
if not (network and a.network_bindings_only):
glsl(generated, spv)
origin = 'generated from the DLL'
binding = json.loads(generated.with_suffix('.json').read_text())
textures, surfaces = binding['textures'], [tuple(x) for x in binding['surfaces']]
tex = ','.join(map(str, textures)) or '-'
surf = ','.join(f'{off}:{fmt}' for off, fmt in surfaces) or '-'
lines.append(f'stage {stage} {stage}.spv 32 weights-{stage}.bin {tex} {surf}' + (' noopt' if stage in NOOPT else ''))
info['stages'][stage] = {'kernel': kernels[stage], 'shader': origin, 'weights': len(launch.weights)}
print(f'{stage}: {kernels[stage]} ({origin}), {len(launch.weights)} weight bytes')
n = len(weights) if weights is not None else 0
info['stages'][stage] = {'kernel': kernels[stage], 'shader': origin, 'weights': n}
print(f'{stage}: {kernels[stage]} ({origin}), {n} weight bytes')
lines += ['template enc0 enc0.args', 'template hkpn hkpn.args']
for p in (0, 1):
glsl(HERE / 'exposure_k.comp', a.output / f'exposure_rr{p}.spv', '-DMODE=2', f'-DPASS={p}')