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:
1 parent
e3606cedf9
commit
5f520a186d
7 files changed
+483
-20
No files matched your search
+37
-5
@@ -1,7 +1,8 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
"""Build the engine's Ray Reconstruction model directory (preset D or E) from a capture of the denoiser.
|
"""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
|
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:
|
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).
|
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
|
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).
|
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 argparse
|
||||||
import json
|
import json
|
||||||
@@ -63,12 +70,26 @@ def find_module(ptx_dir, kernel):
|
|||||||
|
|
||||||
def main():
|
def main():
|
||||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
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('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('--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('--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')
|
ap.add_argument('--reset-only', action='store_true', help='use the own output shader that has no history path')
|
||||||
a = ap.parse_args()
|
a = ap.parse_args()
|
||||||
|
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)
|
cap = Capture(a.capture)
|
||||||
preset_d = any('diamond_wallaby' in l.kernel for l in cap.launches)
|
preset_d = any('diamond_wallaby' in l.kernel for l in cap.launches)
|
||||||
kernels = dict(KERNEL, **(D_KERNEL if preset_d else {}))
|
kernels = dict(KERNEL, **(D_KERNEL if preset_d else {}))
|
||||||
@@ -76,13 +97,22 @@ def main():
|
|||||||
lines = ['D4RR0001', f'preset {4 if preset_d else 5}']
|
lines = ['D4RR0001', f'preset {4 if preset_d else 5}']
|
||||||
info = {'preset': 'D' if preset_d else 'E', 'stages': {}}
|
info = {'preset': 'D' if preset_d else 'E', 'stages': {}}
|
||||||
for stage in STAGES:
|
for stage in STAGES:
|
||||||
|
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]]
|
launches = [l for l in cap.launches if l.kernel == kernels[stage]]
|
||||||
if len(launches) < a.frame:
|
if len(launches) < a.frame:
|
||||||
raise SystemExit(f'the capture has {len(launches)} launches of {kernels[stage]}; frame {a.frame} needed')
|
raise SystemExit(f'the capture has {len(launches)} launches of {kernels[stage]}; frame {a.frame} needed')
|
||||||
launch = launches[a.frame - 1]
|
launch = launches[a.frame - 1]
|
||||||
(a.output / f'weights-{stage}.bin').write_bytes(launch.weights)
|
weights = launch.weights
|
||||||
if stage in ('enc0', 'hkpn'):
|
if stage in ('enc0', 'hkpn'):
|
||||||
(a.output / f'{stage}.args').write_bytes(stage_params(launch))
|
(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')
|
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'
|
spv = a.output / f'{stage}.spv'
|
||||||
if own.exists() and (stage not in RESET_ONLY or a.reset_only):
|
if own.exists() and (stage not in RESET_ONLY or a.reset_only):
|
||||||
@@ -94,6 +124,7 @@ def main():
|
|||||||
formats = [f'--surface={off}={fmt}' for off, fmt in SURFACE_FORMATS.get(stage, {}).items()]
|
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])),
|
subprocess.run([sys.executable, str(HERE / 'rr' / 'ptx2rr.py'), str(find_module(a.ptx_dir, kernels[stage])),
|
||||||
kernels[stage], str(generated), *formats], check=True)
|
kernels[stage], str(generated), *formats], check=True)
|
||||||
|
if not (network and a.network_bindings_only):
|
||||||
glsl(generated, spv)
|
glsl(generated, spv)
|
||||||
origin = 'generated from the DLL'
|
origin = 'generated from the DLL'
|
||||||
binding = json.loads(generated.with_suffix('.json').read_text())
|
binding = json.loads(generated.with_suffix('.json').read_text())
|
||||||
@@ -101,8 +132,9 @@ def main():
|
|||||||
tex = ','.join(map(str, textures)) or '-'
|
tex = ','.join(map(str, textures)) or '-'
|
||||||
surf = ','.join(f'{off}:{fmt}' for off, fmt in surfaces) 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 ''))
|
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)}
|
n = len(weights) if weights is not None else 0
|
||||||
print(f'{stage}: {kernels[stage]} ({origin}), {len(launch.weights)} weight bytes')
|
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']
|
lines += ['template enc0 enc0.args', 'template hkpn hkpn.args']
|
||||||
for p in (0, 1):
|
for p in (0, 1):
|
||||||
glsl(HERE / 'exposure_k.comp', a.output / f'exposure_rr{p}.spv', '-DMODE=2', f'-DPASS={p}')
|
glsl(HERE / 'exposure_k.comp', a.output / f'exposure_rr{p}.spv', '-DMODE=2', f'-DPASS={p}')
|
||||||
|
|||||||
+40
-7
@@ -3,6 +3,7 @@
|
|||||||
nvngx_dlssd.dll 310.9) from a capture of the denoiser.
|
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]
|
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
|
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:
|
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
|
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
|
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.
|
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.
|
Everything but the exposure shader is derived from NVIDIA's DLL: it stays on the user's machine and is never packaged.
|
||||||
"""
|
"""
|
||||||
import argparse
|
import argparse
|
||||||
@@ -88,28 +92,49 @@ def find_module(ptx_dir, kernel):
|
|||||||
|
|
||||||
def main():
|
def main():
|
||||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
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('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('--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('--frame', type=int, default=2, help='which captured frame supplies the templates (1-based)')
|
||||||
ap.add_argument('--particles', metavar='CAPTURE_DIR',
|
ap.add_argument('--particles', metavar='CAPTURE_DIR',
|
||||||
help='a capture taken with DLSSD.ColorBeforeParticles: adds the downsample kernel NGX runs then (downp)')
|
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')
|
ap.add_argument('--native', metavar='DLL', help='nvngx_dlssd.dll 310.9: build the native stages from its weights')
|
||||||
a = ap.parse_args()
|
a = ap.parse_args()
|
||||||
|
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)
|
cap = Capture(a.capture)
|
||||||
out = a.output
|
out = a.output
|
||||||
out.mkdir(parents=True, exist_ok=True)
|
out.mkdir(parents=True, exist_ok=True)
|
||||||
capped = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0']
|
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']
|
lines = ['D4RR0002', 'preset 6']
|
||||||
pcap = Capture(a.particles) if a.particles else None
|
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]
|
kernel = KERNEL[stage]
|
||||||
|
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':
|
if stage == 'downp':
|
||||||
cap = pcap
|
cap = pcap
|
||||||
launches = [l for l in cap.launches if l.kernel == kernel]
|
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']])
|
per_frame = len(launches) // len([l for l in cap.launches if l.kernel == KERNEL['enc0']])
|
||||||
launch = launches[(a.frame - 1) * per_frame]
|
launch = launches[(a.frame - 1) * per_frame]
|
||||||
network = stage in NETWORK
|
|
||||||
if a.native and stage in NATIVE:
|
if a.native and stage in NATIVE:
|
||||||
work = out / 'native-work' / stage
|
work = out / 'native-work' / stage
|
||||||
work.mkdir(parents=True, exist_ok=True)
|
work.mkdir(parents=True, exist_ok=True)
|
||||||
@@ -123,6 +148,10 @@ def main():
|
|||||||
if RELAXED < 2: # (the relaxed stages leave their conversions to the driver)
|
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')],
|
subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}_m.spv')],
|
||||||
check=True, stdout=subprocess.DEVNULL)
|
check=True, stdout=subprocess.DEVNULL)
|
||||||
|
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)])
|
(out / f'{stage}.args').write_bytes(stage_params(launch, 24)[:len(launch.args)])
|
||||||
arena = launch.arena_alloc
|
arena = launch.arena_alloc
|
||||||
unit = launch.allocs[arena][1] / 3152
|
unit = launch.allocs[arena][1] / 3152
|
||||||
@@ -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
|
# 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
|
# (three ranges) is six words appended to the argument block
|
||||||
split = stage == 'enc0' and SPLIT and GROUP == 1
|
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:
|
if split:
|
||||||
subprocess.run([sys.executable, str(HERE / 'rr' / 'gen_rr2_enc0.py'), str(find_module(a.ptx_dir, kernel)),
|
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:]],
|
str(generated), str(work / 'enc0_g.comp'), '--split', 'gather', *how[2:]],
|
||||||
@@ -183,12 +213,15 @@ def main():
|
|||||||
weights = f'weights-{stage}.bin'
|
weights = f'weights-{stage}.bin'
|
||||||
(out / weights).write_bytes(generated.with_suffix('.globals.bin').read_bytes())
|
(out / weights).write_bytes(generated.with_suffix('.globals.bin').read_bytes())
|
||||||
if stage == 'enc0':
|
if stage == 'enc0':
|
||||||
(out / 'enc0-aux.bin').write_bytes(launch.pre(launch.pointer(104).alloc))
|
(out / 'enc0-aux.bin').write_bytes(aux if launch is None else launch.pre(launch.pointer(104).alloc))
|
||||||
args = stage_params(launch, 24 if network else 9999)[:len(launch.args)]
|
args = blocks[stage] if launch is None else stage_params(launch, 24 if network else 9999)[:len(launch.args)]
|
||||||
if hybrid and split:
|
if hybrid and split:
|
||||||
args = args + bytes(-len(args) % 8) + bytes(24) # the gather records' three ranges (the runtime's)
|
args = args + bytes(-len(args) % 8) + bytes(24) # the gather records' three ranges (the runtime's)
|
||||||
(out / f'{stage}.args').write_bytes(args)
|
(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)
|
# the stage's arena regions, in units of 1/3152 of NGX's arena (W0 * H0 / 32 bytes: the regions scale with it)
|
||||||
|
if launch is None:
|
||||||
|
ptrs = [ARENA[stage]] if network else []
|
||||||
|
else:
|
||||||
arena = launch.arena_alloc
|
arena = launch.arena_alloc
|
||||||
unit = launch.allocs[arena][1] / 3152 if arena is not None else 1
|
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
|
ptrs = [f'{p.offset}:{p.delta / unit:g}' for p in launch.pointers
|
||||||
@@ -205,7 +238,7 @@ def main():
|
|||||||
(f' x{GROUP}' if hybrid and GROUP > 1 else '')) # windows per workgroup (engine/rr/gen_rr2_enc0.py)
|
(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
|
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')
|
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):
|
for p in (0, 1):
|
||||||
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', f'-DPASS={p}',
|
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')],
|
str(HERE / 'rr' / 'exposure_f.comp'), '-o', str(out / f'exposure_f{p}.spv')],
|
||||||
|
|||||||
@@ -5,6 +5,11 @@ set -euo pipefail
|
|||||||
# for the preset the capture was taken with (D or E).
|
# for the preset the capture was taken with (D or E).
|
||||||
#
|
#
|
||||||
# usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL]
|
# 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
|
# 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:
|
# 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).
|
# 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)"
|
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||||
CAP="$(realpath "${1:?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"
|
||||||
OUT="$(realpath -m "${2:?usage: build_native_model.sh CAPTURE_DIR 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"
|
RR="$ROOT/engine/rr"
|
||||||
WORK="$OUT/native-work"
|
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; }
|
[[ -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}"
|
FRAME="${D4R_RR_MODEL_FRAME:-2}"
|
||||||
|
# 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)
|
CAPPED=(systemd-run --user --scope -q -p MemoryMax=24G -p MemorySwapMax=0)
|
||||||
|
fi
|
||||||
mkdir -p "$OUT" "$WORK"
|
mkdir -p "$OUT" "$WORK"
|
||||||
cd "$ROOT"
|
cd "$ROOT"
|
||||||
|
|
||||||
|
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 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}
|
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")"
|
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
|
# 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
|
# 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
|
b=1,1; case $s in enc1|enc2|dec1|dec2) b=5,5 ;; esac
|
||||||
d="$WORK/$s"; mkdir -p "$d"
|
d="$WORK/$s"; mkdir -p "$d"
|
||||||
cp "$WORK/dll/$s.npz" "$d/weights.npz"
|
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
|
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"
|
glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR"
|
||||||
native "$s" "weights-${s}m.bin"
|
native "$s" "weights-${s}m.bin"
|
||||||
@@ -77,7 +101,7 @@ done
|
|||||||
# decoder 0
|
# decoder 0
|
||||||
d="$WORK/dec0"; mkdir -p "$d"
|
d="$WORK/dec0"; mkdir -p "$d"
|
||||||
cp "$WORK/dll/dec0.npz" "$d/weights.npz"
|
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
|
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"
|
glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR"
|
||||||
native dec0 weights-dec0m.bin
|
native dec0 weights-dec0m.bin
|
||||||
|
|||||||
@@ -133,8 +133,11 @@ def noise_table(dll):
|
|||||||
if b - a >= 2048:
|
if b - a >= 2048:
|
||||||
t = u[b - 2048:b]
|
t = u[b - 2048:b]
|
||||||
if len(np.unique(t)) > 1024 and (t & 31 == 31).sum() > 256:
|
if len(np.unique(t)) > 1024 and (t & 31 == 31).sum() > 256:
|
||||||
found.append(t.tobytes())
|
found.append((2 * (b - 2048), t.tobytes()))
|
||||||
return found[0] if len(set(found)) == 1 else None # some DLLs hold the same table twice
|
# 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:
|
class Reader:
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -52,7 +52,12 @@ ZSTD_MAGIC = b'\x28\xb5\x2f\xfd'
|
|||||||
|
|
||||||
|
|
||||||
def zstd_frame(src: bytes) -> bytes:
|
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:
|
try:
|
||||||
import zstandard
|
import zstandard
|
||||||
return zstandard.ZstdDecompressor().decompressobj().decompress(src)
|
return zstandard.ZstdDecompressor().decompressobj().decompress(src)
|
||||||
|
|||||||
Reference in new issue
Block a user