1use bin_rs::{Ntt, NttBackendU64, NttInit};
  2use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
  3use itertools::Itertools;
  4use rand::{thread_rng, Rng};
  5use rand_distr::Uniform;
  6
  7fn forward_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
  8    a.iter_mut().for_each(|r| nttop.forward(r.as_mut_slice()));
  9}
 10
 11fn forward_lazy_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
 12    a.iter_mut()
 13        .for_each(|r| nttop.forward_lazy(r.as_mut_slice()));
 14}
 15
 16fn benchmark(c: &mut Criterion) {
 17    let mut group = c.benchmark_group("ntt");
 18    // 55
 19    for prime in [36028797017456641] {
 20        for ring_size in [1 << 11] {
 21            let ntt = NttBackendU64::new(&prime, ring_size);
 22            let mut rng = thread_rng();
 23
 24            let a = (&mut rng)
 25                .sample_iter(Uniform::new(0, prime))
 26                .take(ring_size)
 27                .collect_vec();
 28            let d = 3;
 29            let a_matrix = (0..d)
 30                .map(|_| {
 31                    (&mut rng)
 32                        .sample_iter(Uniform::new(0, prime))
 33                        .take(ring_size)
 34                        .collect_vec()
 35                })
 36                .collect_vec();
 37
 38            {
 39                group.bench_function(
 40                    BenchmarkId::new("forward", format!("q={prime}/N={ring_size}")),
 41                    |b| {
 42                        b.iter_batched_ref(
 43                            || a.clone(),
 44                            |mut a| black_box(ntt.forward(&mut a)),
 45                            criterion::BatchSize::PerIteration,
 46                        )
 47                    },
 48                );
 49
 50                group.bench_function(
 51                    BenchmarkId::new("forward_lazy", format!("q={prime}/N={ring_size}")),
 52                    |b| {
 53                        b.iter_batched_ref(
 54                            || a.clone(),
 55                            |mut a| black_box(ntt.forward_lazy(&mut a)),
 56                            criterion::BatchSize::PerIteration,
 57                        )
 58                    },
 59                );
 60
 61                group.bench_function(
 62                    BenchmarkId::new("forward_matrix", format!("q={prime}/N={ring_size}/d={d}")),
 63                    |b| {
 64                        b.iter_batched_ref(
 65                            || a_matrix.clone(),
 66                            |a_matrix| black_box(forward_matrix(a_matrix, &ntt)),
 67                            criterion::BatchSize::PerIteration,
 68                        )
 69                    },
 70                );
 71
 72                group.bench_function(
 73                    BenchmarkId::new(
 74                        "forward_lazy_matrix",
 75                        format!("q={prime}/N={ring_size}/d={d}"),
 76                    ),
 77                    |b| {
 78                        b.iter_batched_ref(
 79                            || a_matrix.clone(),
 80                            |a_matrix| black_box(forward_lazy_matrix(a_matrix, &ntt)),
 81                            criterion::BatchSize::PerIteration,
 82                        )
 83                    },
 84                );
 85            }
 86
 87            {
 88                group.bench_function(
 89                    BenchmarkId::new("backward_lazy", format!("q={prime}/N={ring_size}")),
 90                    |b| {
 91                        b.iter_batched_ref(
 92                            || a.clone(),
 93                            |mut a| black_box(ntt.backward_lazy(&mut a)),
 94                            criterion::BatchSize::PerIteration,
 95                        )
 96                    },
 97                );
 98            }
 99        }
100    }
101
102    group.finish();
103}
104
105criterion_group!(ntt, benchmark);
106criterion_main!(ntt);