ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
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()
|