Skip to main content

bench/
bench.rs

1use std::{env, process};
2
3use openvm_cuda_backend::{
4    logup_zerocheck::{fractional_sumcheck_gpu, make_synthetic_leaves, FractionalInputSize},
5    prelude::EF,
6    sponge::DuplexSpongeGpu,
7};
8use openvm_cuda_common::{
9    common::get_device,
10    copy::MemCopyD2D,
11    stream::{CudaStream, GpuDeviceCtx, StreamGuard},
12};
13use p3_field::PrimeCharacteristicRing;
14
15fn parse_usize(var: &str, default: usize) -> usize {
16    env::var(var)
17        .ok()
18        .and_then(|x| x.parse::<usize>().ok())
19        .unwrap_or(default)
20}
21
22/// All benchmarks output the same CSV format: `n,run_idx,is_warmup,elapsed_ms`
23fn bench_fractional_sumcheck() -> Result<(), Box<dyn std::error::Error>> {
24    let n = parse_usize("SWIRL_BENCH_N", 24);
25    let repeats = parse_usize("SWIRL_BENCH_REPEATS", 3);
26    let warmups = parse_usize("SWIRL_BENCH_WARMUPS", 1);
27
28    let device_ctx = GpuDeviceCtx {
29        device_id: get_device()? as u32,
30        stream: StreamGuard::new(CudaStream::new_non_blocking()?),
31    };
32    let template = make_synthetic_leaves(n, &device_ctx)?;
33
34    println!("run_idx,is_warmup,elapsed_ms");
35
36    for run_idx in 0..(warmups + repeats) {
37        let is_warmup = run_idx < warmups;
38        let leaves = template
39            .device_copy_on(&device_ctx)
40            .expect("device copy leaves");
41
42        let mut transcript = DuplexSpongeGpu::default();
43        let mut mem = openvm_cuda_common::memory_manager::MemTracker::start("bench.fractional");
44
45        device_ctx.stream.synchronize().expect("sync before timing");
46        let t0 = std::time::Instant::now();
47        let _ = fractional_sumcheck_gpu(
48            &mut transcript,
49            leaves,
50            FractionalInputSize::dense(1usize << n),
51            EF::ZERO,
52            false,
53            &mut mem,
54            &device_ctx,
55        )?;
56        device_ctx.stream.synchronize().expect("sync after timing");
57        let ms = t0.elapsed().as_secs_f64() * 1000.0;
58
59        println!("{run_idx},{},{:.4}", is_warmup as u8, ms);
60    }
61
62    Ok(())
63}
64
65#[allow(clippy::type_complexity)]
66const BENCHMARKS: &[(&str, fn() -> Result<(), Box<dyn std::error::Error>>)] =
67    &[("fractional-sumcheck", bench_fractional_sumcheck)];
68
69fn main() {
70    let args: Vec<String> = env::args().collect();
71    if args.len() < 2 {
72        eprintln!("Usage: bench <benchmark>");
73        eprintln!();
74        eprintln!("Available benchmarks:");
75        for (name, _) in BENCHMARKS {
76            eprintln!("  {name}");
77        }
78        process::exit(1);
79    }
80
81    let name = &args[1];
82    for (bench_name, bench_fn) in BENCHMARKS {
83        if name == bench_name {
84            if let Err(e) = bench_fn() {
85                eprintln!("Error: {e}");
86                process::exit(1);
87            }
88            return;
89        }
90    }
91
92    eprintln!("Unknown benchmark: {name}");
93    eprintln!();
94    eprintln!("Available benchmarks:");
95    for (name, _) in BENCHMARKS {
96        eprintln!("  {name}");
97    }
98    process::exit(1);
99}