diff --git a/engine/compile_rr.py b/engine/compile_rr.py index 6d34067..e204c8b 100755 --- a/engine/compile_rr.py +++ b/engine/compile_rr.py @@ -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-.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}') diff --git a/engine/compile_rr2.py b/engine/compile_rr2.py index 9f74e0a..eb97e0e 100644 --- a/engine/compile_rr2.py +++ b/engine/compile_rr2.py @@ -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')], diff --git a/engine/rr/build_native_model.sh b/engine/rr/build_native_model.sh index 85df276..88f4dda 100755 --- a/engine/rr/build_native_model.sh +++ b/engine/rr/build_native_model.sh @@ -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 diff --git a/engine/rr/dll_weights.py b/engine/rr/dll_weights.py index 2382e5e..c1fd50b 100644 --- a/engine/rr/dll_weights.py +++ b/engine/rr/dll_weights.py @@ -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: diff --git a/engine/rr/rr2_templates.py b/engine/rr/rr2_templates.py new file mode 100644 index 0000000..2708bfe --- /dev/null +++ b/engine/rr/rr2_templates.py @@ -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 +(.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('> 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('> 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('= 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(' 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() diff --git a/engine/rr/rr_templates.py b/engine/rr/rr_templates.py new file mode 100644 index 0000000..9718134 --- /dev/null +++ b/engine/rr/rr_templates.py @@ -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(' bytes: - """One zstd frame, through the zstandard module when it is installed and the zstd tool otherwise.""" + """One zstd frame, through Python's own zstd (3.14), the zstandard module or the zstd tool.""" + try: + from compression import zstd # Python 3.14 + return zstd.ZstdDecompressor().decompress(src) + except ImportError: + pass try: import zstandard return zstandard.ZstdDecompressor().decompressobj().decompress(src)