poulpy_hal/reference/vec_znx/
automorphism.rs

1use std::hint::black_box;
2
3use criterion::{BenchmarkId, Criterion};
4
5use crate::{
6    api::{
7        ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAutomorphism, VecZnxAutomorphismInplace,
8        VecZnxAutomorphismInplaceTmpBytes,
9    },
10    layouts::{Backend, FillUniform, Module, ScratchOwned, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxView, ZnxViewMut},
11    reference::znx::{ZnxAutomorphism, ZnxCopy, ZnxZero},
12    source::Source,
13};
14
15pub fn vec_znx_automorphism_inplace_tmp_bytes(n: usize) -> usize {
16    n * size_of::<i64>()
17}
18
19pub fn vec_znx_automorphism<R, A, ZNXARI>(p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
20where
21    R: VecZnxToMut,
22    A: VecZnxToRef,
23    ZNXARI: ZnxAutomorphism + ZnxZero,
24{
25    let a: VecZnx<&[u8]> = a.to_ref();
26    let mut res: VecZnx<&mut [u8]> = res.to_mut();
27
28    #[cfg(debug_assertions)]
29    {
30        use crate::layouts::ZnxInfos;
31
32        assert_eq!(a.n(), res.n());
33    }
34
35    let min_size: usize = res.size().min(a.size());
36
37    for j in 0..min_size {
38        ZNXARI::znx_automorphism(p, res.at_mut(res_col, j), a.at(a_col, j));
39    }
40
41    for j in min_size..res.size() {
42        ZNXARI::znx_zero(res.at_mut(res_col, j));
43    }
44}
45
46pub fn vec_znx_automorphism_inplace<R, ZNXARI>(p: i64, res: &mut R, res_col: usize, tmp: &mut [i64])
47where
48    R: VecZnxToMut,
49    ZNXARI: ZnxAutomorphism + ZnxCopy,
50{
51    let mut res: VecZnx<&mut [u8]> = res.to_mut();
52    #[cfg(debug_assertions)]
53    {
54        assert_eq!(res.n(), tmp.len());
55    }
56    for j in 0..res.size() {
57        ZNXARI::znx_automorphism(p, tmp, res.at(res_col, j));
58        ZNXARI::znx_copy(res.at_mut(res_col, j), tmp);
59    }
60}
61
62pub fn bench_vec_znx_automorphism<B: Backend>(c: &mut Criterion, label: &str)
63where
64    Module<B>: VecZnxAutomorphism + ModuleNew<B>,
65{
66    let group_name: String = format!("vec_znx_automorphism::{label}");
67
68    let mut group = c.benchmark_group(group_name);
69
70    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
71    where
72        Module<B>: VecZnxAutomorphism + ModuleNew<B>,
73    {
74        let n: usize = 1 << params[0];
75        let cols: usize = params[1];
76        let size: usize = params[2];
77
78        let module: Module<B> = Module::<B>::new(n as u64);
79
80        let mut source: Source = Source::new([0u8; 32]);
81
82        let mut a: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
83        let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
84
85        // Fill a with random i64
86        a.fill_uniform(50, &mut source);
87        res.fill_uniform(50, &mut source);
88
89        move || {
90            for i in 0..cols {
91                module.vec_znx_automorphism(-7, &mut res, i, &a, i);
92            }
93            black_box(());
94        }
95    }
96
97    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
98        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
99        let mut runner = runner::<B>(params);
100        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
101    }
102
103    group.finish();
104}
105
106pub fn bench_vec_znx_automorphism_inplace<B: Backend>(c: &mut Criterion, label: &str)
107where
108    Module<B>: VecZnxAutomorphismInplace<B> + VecZnxAutomorphismInplaceTmpBytes + ModuleNew<B>,
109    ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
110{
111    let group_name: String = format!("vec_znx_automorphism_inplace::{label}");
112
113    let mut group = c.benchmark_group(group_name);
114
115    fn runner<B: Backend>(params: [usize; 3]) -> impl FnMut()
116    where
117        Module<B>: VecZnxAutomorphismInplace<B> + ModuleNew<B> + VecZnxAutomorphismInplaceTmpBytes,
118        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
119    {
120        let n: usize = 1 << params[0];
121        let cols: usize = params[1];
122        let size: usize = params[2];
123
124        let module: Module<B> = Module::<B>::new(n as u64);
125
126        let mut source: Source = Source::new([0u8; 32]);
127
128        let mut res: VecZnx<Vec<u8>> = VecZnx::alloc(n, cols, size);
129
130        let mut scratch = ScratchOwned::alloc(module.vec_znx_automorphism_inplace_tmp_bytes());
131
132        // Fill a with random i64
133        res.fill_uniform(50, &mut source);
134
135        move || {
136            for i in 0..cols {
137                module.vec_znx_automorphism_inplace(-7, &mut res, i, scratch.borrow());
138            }
139            black_box(());
140        }
141    }
142
143    for params in [[10, 2, 2], [11, 2, 4], [12, 2, 8], [13, 2, 16], [14, 2, 32]] {
144        let id: BenchmarkId = BenchmarkId::from_parameter(format!("{}x({}x{})", 1 << params[0], params[1], params[2],));
145        let mut runner = runner::<B>(params);
146        group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
147    }
148
149    group.finish();
150}