1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4use rand::RngCore;
  5
  6use crate::{
  7    api::{
  8        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxBigAlloc, VecZnxDftAdd, VecZnxDftAddInplace, VecZnxDftAlloc,
  9        VecZnxDftApply, VecZnxDftSub, VecZnxDftSubInplace, VecZnxDftSubNegateInplace, VecZnxIdftApply, VecZnxIdftApplyTmpA,
 10        VecZnxIdftApplyTmpBytes,
 11    },
 12    layouts::{Backend, DataViewMut, Module, ScratchOwned, VecZnx, VecZnxBig, VecZnxDft},
 13    source::Source,
 14};
 15
 16pub fn bench_vec_znx_dft_add<B: Backend>(c: &mut Criterion, label: &str)
 17where
 18    Module<B>: VecZnxDftAdd<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
 19{
 20    let group_name: String = format!("vec_znx_dft_add::{label}");
 21
 22    let mut group = c.benchmark_group(group_name);
 23
 24    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
 25    where
 26        Module<B>: VecZnxDftAdd<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
 27    {
 28        let n: usize = params[0];
 29        let cols: usize = params[1];
 30        let size: usize = params[2];
 31
 32        let module: Module<B> = Module::<B>::new(n as u64);
 33
 34        let mut source: Source = Source::new([0u8; 32]);
 35
 36        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
 37        let mut b: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
 38        let mut c: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
 39
 40        source.fill_bytes(a.data_mut());
 41        source.fill_bytes(b.data_mut());
 42        source.fill_bytes(c.data_mut());
 43
 44        move || {
 45            for i in 0..cols {
 46                module.vec_znx_dft_add(&mut c, i, &a, i, &b, i);
 47            }
 48            black_box(());
 49        }
 50    }
 51
 52    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
 53        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
 54        let mut runner = runner::<B>(params);
 55        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 56    }
 57
 58    group.finish();
 59}
 60
 61pub fn bench_vec_znx_dft_add_inplace<B: Backend>(c: &mut Criterion, label: &str)
 62where
 63    Module<B>: VecZnxDftAddInplace<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
 64{
 65    let group_name: String = format!("vec_znx_dft_add_inplace::{label}");
 66
 67    let mut group = c.benchmark_group(group_name);
 68
 69    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
 70    where
 71        Module<B>: VecZnxDftAddInplace<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
 72    {
 73        let n: usize = params[0];
 74        let cols: usize = params[1];
 75        let size: usize = params[2];
 76
 77        let module: Module<B> = Module::<B>::new(n as u64);
 78
 79        let mut source: Source = Source::new([0u8; 32]);
 80
 81        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
 82        let mut c: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
 83
 84        // Fill a with random i64
 85        source.fill_bytes(a.data_mut());
 86        source.fill_bytes(c.data_mut());
 87
 88        move || {
 89            for i in 0..cols {
 90                module.vec_znx_dft_add_inplace(&mut c, i, &a, i);
 91            }
 92            black_box(());
 93        }
 94    }
 95
 96    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
 97        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
 98        let mut runner = runner::<B>(params);
 99        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
100    }
101
102    group.finish();
103}
104
105pub fn bench_vec_znx_dft_apply<B: Backend>(c: &mut Criterion, label: &str)
106where
107    Module<B>: VecZnxDftApply<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
108{
109    let group_name: String = format!("vec_znx_dft_apply::{label}");
110
111    let mut group = c.benchmark_group(group_name);
112
113    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
114    where
115        Module<B>: VecZnxDftApply<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
116    {
117        let n: usize = params[0];
118        let cols: usize = params[1];
119        let size: usize = params[2];
120
121        let module: Module<B> = Module::<B>::new(n as u64);
122
123        let mut source: Source = Source::new([0u8; 32]);
124
125        let mut res: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
126        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
127        source.fill_bytes(res.data_mut());
128        source.fill_bytes(a.data_mut());
129
130        move || {
131            for i in 0..cols {
132                module.vec_znx_dft_apply(1, 0, &mut res, i, &a, i);
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_idft_apply<B: Backend>(c: &mut Criterion, label: &str)
148where
149    Module<B>: VecZnxIdftApply<B> + ModuleNew<B> + VecZnxIdftApplyTmpBytes + VecZnxDftAlloc<B> + VecZnxBigAlloc<B>,
150    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
151{
152    let group_name: String = format!("vec_znx_idft_apply::{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>: VecZnxIdftApply<B> + ModuleNew<B> + VecZnxIdftApplyTmpBytes + VecZnxDftAlloc<B> + VecZnxBigAlloc<B>,
159        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
160    {
161        let n: usize = 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 mut source: Source = Source::new([0u8; 32]);
168
169        let mut res: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(cols, size);
170        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
171        source.fill_bytes(res.data_mut());
172        source.fill_bytes(a.data_mut());
173
174        let mut scratch = ScratchOwned::alloc(module.vec_znx_idft_apply_tmp_bytes());
175
176        move || {
177            for i in 0..cols {
178                module.vec_znx_idft_apply(&mut res, i, &a, i, scratch.borrow());
179            }
180            black_box(());
181        }
182    }
183
184    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
185        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
186        let mut runner = runner::<B>(params);
187        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
188    }
189
190    group.finish();
191}
192
193pub fn bench_vec_znx_idft_apply_tmpa<B: Backend>(c: &mut Criterion, label: &str)
194where
195    Module<B>: VecZnxIdftApplyTmpA<B> + ModuleNew<B> + VecZnxDftAlloc<B> + VecZnxBigAlloc<B>,
196{
197    let group_name: String = format!("vec_znx_idft_apply_tmpa::{label}");
198
199    let mut group = c.benchmark_group(group_name);
200
201    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
202    where
203        Module<B>: VecZnxIdftApplyTmpA<B> + ModuleNew<B> + VecZnxDftAlloc<B> + VecZnxBigAlloc<B>,
204    {
205        let module: Module<B> = Module::<B>::new(1 << params[0]);
206
207        let cols: usize = params[1];
208        let size: usize = params[2];
209
210        let mut source: Source = Source::new([0u8; 32]);
211
212        let mut res: VecZnxBig<Vec<u8>, B> = module.vec_znx_big_alloc(cols, size);
213        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
214        source.fill_bytes(res.data_mut());
215        source.fill_bytes(a.data_mut());
216
217        move || {
218            for i in 0..cols {
219                module.vec_znx_idft_apply_tmpa(&mut res, i, &mut a, i);
220            }
221            black_box(());
222        }
223    }
224
225    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
226        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
227        let mut runner = runner::<B>(params);
228        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
229    }
230
231    group.finish();
232}
233
234pub fn bench_vec_znx_dft_sub<B: Backend>(c: &mut Criterion, label: &str)
235where
236    Module<B>: VecZnxDftSub<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
237{
238    let group_name: String = format!("vec_znx_dft_sub::{label}");
239
240    let mut group = c.benchmark_group(group_name);
241
242    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
243    where
244        Module<B>: VecZnxDftSub<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
245    {
246        let n: usize = params[0];
247        let cols: usize = params[1];
248        let size: usize = params[2];
249
250        let module: Module<B> = Module::<B>::new(n as u64);
251
252        let mut source: Source = Source::new([0u8; 32]);
253
254        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
255        let mut b: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
256        let mut c: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
257
258        source.fill_bytes(a.data_mut());
259        source.fill_bytes(b.data_mut());
260        source.fill_bytes(c.data_mut());
261
262        move || {
263            for i in 0..cols {
264                module.vec_znx_dft_sub(&mut c, i, &a, i, &b, i);
265            }
266            black_box(());
267        }
268    }
269
270    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
271        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
272        let mut runner = runner::<B>(params);
273        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
274    }
275
276    group.finish();
277}
278
279pub fn bench_vec_znx_dft_sub_inplace<B: Backend>(c: &mut Criterion, label: &str)
280where
281    Module<B>: VecZnxDftSubInplace<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
282{
283    let group_name: String = format!("vec_znx_dft_sub_inplace::{label}");
284
285    let mut group = c.benchmark_group(group_name);
286
287    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
288    where
289        Module<B>: VecZnxDftSubInplace<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
290    {
291        let n: usize = params[0];
292        let cols: usize = params[1];
293        let size: usize = params[2];
294
295        let module: Module<B> = Module::<B>::new(n as u64);
296
297        let mut source: Source = Source::new([0u8; 32]);
298
299        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
300        let mut c: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
301
302        // Fill a with random i64
303        source.fill_bytes(a.data_mut());
304        source.fill_bytes(c.data_mut());
305
306        move || {
307            for i in 0..cols {
308                module.vec_znx_dft_sub_inplace(&mut c, i, &a, i);
309            }
310            black_box(());
311        }
312    }
313
314    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
315        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
316        let mut runner = runner::<B>(params);
317        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
318    }
319
320    group.finish();
321}
322
323pub fn bench_vec_znx_dft_sub_negate_inplace<B: Backend>(c: &mut Criterion, label: &str)
324where
325    Module<B>: VecZnxDftSubNegateInplace<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
326{
327    let group_name: String = format!("vec_znx_dft_sub_negate_inplace::{label}");
328
329    let mut group = c.benchmark_group(group_name);
330
331    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
332    where
333        Module<B>: VecZnxDftSubNegateInplace<B> + ModuleNew<B> + VecZnxDftAlloc<B>,
334    {
335        let n: usize = params[0];
336        let cols: usize = params[1];
337        let size: usize = params[2];
338
339        let module: Module<B> = Module::<B>::new(n as u64);
340
341        let mut source: Source = Source::new([0u8; 32]);
342
343        let mut a: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
344        let mut c: VecZnxDft<Vec<u8>, B> = module.vec_znx_dft_alloc(cols, size);
345
346        // Fill a with random i64
347        source.fill_bytes(a.data_mut());
348        source.fill_bytes(c.data_mut());
349
350        move || {
351            for i in 0..cols {
352                module.vec_znx_dft_sub_negate_inplace(&mut c, i, &a, i);
353            }
354            black_box(());
355        }
356    }
357
358    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
359        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
360        let mut runner = runner::<B>(params);
361        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
362    }
363
364    group.finish();
365}