1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4
  5use crate::{
  6    api::{CnvPVecAlloc, Convolution, ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxBigAlloc, VecZnxDftAlloc},
  7    layouts::{Backend, CnvPVecL, CnvPVecR, FillUniform, Module, ScratchOwned, VecZnx, VecZnxBig},
  8    source::Source,
  9};
 10
 11pub fn bench_cnv_prepare_left<BE: Backend>(c: &mut Criterion, label: &str)
 12where
 13    Module<BE>: ModuleNew<BE> + Convolution<BE> + CnvPVecAlloc<BE>,
 14    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 15{
 16    let group_name: String = format!("cnv_prepare_left::{label}");
 17
 18    let mut group = c.benchmark_group(group_name);
 19
 20    fn runner<BE: Backend>(n: usize, size: usize) -> impl FnMut()
 21    where
 22        Module<BE>: ModuleNew<BE> + Convolution<BE> + CnvPVecAlloc<BE>,
 23        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 24    {
 25        let mut source: Source = Source::new([0u8; 32]);
 26
 27        let base2k: usize = 12;
 28
 29        let c_size: usize = size + size - 1;
 30
 31        let module: Module<BE> = Module::<BE>::new(n as u64);
 32
 33        let mut a_prep: CnvPVecL<Vec<u8>, BE> = module.cnv_pvec_left_alloc(1, size);
 34
 35        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, size);
 36
 37        a.fill_uniform(base2k, &mut source);
 38
 39        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(module.cnv_prepare_left_tmp_bytes(c_size, size));
 40
 41        move || {
 42            module.cnv_prepare_left(&mut a_prep, &a, scratch.borrow());
 43            black_box(());
 44        }
 45    }
 46
 47    for params in [[10, 1], [11, 2], [12, 4], [13, 8], [14, 16], [15, 32], [16, 64]] {
 48        let log_n: usize = params[0];
 49        let size: usize = params[1];
 50        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x{}", 1 << log_n, size));
 51        let mut runner = runner(1 << log_n, size);
 52        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 53    }
 54
 55    group.finish();
 56}
 57
 58pub fn bench_cnv_prepare_right<BE: Backend>(c: &mut Criterion, label: &str)
 59where
 60    Module<BE>: ModuleNew<BE> + Convolution<BE> + CnvPVecAlloc<BE>,
 61    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 62{
 63    let group_name: String = format!("cnv_prepare_right::{label}");
 64
 65    let mut group = c.benchmark_group(group_name);
 66
 67    fn runner<BE: Backend>(n: usize, size: usize) -> impl FnMut()
 68    where
 69        Module<BE>: ModuleNew<BE> + Convolution<BE> + CnvPVecAlloc<BE>,
 70        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 71    {
 72        let mut source: Source = Source::new([0u8; 32]);
 73
 74        let base2k: usize = 12;
 75
 76        let c_size: usize = size + size - 1;
 77
 78        let module: Module<BE> = Module::<BE>::new(n as u64);
 79
 80        let mut a_prep: CnvPVecR<Vec<u8>, BE> = module.cnv_pvec_right_alloc(1, size);
 81
 82        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, size);
 83
 84        a.fill_uniform(base2k, &mut source);
 85
 86        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(module.cnv_prepare_right_tmp_bytes(c_size, size));
 87
 88        move || {
 89            module.cnv_prepare_right(&mut a_prep, &a, scratch.borrow());
 90            black_box(());
 91        }
 92    }
 93
 94    for params in [[10, 1], [11, 2], [12, 4], [13, 8], [14, 16], [15, 32], [16, 64]] {
 95        let log_n: usize = params[0];
 96        let size: usize = params[1];
 97        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x{}", 1 << log_n, size));
 98        let mut runner = runner(1 << log_n, size);
 99        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
100    }
101
102    group.finish();
103}
104
105pub fn bench_cnv_apply_dft<BE: Backend>(c: &mut Criterion, label: &str)
106where
107    Module<BE>: ModuleNew<BE> + Convolution<BE> + VecZnxDftAlloc<BE> + CnvPVecAlloc<BE>,
108    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
109{
110    let group_name: String = format!("cnv_apply_dft::{label}");
111
112    let mut group = c.benchmark_group(group_name);
113
114    fn runner<BE: Backend>(n: usize, size: usize) -> impl FnMut()
115    where
116        Module<BE>: ModuleNew<BE> + Convolution<BE> + VecZnxDftAlloc<BE> + CnvPVecAlloc<BE>,
117        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
118    {
119        let mut source: Source = Source::new([0u8; 32]);
120
121        let base2k: usize = 12;
122
123        let c_size: usize = size + size - 1;
124
125        let module: Module<BE> = Module::<BE>::new(n as u64);
126
127        let mut a_prep: CnvPVecL<Vec<u8>, BE> = module.cnv_pvec_left_alloc(1, size);
128        let mut b_prep: CnvPVecR<Vec<u8>, BE> = module.cnv_pvec_right_alloc(1, size);
129
130        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, size);
131        let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), 1, size);
132        let mut c_dft = module.vec_znx_dft_alloc(1, c_size);
133
134        a.fill_uniform(base2k, &mut source);
135        b.fill_uniform(base2k, &mut source);
136
137        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
138            module
139                .cnv_apply_dft_tmp_bytes(c_size, 0, size, size)
140                .max(module.cnv_prepare_left_tmp_bytes(c_size, size))
141                .max(module.cnv_prepare_right_tmp_bytes(c_size, size)),
142        );
143        module.cnv_prepare_left(&mut a_prep, &a, scratch.borrow());
144        module.cnv_prepare_right(&mut b_prep, &b, scratch.borrow());
145        move || {
146            module.cnv_apply_dft(&mut c_dft, 0, 0, &a_prep, 0, &b_prep, 0, scratch.borrow());
147            black_box(());
148        }
149    }
150
151    for params in [[10, 1], [11, 2], [12, 4], [13, 8], [14, 16], [15, 32], [16, 64]] {
152        let log_n: usize = params[0];
153        let size: usize = params[1];
154        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x{}", 1 << log_n, size));
155        let mut runner = runner(1 << log_n, size);
156        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
157    }
158
159    group.finish();
160}
161
162pub fn bench_cnv_pairwise_apply_dft<BE: Backend>(c: &mut Criterion, label: &str)
163where
164    Module<BE>: ModuleNew<BE> + Convolution<BE> + VecZnxDftAlloc<BE> + CnvPVecAlloc<BE>,
165    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
166{
167    let group_name: String = format!("cnv_pairwise_apply_dft::{label}");
168
169    let mut group = c.benchmark_group(group_name);
170
171    fn runner<BE: Backend>(n: usize, size: usize) -> impl FnMut()
172    where
173        Module<BE>: ModuleNew<BE> + Convolution<BE> + VecZnxDftAlloc<BE> + CnvPVecAlloc<BE>,
174        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
175    {
176        let mut source: Source = Source::new([0u8; 32]);
177
178        let base2k: usize = 12;
179
180        let module: Module<BE> = Module::<BE>::new(n as u64);
181
182        let cols = 2;
183        let c_size: usize = size + size - 1;
184
185        let mut a_prep: CnvPVecL<Vec<u8>, BE> = module.cnv_pvec_left_alloc(cols, size);
186        let mut b_prep: CnvPVecR<Vec<u8>, BE> = module.cnv_pvec_right_alloc(cols, size);
187
188        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), cols, size);
189        let mut b: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), cols, size);
190        let mut c_dft = module.vec_znx_dft_alloc(1, c_size);
191
192        a.fill_uniform(base2k, &mut source);
193        b.fill_uniform(base2k, &mut source);
194
195        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
196            module
197                .cnv_pairwise_apply_dft_tmp_bytes(c_size, 0, size, size)
198                .max(module.cnv_prepare_left_tmp_bytes(c_size, size))
199                .max(module.cnv_prepare_right_tmp_bytes(c_size, size)),
200        );
201        module.cnv_prepare_left(&mut a_prep, &a, scratch.borrow());
202        module.cnv_prepare_right(&mut b_prep, &b, scratch.borrow());
203        move || {
204            module.cnv_pairwise_apply_dft(&mut c_dft, 0, 0, &a_prep, &b_prep, 0, 1, scratch.borrow());
205            black_box(());
206        }
207    }
208
209    for params in [[10, 1], [11, 2], [12, 4], [13, 8], [14, 16], [15, 32], [16, 64]] {
210        let log_n: usize = params[0];
211        let size: usize = params[1];
212        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x{}", 1 << log_n, size));
213        let mut runner = runner(1 << log_n, size);
214        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
215    }
216
217    group.finish();
218}
219
220pub fn bench_cnv_by_const_apply<BE: Backend>(c: &mut Criterion, label: &str)
221where
222    Module<BE>: ModuleNew<BE> + Convolution<BE> + VecZnxBigAlloc<BE> + CnvPVecAlloc<BE>,
223    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
224{
225    let group_name: String = format!("cnv_by_const::{label}");
226
227    let mut group = c.benchmark_group(group_name);
228
229    fn runner<BE: Backend>(n: usize, size: usize) -> impl FnMut()
230    where
231        Module<BE>: ModuleNew<BE> + Convolution<BE> + VecZnxBigAlloc<BE> + CnvPVecAlloc<BE>,
232        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
233    {
234        let mut source: Source = Source::new([0u8; 32]);
235
236        let base2k: usize = 12;
237
238        let module: Module<BE> = Module::<BE>::new(n as u64);
239
240        let cols = 2;
241        let c_size: usize = size + size - 1;
242
243        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), cols, size);
244        let mut c_big: VecZnxBig<Vec<u8>, BE> = module.vec_znx_big_alloc(1, c_size);
245
246        a.fill_uniform(base2k, &mut source);
247        let mut b = vec![0i64; size];
248        for x in &mut b {
249            *x = source.next_i64();
250        }
251
252        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(module.cnv_by_const_apply_tmp_bytes(c_size, 0, size, size));
253        move || {
254            module.cnv_by_const_apply(&mut c_big, 0, 0, &a, 0, &b, scratch.borrow());
255            black_box(());
256        }
257    }
258
259    for params in [[10, 1], [11, 2], [12, 4], [13, 8], [14, 16], [15, 32], [16, 64]] {
260        let log_n: usize = params[0];
261        let size: usize = params[1];
262        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x{}", 1 << log_n, size));
263        let mut runner = runner(1 << log_n, size);
264        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
265    }
266
267    group.finish();
268}