1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4
  5use crate::{
  6    api::{ModuleNew, VecZnxSub, VecZnxSubABInplace, VecZnxSubBAInplace},
  7    layouts::{Backend, FillUniform, Module, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
  8    oep::{ModuleNewImpl, VecZnxSubABInplaceImpl, VecZnxSubBAInplaceImpl, VecZnxSubImpl},
  9    reference::znx::{ZnxCopy, ZnxNegate, ZnxNegateInplace, ZnxSub, ZnxSubABInplace, ZnxSubBAInplace, ZnxZero},
 10    source::Source,
 11};
 12
 13pub fn vec_znx_sub<R, A, B, ZNXARI>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 14where
 15    R: VecZnxToMut,
 16    A: VecZnxToRef,
 17    B: VecZnxToRef,
 18    ZNXARI: ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
 19{
 20    let a: VecZnx<&[u8]> = a.to_ref();
 21    let b: VecZnx<&[u8]> = b.to_ref();
 22    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 23
 24    #[cfg(debug_assertions)]
 25    {
 26        assert_eq!(a.n(), res.n());
 27        assert_eq!(b.n(), res.n());
 28    }
 29
 30    let res_size: usize = res.size();
 31    let a_size: usize = a.size();
 32    let b_size: usize = b.size();
 33
 34    if a_size <= b_size {
 35        let sum_size: usize = a_size.min(res_size);
 36        let cpy_size: usize = b_size.min(res_size);
 37
 38        for j in 0..sum_size {
 39            ZNXARI::znx_sub(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
 40        }
 41
 42        for j in sum_size..cpy_size {
 43            ZNXARI::znx_negate(res.at_mut(res_col, j), b.at(b_col, j));
 44        }
 45
 46        for j in cpy_size..res_size {
 47            ZNXARI::znx_zero(res.at_mut(res_col, j));
 48        }
 49    } else {
 50        let sum_size: usize = b_size.min(res_size);
 51        let cpy_size: usize = a_size.min(res_size);
 52
 53        for j in 0..sum_size {
 54            ZNXARI::znx_sub(res.at_mut(res_col, j), a.at(a_col, j), b.at(b_col, j));
 55        }
 56
 57        for j in sum_size..cpy_size {
 58            ZNXARI::znx_copy(res.at_mut(res_col, j), a.at(a_col, j));
 59        }
 60
 61        for j in cpy_size..res_size {
 62            ZNXARI::znx_zero(res.at_mut(res_col, j));
 63        }
 64    }
 65}
 66
 67pub fn vec_znx_sub_ab_inplace<R, A, ZNXARI>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 68where
 69    R: VecZnxToMut,
 70    A: VecZnxToRef,
 71    ZNXARI: ZnxSubABInplace,
 72{
 73    let a: VecZnx<&[u8]> = a.to_ref();
 74    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 75
 76    #[cfg(debug_assertions)]
 77    {
 78        assert_eq!(a.n(), res.n());
 79    }
 80
 81    let res_size: usize = res.size();
 82    let a_size: usize = a.size();
 83
 84    let sum_size: usize = a_size.min(res_size);
 85
 86    for j in 0..sum_size {
 87        ZNXARI::znx_sub_ab_inplace(res.at_mut(res_col, j), a.at(a_col, j));
 88    }
 89}
 90
 91pub fn vec_znx_sub_ba_inplace<R, A, ZNXARI>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 92where
 93    R: VecZnxToMut,
 94    A: VecZnxToRef,
 95    ZNXARI: ZnxSubBAInplace + ZnxNegateInplace,
 96{
 97    let a: VecZnx<&[u8]> = a.to_ref();
 98    let mut res: VecZnx<&mut [u8]> = res.to_mut();
 99
100    #[cfg(debug_assertions)]
101    {
102        assert_eq!(a.n(), res.n());
103    }
104
105    let res_size: usize = res.size();
106    let a_size: usize = a.size();
107
108    let sum_size: usize = a_size.min(res_size);
109
110    for j in 0..sum_size {
111        ZNXARI::znx_sub_ba_inplace(res.at_mut(res_col, j), a.at(a_col, j));
112    }
113
114    for j in sum_size..res_size {
115        ZNXARI::znx_negate_inplace(res.at_mut(res_col, j));
116    }
117}
118
119pub fn bench_vec_znx_sub<B>(c: &mut Criterion, label: &str)
120where
121    B: Backend + ModuleNewImpl<B> + VecZnxSubImpl<B>,
122{
123    let group_name: String = format!("vec_znx_sub::{}", label);
124
125    let mut group = c.benchmark_group(group_name);
126
127    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
128    where
129        Module<B>: VecZnxSub + ModuleNew<B>,
130    {
131        let n: usize = 1 << params[0];
132        let cols: usize = params[1];
133        let size: usize = params[2];
134
135        let module: Module<B> = Module::<B>::new(n as u64);
136
137        let mut source: Source = Source::new([0u8; 32]);
138
139        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
140        let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
141        let mut c: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
142
143        // Fill a with random i64
144        a.fill_uniform(50, &mut source);
145        b.fill_uniform(50, &mut source);
146
147        move || {
148            for i in 0..cols {
149                module.vec_znx_sub(&mut c, i, &a, i, &b, i);
150            }
151            black_box(());
152        }
153    }
154
155    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
156        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
157        let mut runner = runner::<B>(params);
158        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
159    }
160
161    group.finish();
162}
163
164pub fn bench_vec_znx_sub_ab_inplace<B>(c: &mut Criterion, label: &str)
165where
166    B: Backend + ModuleNewImpl<B> + VecZnxSubABInplaceImpl<B>,
167{
168    let group_name: String = format!("vec_znx_sub_ab_inplace::{}", label);
169
170    let mut group = c.benchmark_group(group_name);
171
172    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
173    where
174        Module<B>: VecZnxSubABInplace + ModuleNew<B>,
175    {
176        let n: usize = 1 << params[0];
177        let cols: usize = params[1];
178        let size: usize = params[2];
179
180        let module: Module<B> = Module::<B>::new(n as u64);
181
182        let mut source: Source = Source::new([0u8; 32]);
183
184        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
185        let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
186
187        // Fill a with random i64
188        a.fill_uniform(50, &mut source);
189        b.fill_uniform(50, &mut source);
190
191        move || {
192            for i in 0..cols {
193                module.vec_znx_sub_ab_inplace(&mut b, i, &a, i);
194            }
195            black_box(());
196        }
197    }
198
199    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
200        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2]));
201        let mut runner = runner::<B>(params);
202        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
203    }
204
205    group.finish();
206}
207
208pub fn bench_vec_znx_sub_ba_inplace<B>(c: &mut Criterion, label: &str)
209where
210    B: Backend + ModuleNewImpl<B> + VecZnxSubBAInplaceImpl<B>,
211{
212    let group_name: String = format!("vec_znx_sub_ba_inplace::{}", label);
213
214    let mut group = c.benchmark_group(group_name);
215
216    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
217    where
218        Module<B>: VecZnxSubBAInplace + ModuleNew<B>,
219    {
220        let n: usize = 1 << params[0];
221        let cols: usize = params[1];
222        let size: usize = params[2];
223
224        let module: Module<B> = Module::<B>::new(n as u64);
225
226        let mut source: Source = Source::new([0u8; 32]);
227
228        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
229        let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
230
231        // Fill a with random i64
232        a.fill_uniform(50, &mut source);
233        b.fill_uniform(50, &mut source);
234
235        move || {
236            for i in 0..cols {
237                module.vec_znx_sub_ba_inplace(&mut b, i, &a, i);
238            }
239            black_box(());
240        }
241    }
242
243    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
244        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2]));
245        let mut runner = runner::<B>(params);
246        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
247    }
248
249    group.finish();
250}