1use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
  2use math::modulus::WordOps;
  3use math::poly::Poly;
  4use math::ring::Ring;
  5
  6fn ntt(c: &mut Criterion) {
  7    fn runner<'a, const INPLACE: bool, const LAZY: bool>(
  8        ring: &'a Ring<u64>,
  9    ) -> Box<dyn FnMut() + 'a> {
 10        let mut a: Poly<u64> = ring.new_poly();
 11        for i in 0..a.n() {
 12            a.0[i] = i as u64;
 13        }
 14        if INPLACE {
 15            Box::new(move || ring.ntt_inplace::<LAZY>(&mut a))
 16        } else {
 17            let mut b: Poly<u64> = ring.new_poly();
 18            Box::new(move || ring.ntt::<LAZY>(&a, &mut b))
 19        }
 20    }
 21
 22    let q: u64 = 0x1fffffffffe00001u64;
 23
 24    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 25        c.benchmark_group("ntt");
 26
 27    for log_n in 10..17 {
 28        let ring = Ring::new(1 << log_n, q, 1);
 29
 30        let runners: [(String, Box<dyn FnMut()>); 4] = [
 31            (format!("inplace=true/LAZY=true/q={}", q.log2()), {
 32                runner::<true, true>(&ring)
 33            }),
 34            (format!("inplace=true/LAZY=false/q={}", q.log2()), {
 35                runner::<true, false>(&ring)
 36            }),
 37            (format!("inplace=false/LAZY=true/q={}", q.log2()), {
 38                runner::<false, true>(&ring)
 39            }),
 40            (format!("inplace=false/LAZY=false/q={}", q.log2()), {
 41                runner::<false, false>(&ring)
 42            }),
 43        ];
 44
 45        for (name, mut runner) in runners {
 46            let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
 47            b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
 48                b.iter(&mut runner)
 49            });
 50        }
 51    }
 52}
 53
 54fn intt(c: &mut Criterion) {
 55    fn runner<'a, const INPLACE: bool, const LAZY: bool>(
 56        ring: &'a Ring<u64>,
 57    ) -> Box<dyn FnMut() + 'a> {
 58        let mut a: Poly<u64> = ring.new_poly();
 59        for i in 0..a.n() {
 60            a.0[i] = i as u64;
 61        }
 62        if INPLACE {
 63            Box::new(move || ring.intt_inplace::<LAZY>(&mut a))
 64        } else {
 65            let mut b: Poly<u64> = ring.new_poly();
 66            Box::new(move || ring.intt::<LAZY>(&a, &mut b))
 67        }
 68    }
 69
 70    let q: u64 = 0x1fffffffffe00001u64;
 71
 72    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 73        c.benchmark_group("intt");
 74
 75    for log_n in 10..17 {
 76        let ring = Ring::new(1 << log_n, q, 1);
 77
 78        let runners: [(String, Box<dyn FnMut()>); 4] = [
 79            (format!("inplace=true/LAZY=true/q={}", q.log2()), {
 80                runner::<true, true>(&ring)
 81            }),
 82            (format!("inplace=true/LAZY=false/q={}", q.log2()), {
 83                runner::<true, false>(&ring)
 84            }),
 85            (format!("inplace=false/LAZY=true/q={}", q.log2()), {
 86                runner::<false, true>(&ring)
 87            }),
 88            (format!("inplace=false/LAZY=false/q={}", q.log2()), {
 89                runner::<false, false>(&ring)
 90            }),
 91        ];
 92
 93        for (name, mut runner) in runners {
 94            let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
 95            b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
 96                b.iter(&mut runner)
 97            });
 98        }
 99    }
100}
101
102criterion_group!(benches, ntt, intt,);
103criterion_main!(benches);