96 lines
4.4 KiB
Python
96 lines
4.4 KiB
Python
"""Compare real K captures with the upstream CPU PTX interpreter.
|
|
|
|
The local DLL and captures remain private. Only numerical metrics are public.
|
|
Use this independent reference to audit the NumPy model's reduction uncertainty.
|
|
"""
|
|
import argparse
|
|
import math
|
|
import pathlib
|
|
import re
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[2]
|
|
REFERENCE = pathlib.Path(__file__).resolve().parent / 'reference'
|
|
if not (REFERENCE / 'ptxsim.py').exists():
|
|
REFERENCE = ROOT / 'kernels/tools'
|
|
sys.path.insert(0, str(REFERENCE))
|
|
import pwin_model as pm
|
|
pm.NORM_FRAGMENT_ORDER = True
|
|
import pwin_check as pc
|
|
import ptxsim
|
|
from extract_dlss_ptx import fatbins, ptx_entries
|
|
|
|
|
|
def metrics(label, expected, actual):
|
|
expected = expected.astype(np.float64)
|
|
actual = actual.astype(np.float64)
|
|
if not np.isfinite(expected).all() or not np.isfinite(actual).all():
|
|
raise RuntimeError(f'{label} contains NaN/Inf')
|
|
error = np.abs(expected - actual)
|
|
peak = max(1., float(np.abs(expected).max()))
|
|
maximum = float(error.max())
|
|
relative = float((error / np.maximum(np.abs(expected), .001)).max())
|
|
rms = float(np.sqrt(np.mean(error ** 2)))
|
|
psnr = math.inf if rms == 0 else 20 * math.log10(peak / rms)
|
|
print(f'{label} max_abs={maximum:.9g} max_rel={relative:.9g} rms={rms:.9g} '
|
|
f'peak={peak:.9g} psnr_db={psnr:.6g}', flush=True)
|
|
return maximum, peak, psnr
|
|
|
|
|
|
def validate(args):
|
|
name = f'dltss_pwin_{args.layer}_layer'
|
|
dump = pm.Dump(str(args.capture_dir))
|
|
if dump.kernel != name:
|
|
raise RuntimeError('Capture kernel does not match requested layer')
|
|
native = pc.Outputs(dump, str(args.native_dir))
|
|
candidates = []
|
|
for bx, by in pc.windows(dump.grid, 12):
|
|
model = pc.model_window(dump, args.layer, bx, by)
|
|
expected, actual = pc.stored(dump, native, args.layer, bx, by, *model)
|
|
difference = float(np.abs(expected.astype(np.float64) - actual.astype(np.float64)).max())
|
|
candidates.append((difference, bx, by, model, expected, actual))
|
|
selected = sorted(candidates, key=lambda c: (-c[0], c[1], c[2]))[:args.windows]
|
|
pattern = re.compile(rb'\.entry\s+' + name.encode() + rb'\s*\(')
|
|
sources = [source.rstrip(b'\0') for _, fatbin in fatbins(args.dlss_dll.read_bytes())
|
|
for source in ptx_entries(fatbin) if pattern.search(source)]
|
|
if not sources:
|
|
raise RuntimeError('Local DLSS DLL does not contain the captured kernel')
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
path = args.output_dir / (name + '.ptx')
|
|
path.write_bytes(sources[-1])
|
|
for _, bx, by, model, numpy_expected, actual in selected:
|
|
_, memory, steps = ptxsim.run_block(str(path), name, str(args.capture_dir), (bx, by, 0))
|
|
reference_output = pc.Outputs.__new__(pc.Outputs)
|
|
reference_output.d = dump
|
|
reference_output.allocs = {identifier: (base, data) for identifier, base, data
|
|
in zip(memory.ids, memory.bases, memory.data)}
|
|
_, expected = pc.stored(dump, reference_output, args.layer, bx, by, *model)
|
|
print(f'PTX_REFERENCE layer={args.layer} window={bx},{by} statements={steps}', flush=True)
|
|
maximum, peak, psnr = metrics('PTX_VS_NATIVE', expected, actual)
|
|
metrics('PTX_VS_NUMPY', expected, numpy_expected)
|
|
# Half reductions and approximate reciprocal/rsqrt are not bit-exact
|
|
# across ISAs. Bound both the worst element and total error explicitly.
|
|
if maximum > max(.001, .005 * peak) or psnr < 60:
|
|
raise RuntimeError('Native layer differs from independent PTX reference')
|
|
print(f'PASS K_PTX_REFERENCE layer={args.layer} windows={len(selected)} nan_inf=0', flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--layer', choices=list(pc.LAYERS), required=True)
|
|
parser.add_argument('--capture-dir', type=pathlib.Path, required=True)
|
|
parser.add_argument('--native-dir', type=pathlib.Path, required=True)
|
|
parser.add_argument('--dlss-dll', type=pathlib.Path, required=True)
|
|
parser.add_argument('--output-dir', type=pathlib.Path, required=True)
|
|
parser.add_argument('--windows', type=int, default=2)
|
|
args = parser.parse_args()
|
|
if args.windows < 1:
|
|
parser.error('--windows must be positive')
|
|
try:
|
|
validate(args)
|
|
except Exception as error:
|
|
print(f'FAIL K_PTX_REFERENCE {error}', file=sys.stderr)
|
|
raise SystemExit(5)
|