Skip to content

Commit

Permalink
refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
wenyuzhao committed May 17, 2024
1 parent 27b8af9 commit 8cc51be
Show file tree
Hide file tree
Showing 4 changed files with 134 additions and 112 deletions.
38 changes: 21 additions & 17 deletions harness/src/bencher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,12 @@ use std::{
time::{Duration, Instant},
};

use clap::{Parser, ValueEnum};
use clap::Parser;

use crate::probe::ProbeManager;
use crate::{
probe::ProbeManager,
record::{Record, StatPrintFormat},
};

#[derive(Parser, Debug)]
pub struct BenchArgs {
Expand Down Expand Up @@ -41,13 +44,6 @@ pub struct BenchArgs {
pub current_build: Option<String>,
}

#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
#[clap(rename_all = "kebab_case")]
pub(crate) enum StatPrintFormat {
Table,
Yaml,
}

pub trait Value: ToString + 'static {}

impl Value for f64 {}
Expand Down Expand Up @@ -285,16 +281,24 @@ impl SingleBenchmarkRunner {
}

fn dump_counters(&self, iteration: usize, is_timing_iteration: bool) {
self.bencher.probes.borrow().dump_counters(
&self.bench_name,
self.args.output_csv.as_ref(),
self.args.current_invocation,
self.args.current_build.as_ref(),
std::mem::take(&mut *self.bencher.extra_stats.lock().unwrap()),
StatPrintFormat::Yaml,
let probe_stats = self
.bencher
.probes
.borrow()
.get_counter_values(std::mem::take(
&mut *self.bencher.extra_stats.lock().unwrap(),
));
let record = Record {
name: &self.bench_name,
csv: self.args.output_csv.as_ref(),
invocation: self.args.current_invocation,
build: self.args.current_build.as_ref(),
format: StatPrintFormat::Yaml,
iteration,
is_timing_iteration,
);
stats: probe_stats,
};
record.dump_values();
}

fn run_once_impl(&mut self, iteration: usize) -> f32 {
Expand Down
1 change: 1 addition & 0 deletions harness/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
mod bencher;
pub mod probe;
mod record;
pub mod utils;

pub use bencher::{BenchTimer, Bencher, Value};
Expand Down
100 changes: 5 additions & 95 deletions harness/src/probe.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
use std::io::Write;
use std::time::Duration;
use std::{collections::HashMap, fs::OpenOptions, path::PathBuf, time::Instant};
use std::{collections::HashMap, time::Instant};

use libloading::{Library, Symbol};
use serde::{de::DeserializeOwned, Deserialize, Serialize};

use crate::bencher::{StatPrintFormat, Value};
use crate::bencher::Value;

struct Counters {
counters: Vec<(String, f32)>,
Expand Down Expand Up @@ -173,79 +172,10 @@ impl ProbeManager {
self.counters.get_value(name)
}

fn dump_counters_stderr_table(&self, stats: &[(String, Box<dyn Value>)]) {
for (name, _) in stats {
eprint!("{}\t", name);
}
eprintln!();
for (_, value) in stats {
eprint!("{}\t", value.to_string());
}
eprintln!();
}

fn dump_counters_stderr_yaml(&self, stats: &[(String, Box<dyn Value>)]) {
for (name, value) in stats {
eprintln!("{}: {}", name, value.to_string());
}
}

fn dump_counters_stderr(&self, stats: &[(String, Box<dyn Value>)], format: StatPrintFormat) {
let force_table = std::env::var("HARNESS_LOG_STAT_FORMAT") == Ok("table".to_owned());
if force_table {
return self.dump_counters_stderr_table(stats);
}
match format {
StatPrintFormat::Table => self.dump_counters_stderr_table(stats),
StatPrintFormat::Yaml => self.dump_counters_stderr_yaml(stats),
}
}

fn dump_counters_csv(
pub(crate) fn get_counter_values(
&self,
name: &str,
csv: Option<&PathBuf>,
invocation: Option<usize>,
build: Option<&String>,
stats: &[(String, Box<dyn Value>)],
iteration: usize,
) {
if let Some(csv) = csv {
if !csv.exists() {
let mut headers = "bench,build,invocation,iteration".to_owned();
for (name, _value) in stats {
headers += ",";
headers += name;
}
headers += "\n";
std::fs::write(csv, headers).unwrap();
}
let mut record = format!(
"{},{},{},{}",
name,
build.unwrap(),
invocation.unwrap_or(0),
iteration
);
for (_, value) in stats {
record += &format!(",{}", value.to_string());
}
let mut csv = OpenOptions::new().append(true).open(csv).unwrap();
writeln!(csv, "{record}").unwrap();
}
}

pub(crate) fn dump_counters(
&self,
name: &str,
csv: Option<&PathBuf>,
invocation: Option<usize>,
build: Option<&String>,
extra_stats: Vec<(String, Box<dyn Value>)>,
format: StatPrintFormat,
iteration: usize,
is_timing_iteration: bool,
) {
) -> HashMap<String, Box<dyn Value>> {
// Collect all stats
let mut stats_map: HashMap<String, Box<dyn Value>> = HashMap::new();
for (name, value) in &self.counters.counters {
Expand All @@ -254,26 +184,6 @@ impl ProbeManager {
for (name, value) in extra_stats {
stats_map.insert(name.clone(), value);
}
let time = stats_map.remove("time").unwrap();
let mut stats: Vec<(String, Box<dyn Value>)> = vec![];
for (name, value) in stats_map {
stats.push((name.clone(), value));
}
stats.sort_by_key(|x| x.0.clone());
stats.insert(0, ("time".to_owned(), time));
if is_timing_iteration {
// Print to the log file
let banner_start = std::env::var("HARNESS_LOG_STAT_BANNER_START").unwrap_or_else(|_| {
"============================ Harness Statistics Totals ============================".to_string()
});
eprintln!("{banner_start}");
self.dump_counters_stderr(&stats, format);
let banner_end = std::env::var("HARNESS_LOG_STAT_BANNER_END").unwrap_or_else(|_| {
"------------------------------ End Harness Statistics -----------------------------".to_string()
});
eprintln!("{banner_end}");
}
// Print to the CSV file
self.dump_counters_csv(name, csv, invocation, build, &stats, iteration);
stats_map
}
}
107 changes: 107 additions & 0 deletions harness/src/record.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
use std::io::Write;
use std::{collections::HashMap, fs::OpenOptions, path::PathBuf};

use clap::ValueEnum;

use crate::Value;

#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
#[clap(rename_all = "kebab_case")]
pub(crate) enum StatPrintFormat {
Table,
Yaml,
}

pub(crate) struct Record<'a> {
pub name: &'a str,
pub csv: Option<&'a PathBuf>,
pub invocation: Option<usize>,
pub build: Option<&'a String>,
pub format: StatPrintFormat,
pub iteration: usize,
pub is_timing_iteration: bool,
pub stats: HashMap<String, Box<dyn Value>>,
}

impl<'a> Record<'a> {
fn dump_counters_stderr_table(&self, stats: &[(String, Box<dyn Value>)]) {
for (name, _) in stats {
eprint!("{}\t", name);
}
eprintln!();
for (_, value) in stats {
eprint!("{}\t", value.to_string());
}
eprintln!();
}

fn dump_counters_stderr_yaml(&self, stats: &[(String, Box<dyn Value>)]) {
for (name, value) in stats {
eprintln!("{}: {}", name, value.to_string());
}
}

fn dump_counters_stderr(&self, stats: &[(String, Box<dyn Value>)], format: StatPrintFormat) {
let force_table = std::env::var("HARNESS_LOG_STAT_FORMAT") == Ok("table".to_owned());
if force_table {
return self.dump_counters_stderr_table(stats);
}
match format {
StatPrintFormat::Table => self.dump_counters_stderr_table(stats),
StatPrintFormat::Yaml => self.dump_counters_stderr_yaml(stats),
}
}

fn dump_counters_csv(&self, stats: &[(String, Box<dyn Value>)]) {
if let Some(csv) = self.csv {
if !csv.exists() {
let mut headers = "bench,build,invocation,iteration".to_owned();
for (name, _value) in stats {
headers += ",";
headers += name;
}
headers += "\n";
std::fs::write(csv, headers).unwrap();
}
let mut record = format!(
"{},{},{},{}",
self.name,
self.build.unwrap(),
self.invocation.unwrap_or(0),
self.iteration
);
for (_, value) in stats {
record += &format!(",{}", value.to_string());
}
let mut csv = OpenOptions::new().append(true).open(csv).unwrap();
writeln!(csv, "{record}").unwrap();
}
}

pub fn dump_values(mut self) {
let mut stats_map = std::mem::take(&mut self.stats);
let time = stats_map.remove("time");
let mut stats: Vec<(String, Box<dyn Value>)> = vec![];
for (name, value) in stats_map {
stats.push((name.clone(), value));
}
stats.sort_by_key(|x| x.0.clone());
if let Some(time) = time {
stats.insert(0, ("time".to_owned(), time));
}
if self.is_timing_iteration {
// Print to the log file
let banner_start = std::env::var("HARNESS_LOG_STAT_BANNER_START").unwrap_or_else(|_| {
"============================ Harness Statistics Totals ============================".to_string()
});
eprintln!("{banner_start}");
self.dump_counters_stderr(&stats, self.format);
let banner_end = std::env::var("HARNESS_LOG_STAT_BANNER_END").unwrap_or_else(|_| {
"------------------------------ End Harness Statistics -----------------------------".to_string()
});
eprintln!("{banner_end}");
}
// Print to the CSV file
self.dump_counters_csv(&stats);
}
}

0 comments on commit 8cc51be

Please sign in to comment.