Improvements to benchmarking system #248

matt-graham · 2024-11-29T10:58:10Z

A collection of improvements to benchmarking system made during the DiRAC NVIDIA hackathon

Due to behaviour noted in NumPy generate_precomputes implementation has cubic in L memory overhead #246 switches to always using JAX implementation for generating precomputes for spherical transforms to avoid cubic in L overhead.
Adds benchmark modules for separate precompute versions of transforms.
Some improvements to help text of benchmarking scripts
Adds option to run benchmarked functions once initially without recording time to remove effect of initial JIT compilation from reported times
Fix to issue causing only last parameter set results being recorded in JSON output
JSON output now records additional metadata about host system such as installed package versions and information about GPUs and CPUs available, to make these more useful artefacts to store for longer term comparison.

matt-graham added 15 commits November 27, 2024 11:53

Use JAX function for generating precomputes

fdf93e3

Add benchmarks for precompute versions of spherical transforms

268e621

Show defaults in benchmark help text

1f5182a

Add flag for running once and discarding to eliminate JIT overhead

24d0441

Use module docstring for benchmark help description

3adac86

Make benchmark module docstring consistent

0c87cd6

Add precompute Wigner benchmarks

b3e3283

Record all benchmark results to JSON to not just last param set

dfaa569

Record system info and date time in JSON output

54db6d2

Expose recursion and mode parameters for precompute benchmarks

f4e9458

Record CPU and GPU info in benchmarks output if available

e917f05

Record GPU memory + CUDA info in results

742fdad

Fix Wigner benchmarks

942cf9e

Removing unused imports

64a2c71

Update benchmarking README to reflect new modules and functionality

77fa3a3

Provide feedback