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.
96 lines
5.3 KiB
Python
96 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare a full-frame native K replay from a capture. Does not translate kernels.
|
|
|
|
prepare_k.py CAPTURE_DIR OUTPUT_DIR [--frames N] [--start FRAME]
|
|
prepare_k.py CAPTURE_DIR OUTPUT_DIR --lut-only
|
|
The model's reconstruction tables can be packaged using OUTPUT_DIR/lut.bin. --lut-only writes only that, from the
|
|
first output launch's replay (a capture without the launch dumps: install.sh's weight dump).
|
|
"""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import struct
|
|
from capture import INPUT, OUTPUT, ORDER, launches, input_params, output_params
|
|
|
|
|
|
def quote(path):
|
|
return json.dumps(str(Path(path).resolve()))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('capture', type=Path)
|
|
parser.add_argument('output', type=Path)
|
|
parser.add_argument('--frames', type=int, default=0, help='0 = all captured frames')
|
|
parser.add_argument('--start', type=int, default=0)
|
|
parser.add_argument('--lut-only', action='store_true', help='write only lut.bin')
|
|
a = parser.parse_args()
|
|
if a.lut_only:
|
|
outputs = launches(a.capture, OUTPUT)
|
|
if not outputs:
|
|
raise ValueError('capture holds no K output launch')
|
|
a.output.mkdir(parents=True, exist_ok=True)
|
|
(a.output / 'lut.bin').write_bytes(outputs[0].allocation(280, 16 * 32768))
|
|
return
|
|
inputs, outputs = launches(a.capture, INPUT), launches(a.capture, OUTPUT)
|
|
if not inputs or not outputs or len(outputs) > len(inputs):
|
|
raise ValueError('capture must contain complete K input/output frame pairs')
|
|
pairs = list(zip(inputs, outputs))
|
|
for i, (ip, op) in enumerate(pairs):
|
|
if not ip.sequence < op.sequence or (i+1 < len(inputs) and op.sequence >= inputs[i+1].sequence):
|
|
raise ValueError('capture has a missing or out-of-order input/output launch')
|
|
if a.start < 0 or a.start >= len(pairs) or a.frames < 0:
|
|
raise ValueError('invalid frame selection')
|
|
end = len(pairs) if a.frames == 0 else min(len(pairs), a.start + a.frames)
|
|
selected = pairs[a.start:end]
|
|
# Unselected later frames may be incomplete (for example a capture stopped while writing them).
|
|
for ip, op in selected:
|
|
for resource in (*ip.resources.values(), *op.resources.values()): resource.validate()
|
|
layer_launches = {name: launches(a.capture, 'dltss_pwin_'+name+'_layer') for name in ORDER}
|
|
first_in, first_out = selected[0]
|
|
geometry = []
|
|
for name in ORDER:
|
|
matches = [p for p in layer_launches[name] if first_in.sequence < p.sequence < first_out.sequence]
|
|
if len(matches) != 1:
|
|
raise ValueError(f'frame {a.start}: expected one {name} launch, found {len(matches)}')
|
|
geometry.append((*matches[0].i2(0), *matches[0].i2(56)))
|
|
final = first_out.resources[312]
|
|
a.output.mkdir(parents=True, exist_ok=True)
|
|
(a.output / 'shape.bin').write_bytes(struct.pack('<2I', final.width, final.height) + b''.join(struct.pack('<4i', *s) for s in geometry))
|
|
# NGX's whole allocation: sixteen reconstruction tables, of which a launch uses the one for its upscaling ratio
|
|
(a.output / 'lut.bin').write_bytes(first_out.allocation(280, 16 * 32768))
|
|
lines = [f'shape {quote(a.output / "shape.bin")}',
|
|
f'history {quote(first_in.resources[176].path)} {quote(first_out.resources[232].path)} {quote(first_in.resources[224].path)}']
|
|
resource_ids = {}
|
|
references = []
|
|
for number, (ip, op) in enumerate(selected):
|
|
if not ip.sequence < op.sequence or (number and ip.sequence <= selected[number-1][1].sequence):
|
|
raise ValueError('capture frames are not in network order')
|
|
if op.resources[312].width != final.width or op.resources[312].height != final.height:
|
|
raise ValueError('resolution changes require another Engine')
|
|
frame_geometry = []
|
|
for name, expected in zip(ORDER, geometry):
|
|
matches = [p for p in layer_launches[name] if ip.sequence < p.sequence < op.sequence]
|
|
if len(matches) != 1 or matches[0].i2(0) != expected[:2]:
|
|
raise ValueError(f'frame {number}: missing or changed {name} geometry')
|
|
frame_geometry.append((*matches[0].i2(0), *matches[0].i2(56)))
|
|
ids = []
|
|
for off in (168, 184, 192):
|
|
r = ip.resources[off]
|
|
if r.path not in resource_ids:
|
|
i = len(resource_ids); resource_ids[r.path] = i
|
|
lines.append(f'image {i} {quote(r.path)} {r.width} {r.height} {r.channels} {r.bits}')
|
|
ids.append(resource_ids[r.path])
|
|
inp, outp = a.output / f'frame-{number}.input', a.output / f'frame-{number}.output'
|
|
inp.write_bytes(input_params(ip)); outp.write_bytes(output_params(op))
|
|
windows = a.output / f'frame-{number}.windows'
|
|
windows.write_bytes(b''.join(struct.pack('<4i', *s) for s in frame_geometry))
|
|
lines.append(f'frame {quote(inp)} {quote(outp)} {quote(windows)} ' + ' '.join(map(str, ids)))
|
|
references.append({str(off): str(op.resources[off].path) for off in (288,296,304,312)})
|
|
(a.output / 'case.txt').write_text('\n'.join(lines) + '\n')
|
|
(a.output / 'references.json').write_text(json.dumps(references, indent=2) + '\n')
|
|
print(f'Prepared {len(selected)} frames, {final.width}x{final.height}, token grid {geometry[0][0]}x{geometry[0][1]}')
|
|
|
|
|
|
if __name__ == '__main__': main()
|