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
261 lines
12 KiB
Python
Executable File
261 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""rrsim: run selected thread blocks of a captured Ray Reconstruction CUDA launch on the CPU and check
|
|
every global and surface store against the GPU reference (the capture's post-launch arena and surfaces).
|
|
|
|
USAGE
|
|
python3 kernels/tools/rrsim.py CAPTURE KERNEL_SUBSTRING FRAME --blocks X,Y[,X,Y...]
|
|
[--watch 'LINE:%reg,...' ...] [--mma exact|split16] [--watch-max N]
|
|
|
|
CAPTURE capture directory, e.g. ~/.cache/d4r-rr-work/capE
|
|
KERNEL_SUBSTRING substring of the launch's kernel name, e.g. swin_enc1, swin_dec2, swin_enc0, swin_dec0, hkpn_output
|
|
FRAME 1-based index among the launches that match KERNEL_SUBSTRING (enc1 frame 2 = 2)
|
|
--blocks ctaid.x,ctaid.y pairs, e.g. 0,0,20,12
|
|
--watch print the named registers each time the statement at .ptx line LINE executes
|
|
(LINE is the line in the corpus module, as in grep -n). Repeatable.
|
|
--mma MMA accumulation model (see ptxsim.py module docstring): exact (default) or split16
|
|
--watch-max records printed per --watch (default 16)
|
|
|
|
For each block: runtime, statements executed, then one line per compared store group:
|
|
identical / differing counts, the ulp histogram of differing values, and max ulp.
|
|
|
|
LIBRARY
|
|
import sys; sys.path.insert(0, 'kernels/tools')
|
|
from rrsim import run_block, compare
|
|
r = run_block('~/.cache/d4r-rr-work/capE', 'swin_enc1', 2, (0, 0),
|
|
watch={1234: ['%f1']}) # watch: {ptx line: [registers]}, optional
|
|
r['regs'] final register file: name -> uint64 array (T lanes), bool for predicates
|
|
r['stores'] store log: list of (alloc_id, byte offsets per lane, raw bytes (lanes, bytes), ptx type)
|
|
r['surfaces'] [(Image, Surface)] surface writes of the block (Surface.data, Surface.written)
|
|
r['watch'] list of {line, stmt, reg, width, values, mask} records, in execution order
|
|
r['seconds'], r['steps'], r['block'] (ptxsim.Block), r['launch'], r['capture'], r['ptx']
|
|
compare(r) -> list of per-group dicts (n, identical, differing, max ulp, ulp histogram)
|
|
|
|
FIDELITY (details in ptxsim.py)
|
|
Exact: integer and bit ops, shfl/movmatrix/prmt, sust/tex.base point fetch, cvt, f16/f32 add/mul/sub
|
|
(one rounding per op). The MMA accumulation model and transcendental approximations may differ in the last
|
|
bit from the GPU reference; compare() reports the differences as ulp counts so they can be attributed.
|
|
Surfaces start at zero and textures are the post-launch dump (the capture has no pre-launch surface or
|
|
texture bytes). Only stores of the simulated blocks are compared; the arena is the replay's pre-launch
|
|
copy with only those blocks executed.
|
|
"""
|
|
import argparse
|
|
import os
|
|
import struct
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(HERE))
|
|
sys.path.insert(0, str(HERE.parents[1] / 'engine' / 'rr'))
|
|
import ptxsim # noqa: E402
|
|
from rr_capture import Capture # noqa: E402
|
|
|
|
U64 = np.uint64
|
|
CORPUS = Path.home() / '.cache' / 'd4r-rr-corpus'
|
|
NB = {'b8': 1, 'u8': 1, 's8': 1, 'b16': 2, 'u16': 2, 's16': 2, 'f16': 2,
|
|
'b32': 4, 'u32': 4, 's32': 4, 'f32': 4, 'b64': 8, 'u64': 8, 's64': 8, 'f64': 8}
|
|
_PTX = {}
|
|
|
|
|
|
def find_ptx(kernel):
|
|
"""the corpus module that defines the entry KERNEL (the PREPASS entry is never selected)"""
|
|
if kernel not in _PTX:
|
|
hits = [p for p in sorted(CORPUS.glob('*.ptx')) if f'.entry {kernel}(' in p.read_text(errors='replace')]
|
|
if len(hits) != 1:
|
|
raise FileNotFoundError(f'{len(hits)} corpus modules define entry {kernel}: {hits}')
|
|
_PTX[kernel] = hits[0]
|
|
return _PTX[kernel]
|
|
|
|
|
|
def bind(launch):
|
|
"""resource handles of a launch: handle value (read from the argument block) -> Texture / fresh Surface"""
|
|
handles, surfaces, textures = {}, [], {}
|
|
for img in launch.images:
|
|
h = struct.unpack_from('<Q', launch.args, img.offset)[0]
|
|
if img.kind == 'texture':
|
|
if img.offset not in textures:
|
|
textures[img.offset] = ptxsim.Texture(img.path, img.width, img.height, img.channels, img.bits, img.linear)
|
|
handles[h] = textures[img.offset]
|
|
else:
|
|
s = ptxsim.Surface(img.width, img.height, img.channels, img.bits)
|
|
handles[h] = s
|
|
surfaces.append((img, s))
|
|
return handles, surfaces
|
|
|
|
|
|
def make_watch(spec, records):
|
|
def hook(blk, st, mask):
|
|
for name in spec.get(st.line, ()):
|
|
records.append({'line': st.line, 'stmt': st.text, 'reg': name, 'width': blk.width(name),
|
|
'values': np.array(blk.get(name), copy=True), 'mask': mask.copy()})
|
|
return hook
|
|
|
|
|
|
def resolve(capture, kernel, frame):
|
|
cap = capture if isinstance(capture, Capture) else Capture(Path(os.path.expanduser(str(capture))))
|
|
launch = cap.find(kernel, frame)
|
|
return cap, launch
|
|
|
|
|
|
def run_block(capture, kernel, frame, ctaid, watch=None, mma_hook=None, limit=None):
|
|
"""Run one block (ctaid = (x, y) or (x, y, z)) of a captured launch. See the module docstring."""
|
|
cap, launch = resolve(capture, kernel, frame)
|
|
ptx = find_ptx(launch.kernel)
|
|
handles, surfaces = bind(launch)
|
|
records = []
|
|
hooks = [make_watch(watch, records)] if watch else []
|
|
ctaid = tuple(ctaid) + (0,) * (3 - len(ctaid))
|
|
t0 = time.perf_counter()
|
|
blk, mem, steps = ptxsim.run_block(str(ptx), launch.kernel, str(launch.directory), ctaid,
|
|
hooks=hooks, mma_hook=mma_hook, limit=limit, resources=handles)
|
|
seconds = time.perf_counter() - t0
|
|
return {'capture': cap, 'launch': launch, 'ptx': ptx, 'ctaid': ctaid, 'block': blk, 'regs': blk.regs,
|
|
'stores': mem.stores, 'mem': mem, 'surfaces': surfaces, 'watch': records,
|
|
'steps': steps, 'seconds': seconds}
|
|
|
|
|
|
def le_bits(raw):
|
|
"""(n, nbytes) little-endian bytes -> uint64 bit patterns"""
|
|
out = np.zeros(raw.shape[0], U64)
|
|
for b in range(raw.shape[1]):
|
|
out |= raw[:, b].astype(U64) << U64(8 * b)
|
|
return out
|
|
|
|
|
|
def _decode(bits, ty):
|
|
if ty == 'f16':
|
|
return (bits & U64(0xffff)).astype(np.uint16).view(np.float16).astype(np.float64)
|
|
return (bits & U64(0xffffffff)).astype(np.uint32).view(np.float32).astype(np.float64)
|
|
|
|
|
|
def _ordered(bits, n):
|
|
"""monotone integer key of an IEEE bit pattern: adjacent floats differ by 1"""
|
|
sign = (bits >> U64(n - 1)) & U64(1)
|
|
mag = (bits & U64((1 << (n - 1)) - 1)).astype(np.int64)
|
|
return np.where(sign == 1, -mag, mag)
|
|
|
|
|
|
def compare_elements(sim, ref, ty):
|
|
"""sim/ref: uint64 bit patterns of the same elements (type ty). Returns counts and ulp statistics."""
|
|
same = sim == ref
|
|
out = {'n': int(sim.size), 'identical': int(same.sum()), 'differing': int((~same).sum()),
|
|
'nan_mismatch': 0, 'max_ulp': 0, 'ulp_hist': {}}
|
|
if out['differing'] == 0:
|
|
return out
|
|
diff = np.nonzero(~same)[0]
|
|
if ty in ('f16', 'f32'):
|
|
n = 16 if ty == 'f16' else 32
|
|
with np.errstate(all='ignore'):
|
|
nan = np.isnan(_decode(sim[diff], ty)) | np.isnan(_decode(ref[diff], ty))
|
|
out['nan_mismatch'] = int(nan.sum())
|
|
ulp = np.abs(_ordered(sim[diff], n) - _ordered(ref[diff], n))[~nan]
|
|
else:
|
|
ulp = np.abs(sim[diff].astype(np.int64) - ref[diff].astype(np.int64))
|
|
if ulp.size:
|
|
out['max_ulp'] = int(ulp.max())
|
|
vals, counts = np.unique(np.minimum(ulp, 9), return_counts=True)
|
|
out['ulp_hist'] = {(f'>{8}' if v == 9 else int(v)): int(c) for v, c in zip(vals, counts)}
|
|
return out
|
|
|
|
|
|
def compare(result):
|
|
"""Compare the stores and surface writes of a run_block result with the capture's post-launch state."""
|
|
cap, launch, mem = result['capture'], result['launch'], result['mem']
|
|
groups, refs, report = {}, {}, []
|
|
for aid, off, raw, ty in mem.stores:
|
|
nb = NB[ty]
|
|
m = raw.shape[1] // nb
|
|
groups.setdefault((aid, ty), []).append((off[:, None] + np.arange(m) * nb).reshape(-1))
|
|
for (aid, ty), plist in sorted(groups.items()):
|
|
pos = np.unique(np.concatenate(plist))
|
|
nb = NB[ty]
|
|
idx = pos[:, None] + np.arange(nb)
|
|
if aid not in refs:
|
|
refs[aid] = np.frombuffer(launch.post_buffer(cap, aid), np.uint8)
|
|
k = mem.ids.index(aid)
|
|
st = compare_elements(le_bits(mem.data[k][idx]), le_bits(refs[aid][idx]), ty)
|
|
st['name'] = f'global alloc {aid} ({ty}, {"arena" if aid == launch.arena_alloc else "other"})'
|
|
report.append(st)
|
|
for img, s in result['surfaces']:
|
|
ref = np.fromfile(img.path, np.uint8)
|
|
if ref.size != s.data.size:
|
|
raise ValueError(f'surface arg{img.offset}: dump {ref.size} bytes, surface {s.data.size}')
|
|
ntex, nb = s.W * s.H, s.bits // 8
|
|
written = s.written.reshape(ntex, s.C, nb).all(axis=2)
|
|
ty = {16: 'f16', 32: 'f32'}.get(s.bits, 'u8')
|
|
st = compare_elements(le_bits(s.data.reshape(ntex, s.C, nb)[written]),
|
|
le_bits(ref.reshape(ntex, s.C, nb)[written]), ty)
|
|
st['name'] = f'surface arg{img.offset} {img.width}x{img.height}x{img.channels} {ty}'
|
|
report.append(st)
|
|
return report
|
|
|
|
|
|
def format_group(st):
|
|
n = st['n']
|
|
pct = 100.0 * st['identical'] / n if n else 100.0
|
|
line = f" {st['name']}: elements {n}, identical {st['identical']} ({pct:.3f}%), differing {st['differing']}"
|
|
if st['differing']:
|
|
line += f", max ulp {st['max_ulp']}, ulp histogram {st['ulp_hist']}"
|
|
if st['nan_mismatch']:
|
|
line += f", NaN mismatches {st['nan_mismatch']}"
|
|
return line
|
|
|
|
|
|
def print_watch(records, limit):
|
|
for rec in records[:limit]:
|
|
lanes = np.nonzero(rec['mask'])[0][:8]
|
|
parts = []
|
|
for ln in lanes:
|
|
v = int(rec['values'][ln])
|
|
if rec['width'] == 16:
|
|
dec = f'{float(_decode(np.array([v], U64), "f16")[0]):g}'
|
|
elif rec['width'] == 32:
|
|
dec = f'{float(_decode(np.array([v], U64), "f32")[0]):g}'
|
|
elif rec['width'] == 1:
|
|
dec = str(bool(v))
|
|
else:
|
|
dec = str(v)
|
|
parts.append(f'lane{ln}={v:#x}({dec})')
|
|
print(f" watch line {rec['line']} {rec['reg']} [{rec['stmt'][:60]}] active={int(rec['mask'].sum())}: " + ' '.join(parts))
|
|
if len(records) > limit:
|
|
print(f' ... {len(records) - limit} more watch records')
|
|
|
|
|
|
def main(argv=None):
|
|
ap = argparse.ArgumentParser(description='run captured RR CUDA blocks on the CPU and compare with the capture')
|
|
ap.add_argument('capture')
|
|
ap.add_argument('kernel')
|
|
ap.add_argument('frame', type=int)
|
|
ap.add_argument('--blocks', required=True)
|
|
ap.add_argument('--watch', action='append', default=[])
|
|
ap.add_argument('--mma', choices=['exact', 'split16'], default='exact')
|
|
ap.add_argument('--watch-max', type=int, default=16)
|
|
args = ap.parse_args(argv)
|
|
nums = [int(v) for v in args.blocks.split(',')]
|
|
if len(nums) % 2:
|
|
ap.error('--blocks needs X,Y pairs')
|
|
blocks = list(zip(nums[0::2], nums[1::2]))
|
|
watch = {}
|
|
for spec in args.watch:
|
|
line, names = spec.split(':', 1)
|
|
watch.setdefault(int(line), []).extend(n.strip() for n in names.split(',') if n.strip())
|
|
ptxsim.MMA_MODE = args.mma
|
|
|
|
cap, launch = resolve(args.capture, args.kernel, args.frame)
|
|
print(f'launch {launch.sequence} {launch.kernel} grid {launch.grid} block {launch.block} '
|
|
f'ptx {find_ptx(launch.kernel).name} mma {args.mma}')
|
|
for bx, by in blocks:
|
|
r = run_block(cap, args.kernel, args.frame, (bx, by), watch=watch or None)
|
|
print(f'block ({bx},{by}): {r["steps"]} statements in {r["seconds"]:.2f}s')
|
|
if watch:
|
|
print_watch(r['watch'], args.watch_max)
|
|
for st in compare(r):
|
|
print(format_group(st))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|