1use std::hint::black_box;
  2
  3use criterion::{BenchmarkId, Criterion};
  4use rand::RngCore;
  5
  6use crate::{
  7    api::{
  8        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxDftAlloc, VmpApplyDft, VmpApplyDftTmpBytes, VmpApplyDftToDft,
  9        VmpApplyDftToDftAdd, VmpApplyDftToDftAddTmpBytes, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare, VmpPrepareTmpBytes,
 10    },
 11    layouts::{Backend, DataViewMut, MatZnx, Module, ScratchOwned, VecZnx, VecZnxDft, VmpPMat},
 12    source::Source,
 13};
 14
 15pub fn bench_vmp_prepare<B: Backend>(c: &mut Criterion, label: &str)
 16where
 17    Module<B>: ModuleNew<B> + VmpPMatAlloc<B> + VmpPrepare<B> + VmpPrepareTmpBytes,
 18    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 19{
 20    let group_name: String = format!("vmp_prepare::{label}");
 21
 22    let mut group = c.benchmark_group(group_name);
 23
 24    fn runner<B: Backend>(params: [usize; 5]) -> impl FnMut()
 25    where
 26        Module<B>: ModuleNew<B> + VmpPMatAlloc<B> + VmpPrepare<B> + VmpPrepareTmpBytes,
 27        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 28    {
 29        let module: Module<B> = Module::<B>::new(1 << params[0]);
 30
 31        let rows: usize = params[1];
 32        let cols_in: usize = params[2];
 33        let cols_out: usize = params[3];
 34        let size: usize = params[4];
 35
 36        let mut source: Source = Source::new([0u8; 32]);
 37
 38        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vmp_prepare_tmp_bytes(rows, cols_in, cols_out, size));
 39
 40        let mut mat: MatZnx<Vec<u8>> = MatZnx::alloc(module.n(), rows, cols_in, cols_out, size);
 41        let mut pmat: VmpPMat<Vec<u8>, B> = module.vmp_pmat_alloc(rows, cols_in, cols_out, size);
 42
 43        source.fill_bytes(mat.data_mut());
 44        source.fill_bytes(pmat.data_mut());
 45
 46        move || {
 47            module.vmp_prepare(&mut pmat, &mat, scratch.borrow());
 48            black_box(());
 49        }
 50    }
 51
 52    for params in [
 53        [10, 2, 1, 2, 3],
 54        [11, 4, 1, 2, 5],
 55        [12, 7, 1, 2, 8],
 56        [13, 15, 1, 2, 16],
 57        [14, 31, 1, 2, 32],
 58    ] {
 59        let id = BenchmarkId::from_parameter(format!(
 60            "{}x({}x{})x({}x{})",
 61            1 << params[0],
 62            params[2],
 63            params[1],
 64            params[3],
 65            params[4]
 66        ));
 67        let mut runner = runner::<B>(params);
 68        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
 69    }
 70
 71    group.finish();
 72}
 73
 74pub fn bench_vmp_apply_dft<B: Backend>(c: &mut Criterion, label: &str)
 75where
 76    Module<B>: ModuleNew<B> + VmpApplyDftTmpBytes + VmpApplyDft<B> + VmpPMatAlloc<B> + VecZnxDftAlloc<B>,
 77    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 78{
 79    let group_name: String = format!("vmp_apply_dft::{label}");
 80
 81    let mut group = c.benchmark_group(group_name);
 82
 83    fn runner<B: Backend>(params: [usize; 5]) -> impl FnMut()
 84    where
 85        Module<B>: ModuleNew<B> + VmpApplyDftTmpBytes + VmpApplyDft<B> + VmpPMatAlloc<B> + VecZnxDftAlloc<B>,
 86        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 87    {
 88        let module: Module<B> = Module::<B>::new(1 << params[0]);
 89
 90        let rows: usize = params[1];
 91        let cols_in: usize = params[2];
 92        let cols_out: usize = params[3];
 93        let size: usize = params[4];
 94
 95        let mut source: Source = Source::new([0u8; 32]);
 96
 97        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(1 << 20);
 98
 99        let mut res: VecZnxDft<Vec<u8>, _> = module.vec_znx_dft_alloc(cols_out, size);
100        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(module.n(), cols_in, size);
101        let mut pmat: VmpPMat<Vec<u8>, B> = module.vmp_pmat_alloc(rows, cols_in, cols_out, size);
102
103        source.fill_bytes(pmat.data_mut());
104        source.fill_bytes(res.data_mut());
105        source.fill_bytes(a.data_mut());
106
107        move || {
108            module.vmp_apply_dft(&mut res, &a, &pmat, scratch.borrow());
109            black_box(());
110        }
111    }
112
113    for params in [
114        [10, 2, 1, 2, 3],
115        [11, 4, 1, 2, 5],
116        [12, 7, 1, 2, 8],
117        [13, 15, 1, 2, 16],
118        [14, 31, 1, 2, 32],
119    ] {
120        let id = BenchmarkId::from_parameter(format!(
121            "{}x({}x{})x({}x{})",
122            1 << params[0],
123            params[2],
124            params[1],
125            params[3],
126            params[4]
127        ));
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_vmp_apply_dft_to_dft<B: Backend>(c: &mut Criterion, label: &str)
136where
137    Module<B>: ModuleNew<B> + VecZnxDftAlloc<B> + VmpPMatAlloc<B> + VmpApplyDftToDft<B> + VmpApplyDftToDftTmpBytes,
138    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
139{
140    let group_name: String = format!("vmp_apply_dft_to_dft::{label}");
141
142    let mut group = c.benchmark_group(group_name);
143
144    fn runner<B: Backend>(params: [usize; 5]) -> impl FnMut()
145    where
146        Module<B>: ModuleNew<B> + VecZnxDftAlloc<B> + VmpPMatAlloc<B> + VmpApplyDftToDft<B> + VmpApplyDftToDftTmpBytes,
147        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
148    {
149        let module: Module<B> = Module::<B>::new(1 << params[0]);
150
151        let rows: usize = params[1];
152        let cols_in: usize = params[2];
153        let cols_out: usize = params[3];
154        let size: usize = params[4];
155
156        let mut source: Source = Source::new([0u8; 32]);
157
158        let mut scratch: ScratchOwned<B> =
159            ScratchOwned::alloc(module.vmp_apply_dft_to_dft_tmp_bytes(size, size, rows, cols_in, cols_out, size));
160
161        let mut res: VecZnxDft<Vec<u8>, _> = module.vec_znx_dft_alloc(cols_out, size);
162        let mut a: VecZnxDft<Vec<u8>, _> = module.vec_znx_dft_alloc(cols_in, size);
163
164        let mut pmat: VmpPMat<Vec<u8>, B> = module.vmp_pmat_alloc(rows, cols_in, cols_out, size);
165
166        source.fill_bytes(pmat.data_mut());
167        source.fill_bytes(res.data_mut());
168        source.fill_bytes(a.data_mut());
169
170        move || {
171            module.vmp_apply_dft_to_dft(&mut res, &a, &pmat, scratch.borrow());
172            black_box(());
173        }
174    }
175
176    for params in [
177        [10, 2, 1, 2, 3],
178        [11, 4, 1, 2, 5],
179        [12, 7, 1, 2, 8],
180        [13, 15, 1, 2, 16],
181        [14, 31, 1, 2, 32],
182    ] {
183        let id = BenchmarkId::from_parameter(format!(
184            "{}x({}x{})x({}x{})",
185            1 << params[0], // n
186            params[2],      // cols_in
187            params[1],      // size_in (=rows)
188            params[3],      // cols_out
189            params[4]       // size_out
190        ));
191        let mut runner = runner::<B>(params);
192        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
193    }
194
195    group.finish();
196}
197
198pub fn bench_vmp_apply_dft_to_dft_add<B: Backend>(c: &mut Criterion, label: &str)
199where
200    Module<B>: ModuleNew<B> + VecZnxDftAlloc<B> + VmpPMatAlloc<B> + VmpApplyDftToDftAdd<B> + VmpApplyDftToDftAddTmpBytes,
201    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
202{
203    let group_name: String = format!("vmp_apply_dft_to_dft_add::{label}");
204
205    let mut group = c.benchmark_group(group_name);
206
207    fn runner<B: Backend>(params: [usize; 5]) -> impl FnMut()
208    where
209        Module<B>: ModuleNew<B> + VecZnxDftAlloc<B> + VmpPMatAlloc<B> + VmpApplyDftToDftAdd<B> + VmpApplyDftToDftAddTmpBytes,
210        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
211    {
212        let module: Module<B> = Module::<B>::new(1 << params[0]);
213
214        let rows: usize = params[1];
215        let cols_in: usize = params[2];
216        let cols_out: usize = params[3];
217        let size: usize = params[4];
218
219        let mut source: Source = Source::new([0u8; 32]);
220
221        let mut scratch: ScratchOwned<B> =
222            ScratchOwned::alloc(module.vmp_apply_dft_to_dft_add_tmp_bytes(size, size, rows, cols_in, cols_out, size));
223
224        let mut res: VecZnxDft<Vec<u8>, _> = module.vec_znx_dft_alloc(cols_out, size);
225        let mut a: VecZnxDft<Vec<u8>, _> = module.vec_znx_dft_alloc(cols_in, size);
226
227        let mut pmat: VmpPMat<Vec<u8>, B> = module.vmp_pmat_alloc(rows, cols_in, cols_out, size);
228
229        source.fill_bytes(pmat.data_mut());
230        source.fill_bytes(res.data_mut());
231        source.fill_bytes(a.data_mut());
232
233        move || {
234            module.vmp_apply_dft_to_dft_add(&mut res, &a, &pmat, 1, scratch.borrow());
235            black_box(());
236        }
237    }
238
239    for params in [
240        [10, 2, 1, 2, 3],
241        [11, 4, 1, 2, 5],
242        [12, 7, 1, 2, 8],
243        [13, 15, 1, 2, 16],
244        [14, 31, 1, 2, 32],
245    ] {
246        let id = BenchmarkId::from_parameter(format!(
247            "{}x({}x{})x({}x{})",
248            1 << params[0],
249            params[2],
250            params[1],
251            params[3],
252            params[4]
253        ));
254        let mut runner = runner::<B>(params);
255        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
256    }
257
258    group.finish();
259}