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}')
+53 -20
View File
@@ -3,6 +3,7 @@
nvngx_dlssd.dll 310.9) from a capture of the denoiser.
usage: compile_rr2.py CAPTURE_DIR OUT_DIR [--ptx-dir DIR] [--frame N] [--native NVNGX_DLSSD_DLL]
compile_rr2.py --dll NVNGX_DLSSD_DLL OUT_DIR --ptx-dir DIR
CAPTURE_DIR is a bridge capture of a few preset F frames (engine/rr/rr_capture.py, taken with
D4R_CUDA_REPLAY_DUMP_FILTER=rr2 D4R_CUDA_LAUNCH_STATS_FILTER=rr2 and preset hint 6). Writes:
@@ -16,6 +17,9 @@ D4R_CUDA_REPLAY_DUMP_FILTER=rr2 D4R_CUDA_LAUNCH_STATS_FILTER=rr2 and preset hint
With --native, the stages engine/rr/rr2_m.comp implements (enc1 ... enc5, dec4 ... dec1, dec0's first kernel) are that
shader with weights read from the DLL itself (engine/rr/rr2_pack.py) instead of the translation: bit-identical to
NVIDIA's kernels on the reference captures and about three times faster. D4R_RR_MIX=1 builds the one-rounding variant.
--dll builds the same model without a capture: the native stages from the DLL's weights (as --native), the argument
templates, encoder 0's second table and the arena regions from engine/rr/rr2_templates.py, and the particle variant of
the downsample (downp) always. This is what install.sh runs.
Everything but the exposure shader is derived from NVIDIA's DLL: it stays on the user's machine and is never packaged.
"""
import argparse
@@ -88,28 +92,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 310.9")
ap.add_argument('--ptx-dir', default=os.path.expanduser('~/.cache/d4r-rr-corpus-f'))
ap.add_argument('--frame', type=int, default=2, help='which captured frame supplies the templates (1-based)')
ap.add_argument('--particles', metavar='CAPTURE_DIR',
help='a capture taken with DLSSD.ColorBeforeParticles: adds the downsample kernel NGX runs then (downp)')
ap.add_argument('--native', metavar='DLL', help='nvngx_dlssd.dll 310.9: build the native stages from its weights')
a = ap.parse_args()
cap = Capture(a.capture)
if a.dll:
if a.capture or a.particles:
ap.error('--dll takes no capture')
from rr2_templates import ARENA, aux_table, templates
dll = Path(a.dll).read_bytes()
blocks, aux = templates(dll), aux_table(dll)
a.native = a.native or a.dll
cap = None
elif not a.capture:
ap.error('a capture directory, or --dll')
else:
cap = Capture(a.capture)
out = a.output
out.mkdir(parents=True, exist_ok=True)
capped = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0']
# each compiler job in a memory-capped scope where systemd provides one (D4R_NO_SCOPE=1: none)
if os.environ.get('D4R_NO_SCOPE') == '1' or subprocess.run(capped[:4] + ['true'], stderr=subprocess.DEVNULL,
stdout=subprocess.DEVNULL).returncode:
capped = []
lines = ['D4RR0002', 'preset 6']
pcap = Capture(a.particles) if a.particles else None
for stage in STAGES + (['downp'] if pcap else []):
for stage in STAGES + (['downp'] if pcap or a.dll else []):
kernel = KERNEL[stage]
if stage == 'downp':
cap = pcap
launches = [l for l in cap.launches if l.kernel == kernel]
per_frame = len(launches) // len([l for l in cap.launches if l.kernel == KERNEL['enc0']])
launch = launches[(a.frame - 1) * per_frame]
network = stage in NETWORK
if a.dll:
launch = None
if network and stage not in NATIVE and stage not in HYBRID:
raise SystemExit(f'{stage}: no native stage, and its weights come from a capture')
else:
if stage == 'downp':
cap = pcap
launches = [l for l in cap.launches if l.kernel == kernel]
per_frame = len(launches) // len([l for l in cap.launches if l.kernel == KERNEL['enc0']])
launch = launches[(a.frame - 1) * per_frame]
if a.native and stage in NATIVE:
work = out / 'native-work' / stage
work.mkdir(parents=True, exist_ok=True)
@@ -123,10 +148,14 @@ def main():
if RELAXED < 2: # (the relaxed stages leave their conversions to the driver)
subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}_m.spv')],
check=True, stdout=subprocess.DEVNULL)
(out / f'{stage}.args').write_bytes(stage_params(launch, 24)[:len(launch.args)])
arena = launch.arena_alloc
unit = launch.allocs[arena][1] / 3152
ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers if p.alloc == arena and p.offset != 24]
if launch is None:
(out / f'{stage}.args').write_bytes(blocks[stage])
ptrs = [ARENA[stage]]
else:
(out / f'{stage}.args').write_bytes(stage_params(launch, 24)[:len(launch.args)])
arena = launch.arena_alloc
unit = launch.allocs[arena][1] / 3152
ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers if p.alloc == arena and p.offset != 24]
lines.append(f'stage {stage} {stage}_m.spv 32 {weights} {stage}.args - - - {",".join(ptrs)}')
print(f'{stage}: native (engine/rr/rr2_m.comp)')
continue
@@ -150,7 +179,8 @@ def main():
# encoder 0 in two passes (gen_rr2_enc0.py --split): its gather as the stage enc0g, whose records' place
# (three ranges) is six words appended to the argument block
split = stage == 'enc0' and SPLIT and GROUP == 1
how = ['--split', 'block', '--xf-word', str((len(launch.args) + 7) // 8 * 2)] if split else ['--group', str(GROUP)]
n_args = len(blocks[stage]) if launch is None else len(launch.args)
how = ['--split', 'block', '--xf-word', str((n_args + 7) // 8 * 2)] if split else ['--group', str(GROUP)]
if split:
subprocess.run([sys.executable, str(HERE / 'rr' / 'gen_rr2_enc0.py'), str(find_module(a.ptx_dir, kernel)),
str(generated), str(work / 'enc0_g.comp'), '--split', 'gather', *how[2:]],
@@ -183,16 +213,19 @@ def main():
weights = f'weights-{stage}.bin'
(out / weights).write_bytes(generated.with_suffix('.globals.bin').read_bytes())
if stage == 'enc0':
(out / 'enc0-aux.bin').write_bytes(launch.pre(launch.pointer(104).alloc))
args = stage_params(launch, 24 if network else 9999)[:len(launch.args)]
(out / 'enc0-aux.bin').write_bytes(aux if launch is None else launch.pre(launch.pointer(104).alloc))
args = blocks[stage] if launch is None else stage_params(launch, 24 if network else 9999)[:len(launch.args)]
if hybrid and split:
args = args + bytes(-len(args) % 8) + bytes(24) # the gather records' three ranges (the runtime's)
(out / f'{stage}.args').write_bytes(args)
# the stage's arena regions, in units of 1/3152 of NGX's arena (W0 * H0 / 32 bytes: the regions scale with it)
arena = launch.arena_alloc
unit = launch.allocs[arena][1] / 3152 if arena is not None else 1
ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers
if p.alloc == arena and p.offset not in BUFFERS.get(stage, []) and not (network and p.offset == 24)]
if launch is None:
ptrs = [ARENA[stage]] if network else []
else:
arena = launch.arena_alloc
unit = launch.allocs[arena][1] / 3152 if arena is not None else 1
ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers
if p.alloc == arena and p.offset not in BUFFERS.get(stage, []) and not (network and p.offset == 24)]
if not network:
ptrs = [] # the output passes' regions are the runtime's own
tex = ','.join(map(str, info['textures'])) or '-'
@@ -205,7 +238,7 @@ def main():
(f' x{GROUP}' if hybrid and GROUP > 1 else '')) # windows per workgroup (engine/rr/gen_rr2_enc0.py)
if hybrid and split: # the same block, bindings and regions; "rN": N words per window record, "s3": in three ranges
lines.append(f'stage enc0g enc0g.spv 32 - {stage}.args {tex} {surf} {buf} {",".join(ptrs) or "-"} r{record} s3')
print(f'{stage}: {kernel}, {len(launch.args)} argument bytes' + (' (native block)' if hybrid else ''))
print(f'{stage}: {kernel}, {len(args)} argument bytes' + (' (native block)' if hybrid else ''))
for p in (0, 1):
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', f'-DPASS={p}',
str(HERE / 'rr' / 'exposure_f.comp'), '-o', str(out / f'exposure_f{p}.spv')],
+32 -8
View File
@@ -5,6 +5,11 @@ set -euo pipefail
# for the preset the capture was taken with (D or E).
#
# usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL]
# build_native_model.sh --dll NVNGX_DLSSD_DLL D|E OUT_DIR
#
# With --dll no capture is needed: everything comes from the DLL (its PTX: kernels/tools/extract_dlss_ptx.py; the
# argument-block templates: engine/rr/rr_templates.py; the weights: dll_weights.py, model 0 = D, model 1 = E in 310.7's
# file order) and step 3's checks against a captured frame are skipped. This is what install.sh runs.
#
# CAPTURE_DIR is a bridge capture of a few denoiser frames (engine/rr/rr_capture.py); the DLL is the one it was taken
# with (default: CAPTURE_DIR/run/app/nvngx_dlssd.dll, then $D4R_RR_DLL). The steps:
@@ -23,19 +28,38 @@ set -euo pipefail
#
# GPU jobs run one at a time, each in a memory-capped scope (a translated stage's pipeline compile can take tens of GB).
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
CAP="$(realpath "${1:?usage: build_native_model.sh CAPTURE_DIR OUT_DIR}")"
OUT="$(realpath -m "${2:?usage: build_native_model.sh CAPTURE_DIR OUT_DIR}")"
USAGE="usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL] | --dll NVNGX_DLSSD_DLL D|E OUT_DIR"
CAP=""
if [[ "${1:-}" == --dll ]]; then
DLL="$(realpath "${2:?$USAGE}")"; WANT="${3:?$USAGE}"; OUT="$(realpath -m "${4:?$USAGE}")"
[[ "$WANT" == D || "$WANT" == E ]] || { echo "$USAGE" >&2; exit 2; }
else
CAP="$(realpath "${1:?$USAGE}")"
OUT="$(realpath -m "${2:?$USAGE}")"
DLL="${3:-${D4R_RR_DLL:-$CAP/run/app/nvngx_dlssd.dll}}"
fi
RR="$ROOT/engine/rr"
WORK="$OUT/native-work"
DLL="${3:-${D4R_RR_DLL:-$CAP/run/app/nvngx_dlssd.dll}}"
[[ -f "$DLL" ]] || { echo "no nvngx_dlssd.dll at $DLL: pass the DLL the capture was taken with" >&2; exit 1; }
FRAME="${D4R_RR_MODEL_FRAME:-2}"
CAPPED=(systemd-run --user --scope -q -p MemoryMax=24G -p MemorySwapMax=0)
# each GPU / compiler job in a memory-capped scope where systemd provides one (D4R_NO_SCOPE=1: none)
CAPPED=()
if [[ "${D4R_NO_SCOPE:-0}" != 1 ]] && command -v systemd-run > /dev/null && systemd-run --user --scope -q true 2> /dev/null; then
CAPPED=(systemd-run --user --scope -q -p MemoryMax=24G -p MemorySwapMax=0)
fi
mkdir -p "$OUT" "$WORK"
cd "$ROOT"
python3 engine/compile_rr.py "$CAP" "$OUT" > "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; }
python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --capture "$CAP" ${D4R_RR_ANY_KERNELS:+--any-kernels}
if [[ -n "$CAP" ]]; then
python3 engine/compile_rr.py "$CAP" "$OUT" > "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; }
python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --capture "$CAP" ${D4R_RR_ANY_KERNELS:+--any-kernels}
else
python3 kernels/tools/extract_dlss_ptx.py "$DLL" "$WORK/ptx" > /dev/null
python3 engine/compile_rr.py --dll "$DLL" --preset "$WANT" "$OUT" --ptx-dir "$WORK/ptx" --network-bindings-only \
> "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; }
python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --model "$([[ "$WANT" == D ]] && echo 0 || echo 1)" \
${D4R_RR_ANY_KERNELS:+--any-kernels} > /dev/null
fi
PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")"
# RR_EXACT_CONVERT: the compiled module's float conversions are decorated NoContraction (spv_exact_convert.py), which
# keeps the driver from folding a binary16 rounding away without the opaque OR the shaders otherwise need
@@ -68,7 +92,7 @@ for s in enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1; do
b=1,1; case $s in enc1|enc2|dec1|dec2) b=5,5 ;; esac
d="$WORK/$s"; mkdir -p "$d"
cp "$WORK/dll/$s.npz" "$d/weights.npz"
python3 "$RR/swin_ref.py" "$CAP" "$s" "$FRAME" "$d/weights.npz" --blocks "$b" 0,0 | tail -2 | sed "s/^/$s: /"
[[ -z "$CAP" ]] || python3 "$RR/swin_ref.py" "$CAP" "$s" "$FRAME" "$d/weights.npz" --blocks "$b" 0,0 | tail -2 | sed "s/^/$s: /"
python3 "$RR/swin_pack.py" "$s" "$d/weights.npz" "$OUT/weights-${s}m.bin" "$d/swin_stage.h" > /dev/null
glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR"
native "$s" "weights-${s}m.bin"
@@ -77,7 +101,7 @@ done
# decoder 0
d="$WORK/dec0"; mkdir -p "$d"
cp "$WORK/dll/dec0.npz" "$d/weights.npz"
python3 "$RR/dec0_ref.py" "$CAP" "$FRAME" "$d/weights.npz" --blocks 5,5 0,0 | tail -2 | sed 's/^/dec0: /'
[[ -z "$CAP" ]] || python3 "$RR/dec0_ref.py" "$CAP" "$FRAME" "$d/weights.npz" --blocks 5,5 0,0 | tail -2 | sed 's/^/dec0: /'
python3 "$RR/dec0_pack.py" "$d/weights.npz" "$OUT/weights-dec0m.bin" "$d/dec0_m_w.h" > /dev/null
glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR"
native dec0 weights-dec0m.bin
+5 -2
View File
@@ -133,8 +133,11 @@ def noise_table(dll):
if b - a >= 2048:
t = u[b - 2048:b]
if len(np.unique(t)) > 1024 and (t & 31 == 31).sum() > 256:
found.append(t.tobytes())
return found[0] if len(set(found)) == 1 else None # some DLLs hold the same table twice
found.append((2 * (b - 2048), t.tobytes()))
# 310.9.1 also holds a second, unreferenced array of this kind at an odd address; the table the code uses is aligned
if len({t for _, t in found}) > 1:
found = [(at, t) for at, t in found if at % 16 == 0]
return found[0][1] if len({t for _, t in found}) == 1 else None # some DLLs hold the same table twice
class Reader:
+263
View File
@@ -0,0 +1,263 @@
#!/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()
+103
View File
@@ -0,0 +1,103 @@
#!/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()