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            group.bench_function(
39                BenchmarkId::new("forward", format!("q={prime}/N={ring_size}")),
40                |b| {
41                    b.iter_batched_ref(
42                        || a.clone(),
43                        |mut a| black_box(ntt.forward(&mut a)),
44                        criterion::BatchSize::PerIteration,
45                    )
46                },
47            );
48
49            group.bench_function(
50                BenchmarkId::new("forward_lazy", format!("q={prime}/N={ring_size}")),
51                |b| {
52                    b.iter_batched_ref(
53                        || a.clone(),
54                        |mut a| black_box(ntt.forward_lazy(&mut a)),
55                        criterion::BatchSize::PerIteration,
56                    )
57                },
58            );
59
60            group.bench_function(
61                BenchmarkId::new("forward_matrix", format!("q={prime}/N={ring_size}/d={d}")),
62                |b| {
63                    b.iter_batched_ref(
64                        || a_matrix.clone(),
65                        |a_matrix| black_box(forward_matrix(a_matrix, &ntt)),
66                        criterion::BatchSize::PerIteration,
67                    )
68                },
69            );
70
71            group.bench_function(
72                BenchmarkId::new(
73                    "forward_lazy_matrix",
74                    format!("q={prime}/N={ring_size}/d={d}"),
75                ),
76                |b| {
77                    b.iter_batched_ref(
78                        || a_matrix.clone(),
79                        |a_matrix| black_box(forward_lazy_matrix(a_matrix, &ntt)),
80                        criterion::BatchSize::PerIteration,
81                    )
82                },
83            );
84        }
85    }
86
87    group.finish();
88}
89
90criterion_group!(ntt, benchmark);
91criterion_main!(ntt);