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