1use bin_rs::{
  2    ArithmeticLazyOps, ArithmeticOps, Decomposer, DefaultDecomposer, ModInit, ModularOpsU64,
  3    ShoupMatrixFMA, VectorOps,
  4};
  5use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
  6use itertools::{izip, Itertools};
  7use rand::{thread_rng, Rng};
  8use rand_distr::Uniform;
  9
 10fn decompose_r(r: &[u64], decomp_r: &mut [Vec<u64>], decomposer: &DefaultDecomposer<u64>) {
 11    let ring_size = r.len();
 12    // let d = decomposer.decomposition_count();
 13    // let mut count = 0;
 14    for ri in 0..ring_size {
 15        // let el_decomposed = decomposer.decompose(&r[ri]);
 16        decomposer
 17            .decompose_iter(&r[ri])
 18            .enumerate()
 19            .into_iter()
 20            .for_each(|(j, el)| {
 21                decomp_r[j][ri] = el;
 22            });
 23    }
 24}
 25
 26fn matrix_fma(out: &mut [u64], a: &Vec<Vec<u64>>, b: &Vec<Vec<u64>>, modop: &ModularOpsU64<u64>) {
 27    izip!(a.iter(), b.iter()).for_each(|(a_r, b_r)| {
 28        izip!(out.iter_mut(), a_r.iter(), b_r.iter())
 29            .for_each(|(o, ai, bi)| *o = modop.add_lazy(o, &modop.mul_lazy(ai, bi)));
 30    });
 31}
 32
 33fn benchmark_decomposer(c: &mut Criterion) {
 34    let mut group = c.benchmark_group("decomposer");
 35
 36    // let decomposers = vec![];
 37    // 55
 38    for prime in [36028797017456641] {
 39        for ring_size in [1 << 11] {
 40            let logb = 11;
 41            let decomposer = DefaultDecomposer::new(prime, logb, 2);
 42
 43            let mut rng = thread_rng();
 44            let dist = Uniform::new(0, prime);
 45            let a = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 46
 47            group.bench_function(
 48                BenchmarkId::new(
 49                    "decompose",
 50                    format!(
 51                        "q={prime}/N={ring_size}/logB={logb}/d={}",
 52                        *decomposer.decomposition_count().as_ref()
 53                    ),
 54                ),
 55                |b| {
 56                    b.iter_batched_ref(
 57                        || {
 58                            (
 59                                a.clone(),
 60                                vec![
 61                                    vec![0u64; ring_size];
 62                                    *decomposer.decomposition_count().as_ref()
 63                                ],
 64                            )
 65                        },
 66                        |(r, decomp_r)| (decompose_r(r, decomp_r, &decomposer)),
 67                        criterion::BatchSize::PerIteration,
 68                    )
 69                },
 70            );
 71        }
 72    }
 73
 74    group.finish();
 75}
 76
 77fn benchmark(c: &mut Criterion) {
 78    let mut group = c.benchmark_group("modulus");
 79    // 55
 80    for prime in [36028797017456641] {
 81        for ring_size in [1 << 11] {
 82            let modop = ModularOpsU64::new(prime);
 83
 84            let mut rng = thread_rng();
 85            let dist = Uniform::new(0, prime);
 86
 87            let a0 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 88            let a1 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 89            let a2 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 90
 91            let d = 1;
 92            let a0_matrix = (0..d)
 93                .into_iter()
 94                .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
 95                .collect_vec();
 96            // a0 in shoup representation
 97            let a0_shoup_matrix = a0_matrix
 98                .iter()
 99                .map(|r| {
100                    r.iter()
101                        .map(|v| {
102                            // $(v * 2^{\beta}) / p$
103                            ((*v as u128 * (1u128 << 64)) / prime as u128) as u64
104                        })
105                        .collect_vec()
106                })
107                .collect_vec();
108            let a1_matrix = (0..d)
109                .into_iter()
110                .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
111                .collect_vec();
112
113            group.bench_function(
114                BenchmarkId::new("matrix_fma_lazy", format!("q={prime}/N={ring_size}/d={d}")),
115                |b| {
116                    b.iter_batched_ref(
117                        || (vec![0u64; ring_size]),
118                        |(out)| black_box(matrix_fma(out, &a0_matrix, &a1_matrix, &modop)),
119                        criterion::BatchSize::PerIteration,
120                    )
121                },
122            );
123
124            group.bench_function(
125                BenchmarkId::new(
126                    "matrix_shoup_fma_lazy",
127                    format!("q={prime}/N={ring_size}/d={d}"),
128                ),
129                |b| {
130                    b.iter_batched_ref(
131                        || (vec![0u64; ring_size]),
132                        |(out)| {
133                            black_box(modop.shoup_matrix_fma(
134                                out,
135                                &a0_matrix,
136                                &a0_shoup_matrix,
137                                &a1_matrix,
138                            ))
139                        },
140                        criterion::BatchSize::PerIteration,
141                    )
142                },
143            );
144        }
145    }
146
147    group.finish();
148}
149
150criterion_group!(decomposer, benchmark_decomposer);
151criterion_group!(modulus, benchmark);
152criterion_main!(modulus, decomposer);