engine/rr: native Ray Reconstruction stages with capture/replay and shim integration
- 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
This commit is contained in:
woutermans committed
1 parent
30472ae195
commit
50fa65d87b
67 files changed
+11926
-504
No files matched your search
Executable
+207
@@ -0,0 +1,207 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run one Ray Reconstruction stage shader on a captured launch and compare it with NVIDIA's result.
|
||||
|
||||
usage: test_layer.py CAPTURE KERNEL FRAME SHADER.spv WORK [--kvk build/rr/kvk] [--time N] [--subgroup 32|64]
|
||||
|
||||
KERNEL is a substring of the denoiser kernel's name (e.g. swin_enc1, hkpn_output_kernel_transformer), FRAME the
|
||||
1-based occurrence in the capture. The shader follows the RR stage contract (engine/rr/README.md):
|
||||
|
||||
binding 0 the parameter block: NVIDIA's argument block with every pointer into the plane arena
|
||||
replaced by its byte offset in the arena, the weight pointer and every texture or surface
|
||||
handle by zero (stage_params below; the runtime builds the same block)
|
||||
binding 1 the weight image exactly as NGX uploads it (argument 40), read-only
|
||||
binding 2 the plane arena, read/write
|
||||
bindings 3, 4, ... the textures in ascending argument offset, then the surfaces in ascending argument offset
|
||||
|
||||
Every surface starts zeroed and the arena starts as it was before NVIDIA's launch; afterwards each surface and the
|
||||
whole arena are compared with what NVIDIA's launch left. Exit status 0 only when everything is byte-identical.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from rr_capture import Capture # noqa: E402
|
||||
|
||||
|
||||
# GLSL image formats of the stage shaders' surfaces -> (channels, kvk bits)
|
||||
SURFACE_FORMATS = {'rgba16f': (4, 16), 'rg16f': (2, 16), 'r16f': (1, 16), 'rgba8_snorm': (4, 9), 'rgba8': (4, 8)}
|
||||
|
||||
|
||||
def stage_params(launch):
|
||||
args = bytearray(launch.args)
|
||||
arena = launch.arena_alloc
|
||||
for p in launch.pointers:
|
||||
value = p.delta if p.alloc == arena and p.offset != 40 else 0
|
||||
struct.pack_into('<Q', args, p.offset, value)
|
||||
# texture and surface handles: 1 where NGX bound one (the kernels test optional inputs' handles for null, e.g.
|
||||
# enc0's specular hit distance at argument 248), 0 elsewhere
|
||||
for image in launch.images:
|
||||
struct.pack_into('<Q', args, image.offset, 1)
|
||||
return bytes(args)
|
||||
|
||||
|
||||
def e4m3_table():
|
||||
t = np.zeros(256)
|
||||
for c in range(256):
|
||||
s, e, m = c >> 7, (c >> 3) & 15, c & 7
|
||||
v = np.nan if (e == 15 and m == 7) else (m * 2.0 ** -9 if e == 0 else (1 + m / 8) * 2.0 ** (e - 7))
|
||||
t[c] = -v if s else v
|
||||
return t
|
||||
|
||||
|
||||
def psnr(a, b):
|
||||
finite = np.isfinite(a) & np.isfinite(b)
|
||||
if not finite.any():
|
||||
return float('nan')
|
||||
err = np.mean((a[finite] - b[finite]) ** 2)
|
||||
peak = max(np.max(np.abs(b[finite])), 1e-6)
|
||||
return float('inf') if err == 0 else 10 * np.log10(peak * peak / err)
|
||||
|
||||
|
||||
def compare_image(name, image, got, ref, ignore_channels=()):
|
||||
"""IGNORE_CHANNELS: channels NVIDIA's kernel fills from an undefined register (left out of the comparison)"""
|
||||
if image.bits == 16:
|
||||
a, b = np.frombuffer(got, np.float16).astype(np.float64), np.frombuffer(ref, np.float16).astype(np.float64)
|
||||
elif image.bits == 32:
|
||||
a, b = np.frombuffer(got, np.float32).astype(np.float64), np.frombuffer(ref, np.float32).astype(np.float64)
|
||||
else:
|
||||
dt = np.int8 if image.bits == 9 else np.uint8
|
||||
a, b = np.frombuffer(got, dt).astype(np.float64), np.frombuffer(ref, dt).astype(np.float64)
|
||||
keep = np.ones(len(a), bool)
|
||||
for c in ignore_channels:
|
||||
keep[c::image.channels] = False
|
||||
a, b = a[keep], b[keep]
|
||||
element_bytes = len(got) // len(keep)
|
||||
same = (np.frombuffer(got, np.uint8).reshape(-1, element_bytes) == np.frombuffer(ref, np.uint8).reshape(-1, element_bytes))[keep]
|
||||
elements = len(a)
|
||||
equal = np.sum(a == b) + np.sum(np.isnan(a) & np.isnan(b))
|
||||
worst = np.nanmax(np.abs(a - b)) if elements else 0
|
||||
print(f' {name}: {image.width}x{image.height}x{image.channels} fmt{image.format}'
|
||||
f'{" (channels " + ",".join(map(str, ignore_channels)) + " ignored)" if ignore_channels else ""}: '
|
||||
f'{100.0 * equal / elements:.4f}% identical values, max |diff| {worst:.6g}, PSNR {psnr(a, b):.2f} dB, '
|
||||
f'nonfinite ours/ref {int(np.sum(~np.isfinite(a)))}/{int(np.sum(~np.isfinite(b)))}')
|
||||
return bool(same.all())
|
||||
|
||||
|
||||
def compare_arena(got, before, ref):
|
||||
g, p, r = (np.frombuffer(x, np.uint8) for x in (got, before, ref))
|
||||
written = (r != p) | (g != p)
|
||||
differ = g != r
|
||||
print(f' arena: {int(written.sum())} bytes written by either side, {int(differ.sum())} differ', end='')
|
||||
if differ.any():
|
||||
idx = np.nonzero(differ)[0]
|
||||
t = e4m3_table()
|
||||
dv = np.abs(t[g[idx]] - t[r[idx]])
|
||||
print(f' (first at byte {idx[0]:#x}, last at {idx[-1]:#x}; as e4m3 values: max |diff| {np.nanmax(dv):.6g}, '
|
||||
f'{int(np.sum(dv <= 0.0625 * np.abs(t[r[idx]]) + 1e-12))} within one e4m3 step)')
|
||||
else:
|
||||
print()
|
||||
return not differ.any()
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument('capture')
|
||||
ap.add_argument('kernel')
|
||||
ap.add_argument('frame', type=int)
|
||||
ap.add_argument('shader')
|
||||
ap.add_argument('work')
|
||||
ap.add_argument('--kvk', default=str(Path(__file__).resolve().parents[2] / 'build/rr/kvk'))
|
||||
ap.add_argument('--time', type=int, default=1)
|
||||
ap.add_argument('--subgroup', type=int, default=32)
|
||||
ap.add_argument('--weights', help='bind this file at binding 1 instead of NGX\'s weight image (native stages)')
|
||||
ap.add_argument('--grid', help='X,Y workgroups (default: the launch grid)')
|
||||
ap.add_argument('--ignore', action='append', default=[], metavar='SURFACE:CHANNEL',
|
||||
help='leave a channel out of the comparison, e.g. 112:3 for a value NVIDIA stores from an '
|
||||
'undefined register (repeatable)')
|
||||
o = ap.parse_args()
|
||||
|
||||
cap = Capture(o.capture)
|
||||
launch = [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)
|
||||
print(f'{launch.kernel} (launch {launch.sequence}, grid {launch.grid}, block {launch.block})')
|
||||
|
||||
(work / 'params.bin').write_bytes(stage_params(launch))
|
||||
(work / 'weights.bin').write_bytes(Path(o.weights).read_bytes() if o.weights else launch.weights)
|
||||
arena = launch.arena_alloc
|
||||
arena_before = launch.pre(arena)
|
||||
(work / 'arena.bin').write_bytes(arena_before)
|
||||
gx, gy = (int(v) for v in o.grid.split(',')) if o.grid else launch.grid[:2]
|
||||
lines = [f'param {work}/params.bin', f'grid {gx} {gy}']
|
||||
# Bind in the shader's order (its ptx2rr.py sidecar, SHADER.json: textures then surfaces by argument offset).
|
||||
# A handle the shader reads but NGX left null gets a zeroed 1x1 image.
|
||||
sidecar = Path(o.shader).with_suffix('.json')
|
||||
if sidecar.exists():
|
||||
info = json.loads(sidecar.read_text())
|
||||
order = [('texture', off, None) for off in info['textures']] + [('surface', off, f) for off, f in info['surfaces']]
|
||||
else:
|
||||
order = [(i.kind, i.offset, None) for i in launch.images]
|
||||
by_key = {(i.kind, i.offset): i for i in launch.images}
|
||||
(work / 'zero.bin').write_bytes(bytes(64))
|
||||
for kind, offset, fmt in order:
|
||||
image = by_key.get((kind, offset))
|
||||
if image is None:
|
||||
ch, bits = SURFACE_FORMATS.get(fmt, (4, 16))
|
||||
lines.append(f'tex {offset} {work}/zero.bin 1 1 4 16 nearest' if kind == 'texture' else
|
||||
f'surf {offset} 1 1 {ch} {bits} - {work}/dummy{offset}.bin')
|
||||
elif kind == 'texture':
|
||||
lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} {image.bits} '
|
||||
f'{"linear" if image.linear else "nearest"}')
|
||||
else:
|
||||
lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {image.bits} - '
|
||||
f'{work}/surf{image.offset}.bin')
|
||||
n_images = len(order)
|
||||
lines.append(f'lut {work}/weights.bin')
|
||||
lines.append(f'inout {work}/arena.bin {work}/arena.out.bin')
|
||||
probes = info.get('probes', []) if sidecar.exists() else []
|
||||
oob = info.get('oob', False) if sidecar.exists() else False
|
||||
if probes: # ptx2rr.py --probe: the debug buffer after the surfaces
|
||||
threads = launch.block[0] * launch.block[1] * launch.block[2]
|
||||
lines.append(f'outbuf {work}/debug.bin {4 * threads * len(probes)}')
|
||||
elif oob: # ptx2rr.py --oob-report: count, lowest and highest address per buffer
|
||||
(work / 'oob.init').write_bytes(struct.pack('<9I', *([0, 0xffffffff, 0] * 3)))
|
||||
lines.append(f'inout {work}/oob.init {work}/oob.bin')
|
||||
lines.append('bindings 0 ' + ' '.join(str(3 + i) for i in range(n_images)) + ' 1 2' +
|
||||
(f' {3 + n_images}' if probes or oob else ''))
|
||||
(work / 'manifest').write_text('\n'.join(lines) + '\n')
|
||||
r = subprocess.run([o.kvk, o.shader, str(work / 'manifest'), str(o.time), str(o.subgroup)],
|
||||
capture_output=True, text=True)
|
||||
sys.stdout.write(r.stdout)
|
||||
if r.returncode != 0:
|
||||
sys.stderr.write(r.stderr)
|
||||
sys.exit(f'kvk failed ({r.returncode})')
|
||||
|
||||
identical = True
|
||||
bound = {(kind, offset) for kind, offset, _ in order}
|
||||
for image in launch.images:
|
||||
if image.kind != 'surface':
|
||||
continue
|
||||
if ('surface', image.offset) not in bound:
|
||||
print(f' surface {image.offset}: not written by this kernel')
|
||||
continue
|
||||
got = (work / f'surf{image.offset}.bin').read_bytes()
|
||||
waived = [int(s.split(':')[1]) for s in o.ignore if int(s.split(':')[0]) == image.offset]
|
||||
identical &= compare_image(f'surface {image.offset}', image, got, image.path.read_bytes(), waived)
|
||||
if arena is not None:
|
||||
identical &= compare_arena((work / 'arena.out.bin').read_bytes(), arena_before,
|
||||
launch.post_buffer(cap, arena))
|
||||
if oob:
|
||||
r = struct.unpack('<9I', (work / 'oob.bin').read_bytes()[:36])
|
||||
for k, name in enumerate(('weight reads', 'arena reads', 'arena writes')):
|
||||
print(f' out of range {name}: {r[3 * k]}' + (f' (bytes {r[3 * k + 1]:#x}..{r[3 * k + 2]:#x}, size '
|
||||
f'{len(launch.weights) if k == 0 else len(arena_before):#x})' if r[3 * k] else ''))
|
||||
print('IDENTICAL' if identical else 'DIFFERENT')
|
||||
sys.exit(0 if identical else 1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in new issue
Block a user