woutermans
50fa65d87b
- engine/rr: d4r's own shaders for the denoiser frame (exposure, six encoders,
five decoders, hkpn output stage) at NGX layouts, plus PTX->GLSL/HIP packing
scripts, reference implementations and weight extraction
- runtime.h/runtime_rr.inc, game_rr.{h,cpp}: GameDenoiser wired into the engine
build (d4r-rr replay binary) and the D3D12 shim's feature-13 path
- shim: answer GetFeatureRequirements from d4r's own state only (the old
cross-binary CUDA handoff crashed OptiScaler's scan); drop the core's
CUDA GetFeatureRequirements import
- shim/harness: exact guide capture/replay of a frame's RR inputs
- kernels/rr: expanded enc0/engine prototypes; kernels/tools: ptxsim rewrite
plus rrsim and an engine test; vulkan: ptx2glsl FP8/cooperative-matrix
support, kvk inout buffers and KVK_DISABLE_OPTIMIZATION
- bridge: report the surface array in cuSurfObjectCreate traces
- packaging: RR settings and the ICU launch overrides in D4R_README
111 lines
4.9 KiB
Python
111 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Turn a preset-E Ray Reconstruction model directory (engine/compile_rr.py) into the native one: every Swin stage
|
|
from d4r's own shaders instead of translated PTX.
|
|
|
|
usage: build_native.py CAPTURE MODEL_DIR [--frame N] [--work DIR]
|
|
|
|
enc0 engine/rr/enc0_m.comp, generated (gen_enc0_m.py) from the model's translated enc0.generated.comp,
|
|
whose per-pixel feature gather it keeps; weights: NGX's image + enc0_pack.py's blob
|
|
enc1-enc5, engine/rr/swin_m.comp with swin_pack.py's per-stage header and weights
|
|
dec4-dec1
|
|
dec0 engine/rr/dec0_m.comp with dec0_pack.py's weights
|
|
The weights are extracted from ZLUDA replays of the capture's launches (engine/rr/phase_dump.py: every mma's operands
|
|
of one window) and the stages' weight images; extraction checks itself against the dump. Like the model directory,
|
|
everything written is derived from the user's DLL and stays on the machine. Runs one replay or compile at a time
|
|
(phase_dump.py caps its replay at 24 GB).
|
|
"""
|
|
import argparse
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(HERE))
|
|
from rr_capture import Capture # noqa: E402
|
|
import swin_ref # noqa: E402
|
|
import dec0_ref # noqa: E402
|
|
import enc0_ref # noqa: E402
|
|
|
|
SWIN = ['enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4', 'dec3', 'dec2', 'dec1']
|
|
|
|
|
|
def run(*cmd, **kw):
|
|
subprocess.run([str(c) for c in cmd], check=True, **kw)
|
|
|
|
|
|
def glsl(source, output, *includes):
|
|
run('glslangValidator', '-V', '--target-env', 'vulkan1.3', *[f'-I{i}' for i in includes], source, '-o', output,
|
|
stdout=subprocess.DEVNULL)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('capture')
|
|
ap.add_argument('model')
|
|
ap.add_argument('--frame', type=int, default=2)
|
|
ap.add_argument('--work', help='dumps and build files (default MODEL/native-work)')
|
|
o = ap.parse_args()
|
|
model = Path(o.model)
|
|
work = Path(o.work) if o.work else model / 'native-work'
|
|
work.mkdir(parents=True, exist_ok=True)
|
|
cap = Capture(os.path.expanduser(o.capture))
|
|
manifest = (model / 'rr.txt').read_text().split('\n')
|
|
|
|
def dump(stage):
|
|
out = work / f'{stage}_all.npz'
|
|
launch = cap.find(f'swin_{stage}', o.frame)
|
|
block = '5,5' if launch.grid[0] > 5 and launch.grid[1] > 5 else '1,1'
|
|
if not out.exists():
|
|
run(sys.executable, HERE / 'phase_dump.py', o.capture, stage, o.frame, block, out)
|
|
return out, block
|
|
|
|
def stage_line(stage, shader, weights):
|
|
for i, line in enumerate(manifest):
|
|
f = line.split()
|
|
if len(f) > 4 and f[0] == 'stage' and f[1] == stage:
|
|
f[2], f[4] = shader, weights
|
|
manifest[i] = ' '.join(x for x in f if x != 'noopt')
|
|
|
|
for name in ('rr_exact.glsl', 'rr_attn.glsl', 'rr_decfilter.glsl', 'enc0_m_decl.glsl', 'enc0_m_tail.glsl'):
|
|
shutil.copy(HERE / name, work / name)
|
|
for stage in SWIN:
|
|
d, block = dump(stage)
|
|
print(f'{stage}: extracting', flush=True)
|
|
l, ex, ey, sx, sy, inp, outs, tex = swin_ref.load(o.capture, stage, o.frame)
|
|
bx, by = (int(v) for v in block.split(','))
|
|
swin_ref.extract(stage, d, l.weights, work / f'{stage}_w.npz', inp, 8 * bx - sx, 8 * by - sy, ex, ey)
|
|
sd = work / stage
|
|
sd.mkdir(exist_ok=True)
|
|
run(sys.executable, HERE / 'swin_pack.py', stage, work / f'{stage}_w.npz', model / f'weights-{stage}m.bin',
|
|
sd / 'swin_stage.h')
|
|
glsl(HERE / 'swin_m.comp', model / f'{stage}_m.spv', sd, work)
|
|
stage_line(stage, f'{stage}_m.spv', f'weights-{stage}m.bin')
|
|
# dec0
|
|
d, _ = dump('dec0')
|
|
print('dec0: extracting', flush=True)
|
|
dec0_ref.extract(d, cap.find('swin_dec0', o.frame).weights, work / 'dec0_w.npz')
|
|
run(sys.executable, HERE / 'dec0_pack.py', work / 'dec0_w.npz', model / 'weights-dec0m.bin', work / 'dec0_m_w.h')
|
|
glsl(HERE / 'dec0_m.comp', model / 'dec0_m.spv', work)
|
|
stage_line('dec0', 'dec0_m.spv', 'weights-dec0m.bin')
|
|
# enc0: NGX's weight image (the translated feature gather reads it) followed by the native blob
|
|
d, _ = dump('enc0')
|
|
print('enc0: extracting', flush=True)
|
|
nv = cap.find('swin_enc0', o.frame).weights
|
|
enc0_ref.extract(d, nv, work / 'enc0_w.npz')
|
|
base = (len(nv) + 255) // 256 * 256
|
|
run(sys.executable, HERE / 'enc0_pack.py', work / 'enc0_w.npz', work / 'enc0_native.bin', work / 'enc0_m_w.h', base)
|
|
(model / 'weights-enc0m.bin').write_bytes(nv + bytes(base - len(nv)) + (work / 'enc0_native.bin').read_bytes())
|
|
run(sys.executable, HERE / 'gen_enc0_m.py', model / 'enc0.generated.comp', work / 'enc0_m.comp')
|
|
glsl(work / 'enc0_m.comp', model / 'enc0_m.spv', work)
|
|
stage_line('enc0', 'enc0_m.spv', 'weights-enc0m.bin')
|
|
(model / 'rr.txt').write_text('\n'.join(manifest))
|
|
print(f'{model / "rr.txt"}: native stages written')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|