File size: 4,307 Bytes
9960a46 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | # SPDX-License-Identifier: Apache-2.0
"""Compare original baseline and final kernel on identical inputs and fp64 truth."""
import importlib.util
import json
from pathlib import Path
import sys
import torch
def load(path, name):
spec = importlib.util.spec_from_file_location(name, path)
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module)
return module
kind = sys.argv[1]
release_root = Path(__file__).resolve().parent
here = release_root / kind
bench = load(here / 'benchmark.py', f'{kind}_benchmark')
baseline = load(here / 'baseline.py', f'{kind}_baseline')
candidate = load(release_root / 'build/torch-cuda' / f'_{kind}.py', f'{kind}_candidate')
scope = {}
exec(json.loads((here / 'definition.json').read_text())['reference'], scope)
reference = scope['run']
device = torch.device('cuda')
report = {'component': kind, 'gpu': torch.cuda.get_device_name(), 'torch': torch.__version__,
'method': 'Identical inputs for both implementations, original judge unchanged, fp64 rebuilt truth.',
'global_seed_plus_batch': 20261026, 'relative_tolerance': bench.REL_TOL,
'absolute_cap_factor': bench.CAP, 'reference_error_floor_factor': bench.FLOOR,
'individual_state_check_is_additional': True, 'rows': []}
for batch in (16, 32, 64, 128, 256, 512):
torch.manual_seed(20261026 + batch)
first = bench.structured_inputs(batch, device, seed=batch)
ref1 = reference(*[x.clone() for x in first])
second = (*ref1, *bench.update_buffer(batch, device, torch.Generator(device=device).manual_seed(batch + 1)), first[-1])
random = bench.random_inputs(batch, device, seed=batch + 2)
for family, inputs, ref in (('structured', first, ref1), ('second_window', second, None), ('random', random, None)):
if ref is None:
ref = reference(*[x.clone() for x in inputs])
exact = bench.exact_state(inputs)
ref_error = (bench.decode(*ref) - exact).abs().mean((-2, -1))
magnitude = exact.abs().mean((-2, -1))
tight = (ref_error >= bench.FLOOR * magnitude) & (ref_error > 0)
row = {'batch_size': batch, 'family': family}
for name, impl in (('baseline', baseline), ('current', candidate)):
output = impl.run(*[x.clone() for x in inputs])
torch.cuda.synchronize()
passed, ratio = bench.judge(output, ref, exact)
error = (bench.decode(*output) - exact).abs().mean((-2, -1))
strict = torch.where(tight, error <= (1 + bench.REL_TOL) * ref_error, error <= bench.CAP * magnitude)
row[name] = {
'original_judge_passed': passed,
'worst_head_error_ratio': float(ratio),
'max_individual_head_error_ratio': float((error[tight] / ref_error[tight]).max()),
'individual_states_checked': int(strict.numel()),
'individual_states_passing_same_bound': int(strict.sum()),
'max_state_relative_mae': float((error[magnitude > 0] / magnitude[magnitude > 0]).max()),
}
assert passed, (kind, batch, family, name, ratio)
report['rows'].append(row)
(here / 'precision_comparison.json').write_text(json.dumps(report, indent=2) + '\n')
print(kind, batch, 'baseline / current:', [(r['family'], r['baseline']['worst_head_error_ratio'], r['current']['worst_head_error_ratio']) for r in report['rows'][-3:]], flush=True)
report['status'] = 'original_benchmark_pass'
report['summary'] = {}
for name in ('baseline', 'current'):
values = [row[name] for row in report['rows']]
report['summary'][name] = {
'original_cases_passed': sum(v['original_judge_passed'] for v in values),
'worst_original_benchmark_error_ratio': max(v['worst_head_error_ratio'] for v in values),
'individual_states_checked': sum(v['individual_states_checked'] for v in values),
'additional_individual_state_failures': sum(v['individual_states_checked'] - v['individual_states_passing_same_bound'] for v in values),
'max_individual_head_error_ratio': max(v['max_individual_head_error_ratio'] for v in values),
}
(here / 'precision_comparison.json').write_text(json.dumps(report, indent=2) + '\n')
print(kind, 'complete', flush=True)
|