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.

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-<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}')

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')],

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

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:

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
+(<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()

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('<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()

kernels/tools/extract_dlss_ptx.py

@@ -52,7 +52,12 @@ ZSTD_MAGIC = b'\x28\xb5\x2f\xfd'
 
 
 def zstd_frame(src: bytes) -> 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)