1use bin_rs::{Decomposer, DefaultDecomposer, ModInit, ModularOpsU64, VectorOps};
  2use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
  3use itertools::{izip, Itertools};
  4use rand::{thread_rng, Rng};
  5use rand_distr::Uniform;
  6
  7pub(crate) fn decompose_r(
  8    r: &[u64],
  9    decomp_r: &mut [Vec<u64>],
 10    decomposer: &DefaultDecomposer<u64>,
 11) {
 12    let ring_size = r.len();
 13    // let d = decomposer.decomposition_count();
 14    // let mut count = 0;
 15    for ri in 0..ring_size {
 16        // let el_decomposed = decomposer.decompose(&r[ri]);
 17        decomposer
 18            .decompose_iter(&r[ri])
 19            .enumerate()
 20            .into_iter()
 21            .for_each(|(j, el)| {
 22                decomp_r[j][ri] = el;
 23            });
 24    }
 25}
 26
 27fn benchmark_decomposer(c: &mut Criterion) {
 28    let mut group = c.benchmark_group("decomposer");
 29
 30    // let decomposers = vec![];
 31    // 55
 32    for prime in [36028797017456641] {
 33        for ring_size in [1 << 11] {
 34            let logb = 11;
 35            let decomposer = DefaultDecomposer::new(prime, logb, 2);
 36
 37            let mut rng = thread_rng();
 38            let dist = Uniform::new(0, prime);
 39            let a = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 40
 41            group.bench_function(
 42                BenchmarkId::new(
 43                    "decompose",
 44                    format!(
 45                        "q={prime}/N={ring_size}/logB={logb}/d={}",
 46                        decomposer.decomposition_count()
 47                    ),
 48                ),
 49                |b| {
 50                    b.iter_batched_ref(
 51                        || {
 52                            (
 53                                a.clone(),
 54                                vec![vec![0u64; ring_size]; decomposer.decomposition_count()],
 55                            )
 56                        },
 57                        |(r, decomp_r)| (decompose_r(r, decomp_r, &decomposer)),
 58                        criterion::BatchSize::PerIteration,
 59                    )
 60                },
 61            );
 62        }
 63    }
 64
 65    group.finish();
 66}
 67
 68fn benchmark(c: &mut Criterion) {
 69    let mut group = c.benchmark_group("modulus");
 70    // 55
 71    for prime in [36028797017456641] {
 72        for ring_size in [1 << 11, 1 << 15] {
 73            let modop = ModularOpsU64::new(prime);
 74
 75            let mut rng = thread_rng();
 76            let dist = Uniform::new(0, prime);
 77
 78            let a0 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 79            let a1 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 80            let a2 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 81
 82            group.bench_function(
 83                BenchmarkId::new("elwise_fma", format!("q={prime}/{ring_size}")),
 84                |b| {
 85                    b.iter_batched_ref(
 86                        || (a0.clone(), a1.clone(), a2.clone()),
 87                        |(a0, a1, a2)| black_box(modop.elwise_fma_mut(a0, a1, a2)),
 88                        criterion::BatchSize::PerIteration,
 89                    )
 90                },
 91            );
 92        }
 93    }
 94
 95    group.finish();
 96}
 97
 98criterion_group!(decomposer, benchmark_decomposer);
 99criterion_group!(modulus, benchmark);
100criterion_main!(modulus, decomposer);