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 backward_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
 17    a.iter_mut().for_each(|r| nttop.backward(r.as_mut_slice()));
 18}
 19
 20fn backward_lazy_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
 21    a.iter_mut()
 22        .for_each(|r| nttop.backward_lazy(r.as_mut_slice()));
 23}
 24
 25fn benchmark(c: &mut Criterion) {
 26    let mut group = c.benchmark_group("ntt");
 27    // 55
 28    for prime in [36028797017456641] {
 29        for ring_size in [1 << 11] {
 30            let ntt = NttBackendU64::new(&prime, ring_size);
 31            let mut rng = thread_rng();
 32
 33            let a = (&mut rng)
 34                .sample_iter(Uniform::new(0, prime))
 35                .take(ring_size)
 36                .collect_vec();
 37            let d = 3;
 38            let a_matrix = (0..d)
 39                .map(|_| {
 40                    (&mut rng)
 41                        .sample_iter(Uniform::new(0, prime))
 42                        .take(ring_size)
 43                        .collect_vec()
 44                })
 45                .collect_vec();
 46
 47            {
 48                group.bench_function(
 49                    BenchmarkId::new("forward", format!("q={prime}/N={ring_size}")),
 50                    |b| {
 51                        b.iter_batched_ref(
 52                            || a.clone(),
 53                            |mut a| black_box(ntt.forward(&mut a)),
 54                            criterion::BatchSize::PerIteration,
 55                        )
 56                    },
 57                );
 58
 59                group.bench_function(
 60                    BenchmarkId::new("forward_lazy", format!("q={prime}/N={ring_size}")),
 61                    |b| {
 62                        b.iter_batched_ref(
 63                            || a.clone(),
 64                            |mut a| black_box(ntt.forward_lazy(&mut a)),
 65                            criterion::BatchSize::PerIteration,
 66                        )
 67                    },
 68                );
 69
 70                group.bench_function(
 71                    BenchmarkId::new("forward_matrix", format!("q={prime}/N={ring_size}/d={d}")),
 72                    |b| {
 73                        b.iter_batched_ref(
 74                            || a_matrix.clone(),
 75                            |a_matrix| black_box(forward_matrix(a_matrix, &ntt)),
 76                            criterion::BatchSize::PerIteration,
 77                        )
 78                    },
 79                );
 80
 81                group.bench_function(
 82                    BenchmarkId::new(
 83                        "forward_lazy_matrix",
 84                        format!("q={prime}/N={ring_size}/d={d}"),
 85                    ),
 86                    |b| {
 87                        b.iter_batched_ref(
 88                            || a_matrix.clone(),
 89                            |a_matrix| black_box(forward_lazy_matrix(a_matrix, &ntt)),
 90                            criterion::BatchSize::PerIteration,
 91                        )
 92                    },
 93                );
 94            }
 95
 96            {
 97                group.bench_function(
 98                    BenchmarkId::new("backward", format!("q={prime}/N={ring_size}")),
 99                    |b| {
100                        b.iter_batched_ref(
101                            || a.clone(),
102                            |mut a| black_box(ntt.backward(&mut a)),
103                            criterion::BatchSize::PerIteration,
104                        )
105                    },
106                );
107
108                group.bench_function(
109                    BenchmarkId::new("backward_lazy", format!("q={prime}/N={ring_size}")),
110                    |b| {
111                        b.iter_batched_ref(
112                            || a.clone(),
113                            |mut a| black_box(ntt.backward_lazy(&mut a)),
114                            criterion::BatchSize::PerIteration,
115                        )
116                    },
117                );
118
119                group.bench_function(
120                    BenchmarkId::new("backward_matrix", format!("q={prime}/N={ring_size}")),
121                    |b| {
122                        b.iter_batched_ref(
123                            || a_matrix.clone(),
124                            |a_matrix| black_box(backward_matrix(a_matrix, &ntt)),
125                            criterion::BatchSize::PerIteration,
126                        )
127                    },
128                );
129
130                group.bench_function(
131                    BenchmarkId::new(
132                        "backward_lazy_matrix",
133                        format!("q={prime}/N={ring_size}/d={d}"),
134                    ),
135                    |b| {
136                        b.iter_batched_ref(
137                            || a_matrix.clone(),
138                            |a_matrix| black_box(backward_lazy_matrix(a_matrix, &ntt)),
139                            criterion::BatchSize::PerIteration,
140                        )
141                    },
142                );
143            }
144        }
145    }
146
147    group.finish();
148}
149
150criterion_group!(ntt, benchmark);
151criterion_main!(ntt);