Files
tfhe-rs/ci/parse_integer_benches_to_csv.py
2024-08-14 13:42:14 +02:00

104 lines
3.3 KiB
Python

import argparse
from pathlib import Path
import json
def main(args):
criterion_dir = Path(args.criterion_dir)
output_file = Path(args.output_file)
data = []
for json_file in sorted(criterion_dir.glob("**/*.json")):
if json_file.parent.name == "base" or json_file.name != "benchmark.json":
continue
try:
bench_data = json.loads(json_file.read_text())
estimate_file = json_file.with_name("estimates.json")
estimate_data = json.loads(estimate_file.read_text())
bench_function_id = bench_data["function_id"]
split = bench_function_id.split("::")
if len(split) == 5: # Signed integers
(_, _, function_name, parameter_set, bits) = split
else: # Unsigned integers
(_, function_name, parameter_set, bits) = split
if "_scalar_" in bits:
(bits, scalar) = bits.split("_bits_scalar_")
bits = int(bits)
scalar = int(scalar)
else:
(bits, _) = bits.split("_")
bits = int(bits)
scalar = None
estimate_mean_ms = estimate_data["mean"]["point_estimate"] / 1000000
estimate_lower_bound_ms = (
estimate_data["mean"]["confidence_interval"]["lower_bound"] / 1000000
)
estimate_upper_bound_ms = (
estimate_data["mean"]["confidence_interval"]["upper_bound"] / 1000000
)
data.append(
(
function_name,
parameter_set,
bits,
scalar,
estimate_mean_ms,
estimate_lower_bound_ms,
estimate_upper_bound_ms,
)
)
except Exception as e:
print(e)
pass
if len(data) == 0:
print("No integer bench found, skipping writing output file")
return
with open(output_file, "w", encoding="utf-8") as output:
output.write(
"function_name,parameter_set,bits,scalar,mean_ms,"
"confidence_interval_lower_bound_ms,confidence_interval_upper_bound_ms\n"
)
# Sort by func_name, bit width and then parameters
data.sort(key=lambda x: (x[0], x[2], x[1]))
for dat in data:
(
function_name,
parameter_set,
bits,
scalar,
estimate_mean_ms,
estimate_lower_bound_ms,
estimate_upper_bound_ms,
) = dat
output.write(
f"{function_name},{parameter_set},{bits},{scalar},{estimate_mean_ms},"
f"{estimate_lower_bound_ms},{estimate_upper_bound_ms}\n"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser("Parse criterion results to csv file")
parser.add_argument(
"--criterion-dir",
type=str,
default="target/criterion",
help="Where to look for criterion result json files",
)
parser.add_argument(
"--output-file",
type=str,
default="parsed_benches.csv",
help="Path of the output file, will be csv formatted",
)
main(parser.parse_args())