add more NTT benches

This commit is contained in:
Janmajaya Mall
2024-06-10 12:32:29 +05:30
parent a292f9ad9e
commit 085d5e23a0
2 changed files with 138 additions and 28 deletions

View File

@@ -4,20 +4,39 @@ use itertools::Itertools;
use rand::{thread_rng, Rng};
use rand_distr::Uniform;
fn forward_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
a.iter_mut().for_each(|r| nttop.forward(r.as_mut_slice()));
}
fn forward_lazy_matrix(a: &mut [Vec<u64>], nttop: &NttBackendU64) {
a.iter_mut()
.for_each(|r| nttop.forward_lazy(r.as_mut_slice()));
}
fn benchmark(c: &mut Criterion) {
let mut group = c.benchmark_group("ntt");
// 55
for prime in [36028797017456641] {
for ring_size in [1 << 11, 1 << 15] {
for ring_size in [1 << 11] {
let ntt = NttBackendU64::new(&prime, ring_size);
let mut rng = thread_rng();
let a = rng
let a = (&mut rng)
.sample_iter(Uniform::new(0, prime))
.take(ring_size)
.collect_vec();
let d = 3;
let a_matrix = (0..d)
.map(|_| {
(&mut rng)
.sample_iter(Uniform::new(0, prime))
.take(ring_size)
.collect_vec()
})
.collect_vec();
group.bench_function(
BenchmarkId::new("forward", format!("q={prime}/{ring_size}")),
BenchmarkId::new("forward", format!("q={prime}/N={ring_size}")),
|b| {
b.iter_batched_ref(
|| a.clone(),
@@ -26,6 +45,42 @@ fn benchmark(c: &mut Criterion) {
)
},
);
group.bench_function(
BenchmarkId::new("forward_lazy", format!("q={prime}/N={ring_size}")),
|b| {
b.iter_batched_ref(
|| a.clone(),
|mut a| black_box(ntt.forward_lazy(&mut a)),
criterion::BatchSize::PerIteration,
)
},
);
group.bench_function(
BenchmarkId::new("forward_matrix", format!("q={prime}/N={ring_size}/d={d}")),
|b| {
b.iter_batched_ref(
|| a_matrix.clone(),
|a_matrix| black_box(forward_matrix(a_matrix, &ntt)),
criterion::BatchSize::PerIteration,
)
},
);
group.bench_function(
BenchmarkId::new(
"forward_lazy_matrix",
format!("q={prime}/N={ring_size}/d={d}"),
),
|b| {
b.iter_batched_ref(
|| a_matrix.clone(),
|a_matrix| black_box(forward_lazy_matrix(a_matrix, &ntt)),
criterion::BatchSize::PerIteration,
)
},
);
}
}