energy-bench 0.3.0

use std::{fmt, hint, path::PathBuf, time::Instant};

use crate::{ EnergyAccumulator, IdlePower, Metadata, output::OutputFiles, run_result::RunResult};

#[derive(clap::Parser, Clone)]
pub struct EnergyBenchConfig {
    /// The number of warmup runs to do.
    ///
    /// These are run before the actual benchmarking process starts, and are often necessary for
    /// multiple reasons. These include: warming up the caches, waiting for turboboost behavior
    /// to subside, and waiting for the CPU to physically heat up.
    /// This may or may not be necessary depending on system configuration and the benchmark.
    #[arg(long, default_value_t = 0)]
    pub warmup_runs: usize,
    /// The number of benchmark runs to do.
    #[arg(long, default_value_t = 1)]
    pub benchmark_runs: usize,
    /// Set the time taken to estimate idle power draw.
    #[arg(long, default_value_t = 120)]
    pub idle_seconds: usize,
    /// Set the path to the idle results cache file.
    #[arg(long)]
    pub idle_path: Option<PathBuf>,
}

impl Default for EnergyBenchConfig {
    fn default() -> Self {
        Self {
            warmup_runs: 0,
            benchmark_runs: 1,
            idle_seconds: 120,
            idle_path: None,
        }
    }
}

pub struct EnergyBench<MD: Metadata<COLS>, const COLS: usize> {
    warmup_runs: usize,
    benchmark_runs: usize,
    output: OutputFiles<MD, COLS>,
    idle_power: Option<IdlePower>,
    probes: Box<dyn EnergyAccumulator>,
}

impl<MD: Metadata<COLS>, const COLS: usize> EnergyBench<MD, COLS> {
    pub fn default(name: &str) -> Self {
        let config = EnergyBenchConfig::default();
        Self::new(name, config)
    }

    pub fn new(name: &str, config: EnergyBenchConfig) -> Self {
        #[cfg(feature = "software-energy-lab")]
        let probes = Box::new(crate::software_energy_lab::SoftwareEnergyLab::new());
        #[cfg(not(feature = "software-energy-lab"))]
        let probes = Box::new(crate::energy_acc::DefaultEnergyAccumulator::new());
        Self::new_with(name, config, probes)
    }

    pub fn new_with(name: &str, config: EnergyBenchConfig, mut probes: Box<dyn EnergyAccumulator>) -> Self {
        let idle_power = if config.idle_seconds > 0 {
            Some(IdlePower::init(config.idle_path.as_ref(), config.idle_seconds, &mut probes))
        } else {
            None
        };

        Self {
            warmup_runs: config.warmup_runs,
            benchmark_runs: config.benchmark_runs,
            output: OutputFiles::new(name),
            idle_power,
            probes,
        }
    }

    pub fn benchmark<E>(&mut self, metadata: MD, bench_fn: impl Fn() -> Result<(), E>) -> Vec<RunResult>
    where
        E: fmt::Debug,
    {
        self.benchmark_with(metadata, || (), |()| bench_fn())
    }

    pub fn benchmark_with<T, E>(&mut self, metadata: MD, setup_fn: impl Fn() -> T, bench_fn: impl Fn(T) -> Result<(), E>) -> Vec<RunResult>
    where
        E: fmt::Debug,
    {
        self.warmup(&setup_fn, &bench_fn);
        let measurements = self.benchmark_runs(&setup_fn, &bench_fn);

        if measurements.iter().any(|res| res.runtime.as_millis() < 100) {
            log::warn!("One or more runs took less than 100ms, results may be inaccurate, consider wrapping the benchmarked function in a loop to increase runtime");
        }

        self.write(&metadata, &measurements);
        measurements
    }

    fn write(&mut self, metadata: &MD, measurements: &Vec<RunResult>) {
        if measurements.is_empty() {
            log::error!("All benchmarks failed, no results written");
        } else if let Err(e) = self.output.write_results(metadata, measurements) {
            log::error!("Error writing results: {}", e);
        }
    }

    fn warmup<T, E>(&self, setup_fn: &impl Fn() -> T, bench_fn: &impl Fn(T) -> Result<(), E>)
    where
        E: fmt::Debug,
    {
        for _ in 0..self.warmup_runs {
            let data = setup_fn();
            if let Err(e) = hint::black_box(bench_fn(data)) {
                log::error!("Warmup run failed: {e:?}");
            }
        }
    }

    fn benchmark_runs<T, E>(&mut self, setup_fn: &impl Fn() -> T, bench_fn: &impl Fn(T) -> Result<(), E>) -> Vec<RunResult>
    where
        E: fmt::Debug,
    {
        (1..=self.benchmark_runs)
            .filter_map(|i| {
                log::trace!("Starting benchmark run {i}");
                match self.benchmark_run(&setup_fn, &bench_fn) {
                    Ok(measurements) => Some(measurements),
                    Err(e) => {
                        log::error!("Benchmark run {i} failed: {e:?}");
                        None
                    }
                }
            })
            .collect()
    }

    fn benchmark_run<T, E>(&mut self, setup_fn: &impl Fn() -> T, bench_fn: &impl Fn(T) -> Result<(), E>) -> Result<RunResult, E>
    where
        E: fmt::Debug,
    {
        let inp = setup_fn();

        self.probes.reset();
        let now = Instant::now();
        hint::black_box(bench_fn(inp))?;
        let runtime = now.elapsed();
        let energy = self.probes.elapsed();

        let mut run_result = RunResult::new(runtime, energy);
        if let Some(idle) = &self.idle_power {
            run_result.subtract_idle(idle);
        }

        Ok(run_result)
    }
}