1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4
  5use crate::{
  6    api::{
  7        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAutomorphism, VecZnxAutomorphismInplace,
  8        VecZnxAutomorphismInplaceTmpBytes,
  9    },
 10    layouts::{Backend, FillUniform, Module, ScratchOwned, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
 11    reference::znx::{ZnxAutomorphism, ZnxCopy, ZnxZero},
 12    source::Source,
 13};
 14
 15pub fn vec_znx_automorphism_inplace_tmp_bytes(n: usize) -> usize {
 16    n * size_of::<i64>()
 17}
 18
 19pub fn vec_znx_automorphism<R, A, ZNXARI>(p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
 20where
 21    R: VecZnxToMut,
 22    A: VecZnxToRef,
 23    ZNXARI: ZnxAutomorphism + ZnxZero,
 24{
 25    let a: VecZnx<&[u8]> = a.to_ref();
 26    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 27
 28    #[cfg(debug_assertions)]
 29    {
 30        use crate::layouts::ZnxInfos;
 31
 32        assert_eq!(a.n(), res.n());
 33    }
 34
 35    let min_size: usize = res.size().min(a.size());
 36
 37    for j in 0..min_size {
 38        ZNXARI::znx_automorphism(p, res.at_mut(res_col, j), a.at(a_col, j));
 39    }
 40
 41    for j in min_size..res.size() {
 42        ZNXARI::znx_zero(res.at_mut(res_col, j));
 43    }
 44}
 45
 46pub fn vec_znx_automorphism_inplace<R, ZNXARI>(p: i64, res: &mut R, res_col: usize, tmp: &mut [i64])
 47where
 48    R: VecZnxToMut,
 49    ZNXARI: ZnxAutomorphism + ZnxCopy,
 50{
 51    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 52    #[cfg(debug_assertions)]
 53    {
 54        assert_eq!(res.n(), tmp.len());
 55    }
 56    for j in 0..res.size() {
 57        ZNXARI::znx_automorphism(p, tmp, res.at(res_col, j));
 58        ZNXARI::znx_copy(res.at_mut(res_col, j), tmp);
 59    }
 60}
 61
 62pub fn bench_vec_znx_automorphism<B: Backend>(c: &mut Criterion, label: &str)
 63where
 64    Module<B>: VecZnxAutomorphism + ModuleNew<B>,
 65{
 66    let group_name: String = format!("vec_znx_automorphism::{}", label);
 67
 68    let mut group = c.benchmark_group(group_name);
 69
 70    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
 71    where
 72        Module<B>: VecZnxAutomorphism + ModuleNew<B>,
 73    {
 74        let n: usize = 1 << params[0];
 75        let cols: usize = params[1];
 76        let size: usize = params[2];
 77
 78        let module: Module<B> = Module::<B>::new(n as u64);
 79
 80        let mut source: Source = Source::new([0u8; 32]);
 81
 82        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
 83        let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
 84
 85        // Fill a with random i64
 86        a.fill_uniform(50, &mut source);
 87        res.fill_uniform(50, &mut source);
 88
 89        move || {
 90            for i in 0..cols {
 91                module.vec_znx_automorphism(-7, &mut res, i, &a, i);
 92            }
 93            black_box(());
 94        }
 95    }
 96
 97    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
 98        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
 99        let mut runner = runner::<B>(params);
100        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
101    }
102
103    group.finish();
104}
105
106pub fn bench_vec_znx_automorphism_inplace<B: Backend>(c: &mut Criterion, label: &str)
107where
108    Module<B>: VecZnxAutomorphismInplace<B> + VecZnxAutomorphismInplaceTmpBytes + ModuleNew<B>,
109    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
110{
111    let group_name: String = format!("vec_znx_automorphism_inplace::{}", label);
112
113    let mut group = c.benchmark_group(group_name);
114
115    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
116    where
117        Module<B>: VecZnxAutomorphismInplace<B> + ModuleNew<B> + VecZnxAutomorphismInplaceTmpBytes,
118        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
119    {
120        let n: usize = 1 << params[0];
121        let cols: usize = params[1];
122        let size: usize = params[2];
123
124        let module: Module<B> = Module::<B>::new(n as u64);
125
126        let mut source: Source = Source::new([0u8; 32]);
127
128        let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
129
130        let mut scratch = ScratchOwned::alloc(module.vec_znx_automorphism_inplace_tmp_bytes());
131
132        // Fill a with random i64
133        res.fill_uniform(50, &mut source);
134
135        move || {
136            for i in 0..cols {
137                module.vec_znx_automorphism_inplace(-7, &mut res, i, scratch.borrow());
138            }
139            black_box(());
140        }
141    }
142
143    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
144        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
145        let mut runner = runner::<B>(params);
146        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
147    }
148
149    group.finish();
150}