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.
105 lines
5.9 KiB
Python
105 lines
5.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Replay one captured denoiser launch through ZLUDA on Linux (tools/replay_rr_launch.cpp) and compare it with the
|
|
capture, like test_layer.py does for a stage shader.
|
|
|
|
usage: zluda_replay.py CAPTURE KERNEL FRAME WORK [--ptx MODULE] [--zluda DIR]
|
|
|
|
KERNEL and FRAME select the launch as in test_layer.py. MODULE defaults to the corpus module holding the kernel
|
|
(~/.cache/d4r-rr-corpus); pass an edited copy to add debug stores. ZLUDA's own settings (D4R_ZLUDA_WMMA,
|
|
D4R_ZLUDA_WMMA_FP8, D4R_ZLUDA_WMMA_FP8_NATIVE, ...) are taken from the environment. Outputs land in WORK
|
|
(surfOFFSET.bin, arena.out.bin) so the two replays can be compared with each other too.
|
|
"""
|
|
import argparse
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from rr_capture import Capture # noqa: E402
|
|
from test_layer import compare_arena, compare_image # noqa: E402
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('capture')
|
|
ap.add_argument('kernel')
|
|
ap.add_argument('frame', type=int)
|
|
ap.add_argument('work')
|
|
ap.add_argument('--ptx')
|
|
ap.add_argument('--zluda', default=os.path.expanduser('~/.cache/d4r-zluda-rr'))
|
|
ap.add_argument('--rocm', default='/opt/rocm/lib')
|
|
ap.add_argument('--tex', action='append', default=[], metavar='OFFSET=FILE:W:H',
|
|
help='an RGBA32F texture for a handle the capture has no dump of (preset F pooled histories)')
|
|
ap.add_argument('--pad', type=int, default=0, help='bytes of zeros appended to the arena (debug stores)')
|
|
o = ap.parse_args()
|
|
cap = Capture(o.capture)
|
|
exact = [l for l in cap.launches if l.kernel == o.kernel]
|
|
launch = (exact or [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or
|
|
('diamond' in o.kernel and o.kernel in l.kernel)])[o.frame - 1]
|
|
work = Path(o.work)
|
|
work.mkdir(parents=True, exist_ok=True)
|
|
corpus = '~/.cache/d4r-rr-corpus-f' if launch.kernel.startswith('rr2_') else '~/.cache/d4r-rr-corpus'
|
|
module = o.ptx or next(str(p) for p in sorted(Path(corpus).expanduser().glob('*.ptx'))
|
|
if f'.entry {launch.kernel}(' in p.read_text(errors='replace'))
|
|
(work / 'args.bin').write_bytes(launch.args)
|
|
arena = launch.arena_alloc
|
|
lines = [f'args {work}/args.bin', 'grid ' + ' '.join(map(str, launch.grid)),
|
|
'block ' + ' '.join(map(str, launch.block))]
|
|
if o.pad:
|
|
(work / 'arena-pad.bin').write_bytes(launch.pre(arena) + bytes(o.pad))
|
|
for p in launch.pointers:
|
|
out = f' {work}/arena.out.bin' if p.alloc == arena else ''
|
|
source = f'{work}/arena-pad.bin' if o.pad and p.alloc == arena else f'{launch.directory}/alloc-{p.alloc}.bin'
|
|
lines.append(f'buf {p.offset} {source} {p.delta}{out}')
|
|
# textures the launch names but the bridge could not read back (preset F's 10:10:10:2 pooled histories, which
|
|
# ZLUDA backs with RGBA32F): zeros, which is what they hold on a first frame
|
|
have = {i.offset for i in launch.images}
|
|
for line in (launch.directory / 'manifest.txt').read_text().split('\n'):
|
|
f = line.split()
|
|
given = {int(t.split('=')[0]): t.split('=')[1].split(':') for t in o.tex}
|
|
if len(f) == 2 and f[0] == 'texture' and int(f[1]) in given:
|
|
path, w_, h_ = given[int(f[1])]
|
|
lines.append(f'tex {f[1]} {path} {w_} {h_} 4 32 linear')
|
|
elif len(f) == 2 and f[0] == 'texture' and int(f[1]) not in have:
|
|
(work / 'zero16.bin').write_bytes(bytes(16 * 16 * 16))
|
|
lines.append(f'tex {f[1]} {work}/zero16.bin 16 16 4 32 linear')
|
|
if len(f) == 2 and f[0] == 'surface' and int(f[1]) not in have: # half the largest surface: the pooled history
|
|
big = max((i for i in launch.images if i.kind == 'surface'), key=lambda i: i.width, default=None)
|
|
if big is not None:
|
|
lines.append(f'surf {f[1]} {big.width // 2} {big.height // 2} 4 32 {work}/surf{f[1]}.bin')
|
|
for image in launch.images:
|
|
if image.kind == 'texture':
|
|
# SNORM_INT8X4 (200): a SIGNED_INT8 array read as normalized floats gives the same values
|
|
fmt = 8 if image.format == 200 else 1 if image.format == 194 else image.format
|
|
lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} '
|
|
f'{fmt} {"linear" if image.linear else "nearest"}')
|
|
else:
|
|
# ZLUDA has no SNORM_INT8X4 arrays (format 200); surface stores write the same bytes to a SIGNED_INT8 one
|
|
fmt = 8 if image.format == 200 else 1 if image.format == 194 else image.format # 194: UNORM_INT8X4
|
|
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {fmt} '
|
|
f'{work}/surf{image.offset}.bin')
|
|
(work / 'manifest').write_text('\n'.join(lines) + '\n')
|
|
env = dict(os.environ, LD_LIBRARY_PATH=f'{o.zluda}:{o.rocm}:' + os.environ.get('LD_LIBRARY_PATH', ''))
|
|
print(f'{launch.kernel} (launch {launch.sequence}) through ZLUDA {o.zluda}')
|
|
r = subprocess.run([str(ROOT / 'build/replay_rr_launch'), module, launch.kernel, str(work / 'manifest')],
|
|
env=env, capture_output=True, text=True)
|
|
sys.stdout.write(r.stdout)
|
|
if r.returncode:
|
|
sys.stderr.write(r.stderr[-4000:])
|
|
sys.exit('replay failed')
|
|
same = True
|
|
for image in launch.images:
|
|
if image.kind == 'surface':
|
|
same &= compare_image(f'surface {image.offset}', image, (work / f'surf{image.offset}.bin').read_bytes(),
|
|
image.path.read_bytes())
|
|
if arena is not None and not o.pad:
|
|
same &= compare_arena((work / 'arena.out.bin').read_bytes(), launch.pre(arena), launch.post_buffer(cap, arena))
|
|
print('IDENTICAL' if same else 'DIFFERENT')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|