1use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
  2use math::dft::DFT;
  3use math::{dft::ntt::Table, modulus::prime::Prime};
  4
  5fn forward_inplace(c: &mut Criterion) {
  6    fn runner(prime_instance: Prime<u64>, nth_root: usize) -> Box<dyn FnMut()> {
  7        let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
  8        let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
  9        for i in 0..a.len() {
 10            a[i] = i as u64;
 11        }
 12        Box::new(move || ntt_table.forward_inplace::<false>(&mut a))
 13    }
 14
 15    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 16        c.benchmark_group("forward_inplace");
 17    for log_nth_root in 11..18 {
 18        let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
 19
 20        let runners = [("prime", { runner(prime_instance, 1 << log_nth_root) })];
 21        for (name, mut runner) in runners {
 22            let id = BenchmarkId::new(name, 1 << (log_nth_root - 1));
 23            b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 24        }
 25    }
 26}
 27
 28fn forward_inplace_lazy(c: &mut Criterion) {
 29    fn runner(prime_instance: Prime<u64>, nth_root: usize) -> Box<dyn FnMut()> {
 30        let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
 31        let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
 32        for i in 0..a.len() {
 33            a[i] = i as u64;
 34        }
 35        Box::new(move || ntt_table.forward_inplace_lazy(&mut a))
 36    }
 37
 38    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 39        c.benchmark_group("forward_inplace_lazy");
 40    for log_nth_root in 11..17 {
 41        let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
 42
 43        let runners = [("prime", { runner(prime_instance, 1 << log_nth_root) })];
 44        for (name, mut runner) in runners {
 45            let id = BenchmarkId::new(name, 1 << (log_nth_root - 1));
 46            b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 47        }
 48    }
 49}
 50
 51fn backward_inplace(c: &mut Criterion) {
 52    fn runner(prime_instance: Prime<u64>, nth_root: usize) -> Box<dyn FnMut()> {
 53        let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
 54        let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
 55        for i in 0..a.len() {
 56            a[i] = i as u64;
 57        }
 58        Box::new(move || ntt_table.backward_inplace::<false>(&mut a))
 59    }
 60
 61    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 62        c.benchmark_group("backward_inplace");
 63    for log_nth_root in 11..18 {
 64        let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
 65
 66        let runners = [("prime", { runner(prime_instance, 1 << log_nth_root) })];
 67        for (name, mut runner) in runners {
 68            let id = BenchmarkId::new(name, 1 << (log_nth_root - 1));
 69            b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 70        }
 71    }
 72}
 73
 74fn backward_inplace_lazy(c: &mut Criterion) {
 75    fn runner(prime_instance: Prime<u64>, nth_root: usize) -> Box<dyn FnMut()> {
 76        let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
 77        let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
 78        for i in 0..a.len() {
 79            a[i] = i as u64;
 80        }
 81        Box::new(move || ntt_table.backward_inplace::<true>(&mut a))
 82    }
 83
 84    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
 85        c.benchmark_group("backward_inplace_lazy");
 86    for log_nth_root in 11..17 {
 87        let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
 88
 89        let runners = [("prime", { runner(prime_instance, 1 << log_nth_root) })];
 90        for (name, mut runner) in runners {
 91            let id = BenchmarkId::new(name, 1 << (log_nth_root - 1));
 92            b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 93        }
 94    }
 95}
 96
 97criterion_group!(
 98    benches,
 99    forward_inplace,
100    forward_inplace_lazy,
101    backward_inplace,
102    backward_inplace_lazy
103);
104criterion_main!(benches);