Files
exoplasim/scripts/benchmark-transforms.py

68 lines
4.1 KiB
Python
Executable File

#!/usr/bin/env python3
import argparse
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile
import numpy as np
root = Path(__file__).resolve().parents[1]
parser = argparse.ArgumentParser(description='Compare upstream and current Legendre kernels at the selected precision.')
parser.add_argument('--resolutions', nargs='+', type=int, choices=[21, 42, 85], default=[21, 42, 85])
parser.add_argument('--repeats', type=int, default=20)
parser.add_argument('--timeout', type=int, default=180)
parser.add_argument('--work', type=Path)
parser.add_argument('--reference', default='exoplasim-upstream-import')
parser.add_argument('--precision', type=int, choices=[4, 8], default=8)
args = parser.parse_args()
if args.repeats < 1:
parser.error('--repeats must be positive')
(root / 'tmp').mkdir(exist_ok=True)
work = args.work.resolve() if args.work else Path(tempfile.mkdtemp(prefix='transform-benchmark-', dir=root / 'tmp'))
work.mkdir(parents=True, exist_ok=True)
baseline = subprocess.check_output(['git', 'show', args.reference + ':exoplasim/plasim/src/legmod.f90'], cwd=root, timeout=20)
candidate = (root / 'exoplasim/plasim/src/legmod.f90').read_bytes()
stub, program = (root / 'scripts/transform-benchmark.f90').read_text().split('program transform_benchmark', 1)
results = []
for truncation in args.resolutions:
for partition in [1, 2]:
for mode, flags in [('checked', ['-fcheck=all', '-DEXOPLASIM_CHECKED']), ('native', ['-march=native', '-ffp-contract=off'])]:
outputs = {}
timings = {}
for label, source in [('baseline', baseline), ('candidate', candidate)]:
directory = work / f't{truncation}-p{partition}-{mode}-{label}'
directory.mkdir(exist_ok=True)
(directory / 'pumamod.f90').write_text(stub)
(directory / 'driver.f90').write_text('program transform_benchmark' + program)
(directory / 'legmod.f90').write_bytes(source)
command = [os.environ.get('FC', 'gfortran'), '-O3', '-cpp', f'-DTRUNC={truncation}',
f'-DPARTITION={partition}', *(['-fdefault-real-8'] if args.precision == 8 else []), '-ffree-line-length-none',
'-ffpe-trap=invalid,zero,overflow', *flags, 'pumamod.f90', 'legmod.f90',
'driver.f90', '-o', 'benchmark']
with (directory / 'build.log').open('w') as log:
subprocess.run(command, cwd=directory, check=True, timeout=args.timeout,
stdout=log, stderr=subprocess.STDOUT)
output = subprocess.check_output(['./benchmark', str(args.repeats)], cwd=directory, timeout=args.timeout)
(directory / 'timings.txt').write_bytes(output)
timings[label] = list(map(float, output.decode().split()))
outputs[label] = np.fromfile(directory / 'output.bin', dtype=np.float64 if args.precision == 8 else np.float32)
reference, current = outputs['baseline'], outputs['candidate']
if current.shape != reference.shape or not np.isfinite(current).all() or not np.isfinite(reference).all():
raise AssertionError('Invalid kernel output')
delta = float(np.max(np.abs(current - reference)))
scaled = delta / max(1.0, float(np.max(np.abs(reference))))
record = dict(truncation=truncation, partition=partition, mode=mode, precision=args.precision,
max_absolute_error=delta, scaled_error=scaled,
bitwise_identical=reference.tobytes() == current.tobytes(),
baseline_seconds=timings['baseline'], candidate_seconds=timings['candidate'],
speedup=[a / b for a, b in zip(timings['baseline'], timings['candidate'])])
results.append(record)
(work / 'results.json').write_text(json.dumps(results, indent=2) + '\n')
print(json.dumps(record), flush=True)
if scaled > 1e-12:
raise AssertionError(f'Kernel numerical mismatch: {record}')
print('Results:', work / 'results.json')