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()
 53                    ),
 54                ),
 55                |b| {
 56                    b.iter_batched_ref(
 57                        || {
 58                            (
 59                                a.clone(),
 60                                vec![vec![0u64; ring_size]; decomposer.decomposition_count()],
 61                            )
 62                        },
 63                        |(r, decomp_r)| (decompose_r(r, decomp_r, &decomposer)),
 64                        criterion::BatchSize::PerIteration,
 65                    )
 66                },
 67            );
 68        }
 69    }
 70
 71    group.finish();
 72}
 73
 74fn benchmark(c: &mut Criterion) {
 75    let mut group = c.benchmark_group("modulus");
 76    // 55
 77    for prime in [36028797017456641] {
 78        for ring_size in [1 << 11] {
 79            let modop = ModularOpsU64::new(prime);
 80
 81            let mut rng = thread_rng();
 82            let dist = Uniform::new(0, prime);
 83
 84            let a0 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 85            let a1 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 86            let a2 = (&mut rng).sample_iter(dist).take(ring_size).collect_vec();
 87
 88            let d = 2;
 89            let a0_matrix = (0..d)
 90                .into_iter()
 91                .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
 92                .collect_vec();
 93            // a0 in shoup representation
 94            let a0_shoup_matrix = a0_matrix
 95                .iter()
 96                .map(|r| {
 97                    r.iter()
 98                        .map(|v| {
 99                            // $(v * 2^{\beta}) / p$
100                            ((*v as u128 * (1u128 << 64)) / prime as u128) as u64
101                        })
102                        .collect_vec()
103                })
104                .collect_vec();
105            let a1_matrix = (0..d)
106                .into_iter()
107                .map(|_| (&mut rng).sample_iter(dist).take(ring_size).collect_vec())
108                .collect_vec();
109
110            group.bench_function(
111                BenchmarkId::new("matrix_fma_lazy", format!("q={prime}/N={ring_size}/d={d}")),
112                |b| {
113                    b.iter_batched_ref(
114                        || (vec![0u64; ring_size]),
115                        |(out)| black_box(matrix_fma(out, &a0_matrix, &a1_matrix, &modop)),
116                        criterion::BatchSize::PerIteration,
117                    )
118                },
119            );
120
121            group.bench_function(
122                BenchmarkId::new(
123                    "matrix_shoup_fma_lazy",
124                    format!("q={prime}/N={ring_size}/d={d}"),
125                ),
126                |b| {
127                    b.iter_batched_ref(
128                        || (vec![0u64; ring_size]),
129                        |(out)| {
130                            black_box(modop.shoup_matrix_fma(
131                                out,
132                                &a0_matrix,
133                                &a0_shoup_matrix,
134                                &a1_matrix,
135                            ))
136                        },
137                        criterion::BatchSize::PerIteration,
138                    )
139                },
140            );
141        }
142    }
143
144    group.finish();
145}
146
147criterion_group!(decomposer, benchmark_decomposer);
148criterion_group!(modulus, benchmark);
149criterion_main!(modulus, decomposer);