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
22fn 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}