70 lines
3.6 KiB
Python
70 lines
3.6 KiB
Python
"""Independent PTX interpreter validation of real M captures; keep outputs private."""
|
|
import argparse
|
|
import pathlib
|
|
import re
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
from m_replay_validate import me, LAYERS, blocks, model_block, stored, metrics
|
|
import ptxsim
|
|
from extract_dlss_ptx import fatbins, ptx_entries
|
|
|
|
|
|
def validate(args):
|
|
name = f'rrlite_{args.layer}_4x4'
|
|
parameters = me.Params(str(args.capture_dir))
|
|
if parameters.kernel != name:
|
|
raise RuntimeError('Capture kernel does not match requested M layer')
|
|
native = {i: np.fromfile(args.native_dir / f'alloc-{i}.bin', np.uint8) for i in parameters.allocs}
|
|
candidates = []
|
|
coordinates = args.block or blocks(parameters.grid, args.candidate_blocks)
|
|
for bx, by in coordinates:
|
|
if not (0 <= bx < parameters.grid[0] and 0 <= by < parameters.grid[1]):
|
|
raise RuntimeError(f'Block {bx},{by} is outside the capture grid')
|
|
model = model_block(parameters, args.layer, bx, by)
|
|
expected, actual = stored(parameters, args.layer, bx, by, native, model)
|
|
difference = float(np.abs(expected - actual).max())
|
|
candidates.append((difference, bx, by, model, expected, actual))
|
|
selected = sorted(candidates, key=lambda c: (-c[0], c[1], c[2]))[:args.blocks]
|
|
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 this M 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))
|
|
allocation_data = dict(zip(memory.ids, memory.data))
|
|
_, expected = stored(parameters, args.layer, bx, by, allocation_data, model)
|
|
print(f'PTX_REFERENCE layer={args.layer} block={bx},{by} statements={steps}', flush=True)
|
|
_, _, psnr, fp8_steps = metrics(expected, actual, 'PTX_VS_NATIVE')
|
|
metrics(expected, numpy_expected, 'PTX_VS_NUMPY')
|
|
if fp8_steps.max() > 1 or psnr < 50:
|
|
raise RuntimeError('Native M differs from independent PTX reference')
|
|
print(f'PASS M_PTX_REFERENCE layer={args.layer} blocks={len(selected)} nan_inf=0', flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--layer', choices=list(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('--blocks', type=int, default=2)
|
|
parser.add_argument('--candidate-blocks', type=int, default=12,
|
|
help='NumPy scan size used to choose the worst blocks for independent PTX')
|
|
parser.add_argument('--block', type=lambda value: tuple(map(int, value.split(','))),
|
|
action='append', help='Explicit x,y regression block (repeatable)')
|
|
args = parser.parse_args()
|
|
if args.blocks < 1 or args.candidate_blocks < 1:
|
|
parser.error('--blocks must be positive')
|
|
try:
|
|
validate(args)
|
|
except Exception as error:
|
|
print(f'FAIL M_PTX_REFERENCE {error}', file=sys.stderr)
|
|
raise SystemExit(5)
|