1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4
  5use crate::{
  6    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
  7    layouts::{Backend, FillUniform, Module, ScratchOwned, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
  8    reference::znx::{
  9        ZnxNormalizeFinalStep, ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStep, ZnxNormalizeFirstStepCarryOnly,
 10        ZnxNormalizeFirstStepInplace, ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly, ZnxNormalizeMiddleStepInplace,
 11        ZnxZero,
 12    },
 13    source::Source,
 14};
 15
 16pub fn vec_znx_normalize_tmp_bytes(n: usize) -> usize {
 17    n * size_of::<i64>()
 18}
 19
 20pub fn vec_znx_normalize<R, A, ZNXARI>(basek: usize, res: &mut R, res_col: usize, a: &A, a_col: usize, carry: &mut [i64])
 21where
 22    R: VecZnxToMut,
 23    A: VecZnxToRef,
 24    ZNXARI: ZnxZero
 25        + ZnxNormalizeFirstStepCarryOnly
 26        + ZnxNormalizeMiddleStepCarryOnly
 27        + ZnxNormalizeMiddleStep
 28        + ZnxNormalizeFinalStep
 29        + ZnxNormalizeFirstStep,
 30{
 31    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 32    let a: VecZnx<&[u8]> = a.to_ref();
 33
 34    #[cfg(debug_assertions)]
 35    {
 36        assert!(carry.len() >= res.n());
 37    }
 38
 39    let res_size: usize = res.size();
 40    let a_size = a.size();
 41
 42    if a_size > res_size {
 43        for j in (res_size..a_size).rev() {
 44            if j == a_size - 1 {
 45                ZNXARI::znx_normalize_first_step_carry_only(basek, 0, a.at(a_col, j), carry);
 46            } else {
 47                ZNXARI::znx_normalize_middle_step_carry_only(basek, 0, a.at(a_col, j), carry);
 48            }
 49        }
 50
 51        for j in (1..res_size).rev() {
 52            ZNXARI::znx_normalize_middle_step(basek, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
 53        }
 54
 55        ZNXARI::znx_normalize_final_step(basek, 0, res.at_mut(res_col, 0), a.at(a_col, 0), carry);
 56    } else {
 57        for j in (0..a_size).rev() {
 58            if j == a_size - 1 {
 59                ZNXARI::znx_normalize_first_step(basek, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
 60            } else if j == 0 {
 61                ZNXARI::znx_normalize_final_step(basek, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
 62            } else {
 63                ZNXARI::znx_normalize_middle_step(basek, 0, res.at_mut(res_col, j), a.at(a_col, j), carry);
 64            }
 65        }
 66
 67        for j in a_size..res_size {
 68            ZNXARI::znx_zero(res.at_mut(res_col, j));
 69        }
 70    }
 71}
 72
 73pub fn vec_znx_normalize_inplace<R: VecZnxToMut, ZNXARI>(basek: usize, res: &mut R, res_col: usize, carry: &mut [i64])
 74where
 75    ZNXARI: ZnxNormalizeFirstStepInplace + ZnxNormalizeMiddleStepInplace + ZnxNormalizeFinalStepInplace,
 76{
 77    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 78
 79    #[cfg(debug_assertions)]
 80    {
 81        assert!(carry.len() >= res.n());
 82    }
 83
 84    let res_size: usize = res.size();
 85
 86    for j in (0..res_size).rev() {
 87        if j == res_size - 1 {
 88            ZNXARI::znx_normalize_first_step_inplace(basek, 0, res.at_mut(res_col, j), carry);
 89        } else if j == 0 {
 90            ZNXARI::znx_normalize_final_step_inplace(basek, 0, res.at_mut(res_col, j), carry);
 91        } else {
 92            ZNXARI::znx_normalize_middle_step_inplace(basek, 0, res.at_mut(res_col, j), carry);
 93        }
 94    }
 95}
 96
 97pub fn bench_vec_znx_normalize<B: Backend>(c: &mut Criterion, label: &str)
 98where
 99    Module<B>: VecZnxNormalize<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
100    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
101{
102    let group_name: String = format!("vec_znx_normalize::{}", label);
103
104    let mut group = c.benchmark_group(group_name);
105
106    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
107    where
108        Module<B>: VecZnxNormalize<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
109        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
110    {
111        let n: usize = 1 << params[0];
112        let cols: usize = params[1];
113        let size: usize = params[2];
114
115        let module: Module<B> = Module::<B>::new(n as u64);
116
117        let basek: usize = 50;
118
119        let mut source: Source = Source::new([0u8; 32]);
120
121        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
122        let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
123
124        // Fill a with random i64
125        a.fill_uniform(50, &mut source);
126        res.fill_uniform(50, &mut source);
127
128        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
129
130        move || {
131            for i in 0..cols {
132                module.vec_znx_normalize(basek, &mut res, i, &a, i, scratch.borrow());
133            }
134            black_box(());
135        }
136    }
137
138    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
139        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
140        let mut runner = runner::<B>(params);
141        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
142    }
143
144    group.finish();
145}
146
147pub fn bench_vec_znx_normalize_inplace<B: Backend>(c: &mut Criterion, label: &str)
148where
149    Module<B>: VecZnxNormalizeInplace<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
150    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
151{
152    let group_name: String = format!("vec_znx_normalize_inplace::{}", label);
153
154    let mut group = c.benchmark_group(group_name);
155
156    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
157    where
158        Module<B>: VecZnxNormalizeInplace<B> + ModuleNew<B> + VecZnxNormalizeTmpBytes,
159        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
160    {
161        let n: usize = 1 << params[0];
162        let cols: usize = params[1];
163        let size: usize = params[2];
164
165        let module: Module<B> = Module::<B>::new(n as u64);
166
167        let basek: usize = 50;
168
169        let mut source: Source = Source::new([0u8; 32]);
170
171        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
172
173        // Fill a with random i64
174        a.fill_uniform(50, &mut source);
175
176        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
177
178        move || {
179            for i in 0..cols {
180                module.vec_znx_normalize_inplace(basek, &mut a, i, scratch.borrow());
181            }
182            black_box(());
183        }
184    }
185
186    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
187        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
188        let mut runner = runner::<B>(params);
189        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
190    }
191
192    group.finish();
193}