Add benchmarks for SIMD eval_at_point (#607)

crates/prover/benches/eval_at_point.rs

@@ -1,72 +1,40 @@
-use criterion::{black_box, Criterion};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use rand::rngs::SmallRng;
+use rand::{Rng, SeedableRng};
+use stwo_prover::core::backend::cpu::CPUBackend;
+use stwo_prover::core::backend::simd::SimdBackend;
+use stwo_prover::core::circle::CirclePoint;
+use stwo_prover::core::fields::m31::BaseField;
+use stwo_prover::core::poly::circle::{CirclePoly, PolyOps};
 
-#[cfg(target_arch = "x86_64")]
-pub fn cpu_eval_at_secure_point(c: &mut criterion::Criterion) {
-    use rand::rngs::SmallRng;
-    use rand::{Rng, SeedableRng};
-    use stwo_prover::core::backend::CPUBackend;
-    use stwo_prover::core::circle::CirclePoint;
-    use stwo_prover::core::fields::m31::BaseField;
-    use stwo_prover::core::fields::qm31::QM31;
-    use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps};
-    use stwo_prover::core::poly::NaturalOrder;
-    let log_size = 20;
-    let mut rng = SmallRng::seed_from_u64(0);
-
-    let domain = CanonicCoset::new(log_size as u32).circle_domain();
-    let evaluation = CircleEvaluation::<CPUBackend, _, NaturalOrder>::new(
-        domain,
-        (0..(1 << log_size))
-            .map(BaseField::from_u32_unchecked)
-            .collect(),
-    );
-    let poly = evaluation.bit_reverse().interpolate();
-    let x: QM31 = rng.gen();
-    let y: QM31 = rng.gen();
-
-    let point = CirclePoint { x, y };
-    c.bench_function("cpu eval_at_secure_field_point 2^20", |b| {
-        b.iter(|| {
-            black_box(<CPUBackend as PolyOps>::eval_at_point(&poly, point));
-        })
-    });
-}
+const LOG_SIZE: u32 = 20;
 
-#[cfg(target_arch = "x86_64")]
-pub fn avx512_eval_at_secure_point(c: &mut criterion::Criterion) {
-    use rand::rngs::SmallRng;
-    use rand::{Rng, SeedableRng};
-    use stwo_prover::core::backend::avx512::AVX512Backend;
-    use stwo_prover::core::circle::CirclePoint;
-    use stwo_prover::core::fields::m31::BaseField;
-    use stwo_prover::core::fields::qm31::QM31;
-    use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps};
-    use stwo_prover::core::poly::NaturalOrder;
-    let log_size = 20;
+fn bench_eval_at_secure_point<B: PolyOps>(c: &mut Criterion, id: &str) {
+    let poly = CirclePoly::new((0..1 << LOG_SIZE).map(BaseField::from).collect());
     let mut rng = SmallRng::seed_from_u64(0);
-
-    let domain = CanonicCoset::new(log_size as u32).circle_domain();
-    let evaluation = CircleEvaluation::<AVX512Backend, BaseField, NaturalOrder>::new(
-        domain,
-        (0..(1 << log_size))
-            .map(BaseField::from_u32_unchecked)
-            .collect(),
+    let x = rng.gen();
+    let y = rng.gen();
+    let point = CirclePoint { x, y };
+    c.bench_function(
+        &format!("{id} eval_at_secure_field_point 2^{LOG_SIZE}"),
+        |b| {
+            b.iter(|| B::eval_at_point(black_box(&poly), black_box(point)));
+        },
     );
-    let poly = evaluation.bit_reverse().interpolate();
-    let x: QM31 = rng.gen();
-    let y: QM31 = rng.gen();
+}
 
-    let point = CirclePoint { x, y };
-    c.bench_function("avx eval_at_secure_field_point 2^20", |b| {
-        b.iter(|| {
-            black_box(<AVX512Backend as PolyOps>::eval_at_point(&poly, point));
-        })
-    });
+fn eval_at_secure_point_benches(c: &mut Criterion) {
+    #[cfg(target_arch = "x86_64")]
+    if stwo_prover::platform::avx512_detected() {
+        use stwo_prover::core::backend::avx512::AVX512Backend;
+        bench_eval_at_secure_point::<AVX512Backend>(c, "avx");
+    }
+    bench_eval_at_secure_point::<SimdBackend>(c, "simd");
+    bench_eval_at_secure_point::<CPUBackend>(c, "cpu");
 }
 
-#[cfg(target_arch = "x86_64")]
-criterion::criterion_group!(
-    name=secure_eval;
-    config = Criterion::default().sample_size(10);
-    targets=avx512_eval_at_secure_point, cpu_eval_at_secure_point);
-criterion::criterion_main!(secure_eval);
+criterion_group!(
+        name = benches;
+        config = Criterion::default().sample_size(10);
+        targets = eval_at_secure_point_benches);
+criterion_main!(benches);