poulpy_core/automorphism/
glwe_ct.rs

1use poulpy_hal::{
2    api::{
3        ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace,
4        VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallAInplace, VecZnxBigSubSmallBInplace,
5        VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
6    },
7    layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnxBig},
8};
9
10use crate::layouts::{GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyPrepared};
11
12impl GLWECiphertext<Vec<u8>> {
13    #[allow(clippy::too_many_arguments)]
14    pub fn automorphism_scratch_space<B: Backend>(
15        module: &Module<B>,
16        n: usize,
17        basek: usize,
18        k_out: usize,
19        k_in: usize,
20        k_ksk: usize,
21        digits: usize,
22        rank: usize,
23    ) -> usize
24    where
25        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
26    {
27        Self::keyswitch_scratch_space(module, n, basek, k_out, k_in, k_ksk, digits, rank, rank)
28    }
29
30    pub fn automorphism_inplace_scratch_space<B: Backend>(
31        module: &Module<B>,
32        n: usize,
33        basek: usize,
34        k_out: usize,
35        k_ksk: usize,
36        digits: usize,
37        rank: usize,
38    ) -> usize
39    where
40        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxBigNormalizeTmpBytes,
41    {
42        Self::keyswitch_inplace_scratch_space(module, n, basek, k_out, k_ksk, digits, rank)
43    }
44}
45
46impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
47    pub fn automorphism<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
48        &mut self,
49        module: &Module<B>,
50        lhs: &GLWECiphertext<DataLhs>,
51        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
52        scratch: &mut Scratch<B>,
53    ) where
54        Module<B>: VecZnxDftAllocBytes
55            + VmpApplyTmpBytes
56            + VecZnxBigNormalizeTmpBytes
57            + VmpApply<B>
58            + VmpApplyAdd<B>
59            + VecZnxDftFromVecZnx<B>
60            + VecZnxDftToVecZnxBigConsume<B>
61            + VecZnxBigAddSmallInplace<B>
62            + VecZnxBigNormalize<B>
63            + VecZnxAutomorphismInplace,
64        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
65    {
66        self.keyswitch(module, lhs, &rhs.key, scratch);
67        (0..self.rank() + 1).for_each(|i| {
68            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
69        })
70    }
71
72    pub fn automorphism_inplace<DataRhs: DataRef, B: Backend>(
73        &mut self,
74        module: &Module<B>,
75        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
76        scratch: &mut Scratch<B>,
77    ) where
78        Module<B>: VecZnxDftAllocBytes
79            + VmpApplyTmpBytes
80            + VecZnxBigNormalizeTmpBytes
81            + VmpApply<B>
82            + VmpApplyAdd<B>
83            + VecZnxDftFromVecZnx<B>
84            + VecZnxDftToVecZnxBigConsume<B>
85            + VecZnxBigAddSmallInplace<B>
86            + VecZnxBigNormalize<B>
87            + VecZnxAutomorphismInplace,
88        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
89    {
90        self.keyswitch_inplace(module, &rhs.key, scratch);
91        (0..self.rank() + 1).for_each(|i| {
92            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
93        })
94    }
95
96    pub fn automorphism_add<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
97        &mut self,
98        module: &Module<B>,
99        lhs: &GLWECiphertext<DataLhs>,
100        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
101        scratch: &mut Scratch<B>,
102    ) where
103        Module<B>: VecZnxDftAllocBytes
104            + VmpApplyTmpBytes
105            + VecZnxBigNormalizeTmpBytes
106            + VmpApply<B>
107            + VmpApplyAdd<B>
108            + VecZnxDftFromVecZnx<B>
109            + VecZnxDftToVecZnxBigConsume<B>
110            + VecZnxBigAddSmallInplace<B>
111            + VecZnxBigNormalize<B>
112            + VecZnxBigAutomorphismInplace<B>,
113        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
114    {
115        #[cfg(debug_assertions)]
116        {
117            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
118        }
119        let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
120        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch1);
121        (0..self.cols()).for_each(|i| {
122            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
123            module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i);
124            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
125        })
126    }
127
128    pub fn automorphism_add_inplace<DataRhs: DataRef, B: Backend>(
129        &mut self,
130        module: &Module<B>,
131        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
132        scratch: &mut Scratch<B>,
133    ) where
134        Module<B>: VecZnxDftAllocBytes
135            + VmpApplyTmpBytes
136            + VecZnxBigNormalizeTmpBytes
137            + VmpApply<B>
138            + VmpApplyAdd<B>
139            + VecZnxDftFromVecZnx<B>
140            + VecZnxDftToVecZnxBigConsume<B>
141            + VecZnxBigAddSmallInplace<B>
142            + VecZnxBigNormalize<B>
143            + VecZnxBigAutomorphismInplace<B>,
144        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
145    {
146        unsafe {
147            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
148            self.automorphism_add(module, &*self_ptr, rhs, scratch);
149        }
150    }
151
152    pub fn automorphism_sub_ab<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
153        &mut self,
154        module: &Module<B>,
155        lhs: &GLWECiphertext<DataLhs>,
156        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
157        scratch: &mut Scratch<B>,
158    ) where
159        Module<B>: VecZnxDftAllocBytes
160            + VmpApplyTmpBytes
161            + VecZnxBigNormalizeTmpBytes
162            + VmpApply<B>
163            + VmpApplyAdd<B>
164            + VecZnxDftFromVecZnx<B>
165            + VecZnxDftToVecZnxBigConsume<B>
166            + VecZnxBigAddSmallInplace<B>
167            + VecZnxBigNormalize<B>
168            + VecZnxBigAutomorphismInplace<B>
169            + VecZnxBigSubSmallAInplace<B>,
170        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
171    {
172        #[cfg(debug_assertions)]
173        {
174            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
175        }
176        let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
177        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch1);
178        (0..self.cols()).for_each(|i| {
179            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
180            module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i);
181            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
182        })
183    }
184
185    pub fn automorphism_sub_ab_inplace<DataRhs: DataRef, B: Backend>(
186        &mut self,
187        module: &Module<B>,
188        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
189        scratch: &mut Scratch<B>,
190    ) where
191        Module<B>: VecZnxDftAllocBytes
192            + VmpApplyTmpBytes
193            + VecZnxBigNormalizeTmpBytes
194            + VmpApply<B>
195            + VmpApplyAdd<B>
196            + VecZnxDftFromVecZnx<B>
197            + VecZnxDftToVecZnxBigConsume<B>
198            + VecZnxBigAddSmallInplace<B>
199            + VecZnxBigNormalize<B>
200            + VecZnxBigAutomorphismInplace<B>
201            + VecZnxBigSubSmallAInplace<B>,
202        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
203    {
204        unsafe {
205            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
206            self.automorphism_sub_ab(module, &*self_ptr, rhs, scratch);
207        }
208    }
209
210    pub fn automorphism_sub_ba<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
211        &mut self,
212        module: &Module<B>,
213        lhs: &GLWECiphertext<DataLhs>,
214        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
215        scratch: &mut Scratch<B>,
216    ) where
217        Module<B>: VecZnxDftAllocBytes
218            + VmpApplyTmpBytes
219            + VecZnxBigNormalizeTmpBytes
220            + VmpApply<B>
221            + VmpApplyAdd<B>
222            + VecZnxDftFromVecZnx<B>
223            + VecZnxDftToVecZnxBigConsume<B>
224            + VecZnxBigAddSmallInplace<B>
225            + VecZnxBigNormalize<B>
226            + VecZnxBigAutomorphismInplace<B>
227            + VecZnxBigSubSmallBInplace<B>,
228        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
229    {
230        #[cfg(debug_assertions)]
231        {
232            self.assert_keyswitch(module, lhs, &rhs.key, scratch);
233        }
234        let (res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), self.cols(), rhs.size()); // TODO: optimise size
235        let mut res_big: VecZnxBig<_, B> = lhs.keyswitch_internal(module, res_dft, &rhs.key, scratch1);
236        (0..self.cols()).for_each(|i| {
237            module.vec_znx_big_automorphism_inplace(rhs.p(), &mut res_big, i);
238            module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i);
239            module.vec_znx_big_normalize(self.basek(), &mut self.data, i, &res_big, i, scratch1);
240        })
241    }
242
243    pub fn automorphism_sub_ba_inplace<DataRhs: DataRef, B: Backend>(
244        &mut self,
245        module: &Module<B>,
246        rhs: &GGLWEAutomorphismKeyPrepared<DataRhs, B>,
247        scratch: &mut Scratch<B>,
248    ) where
249        Module<B>: VecZnxDftAllocBytes
250            + VmpApplyTmpBytes
251            + VecZnxBigNormalizeTmpBytes
252            + VmpApply<B>
253            + VmpApplyAdd<B>
254            + VecZnxDftFromVecZnx<B>
255            + VecZnxDftToVecZnxBigConsume<B>
256            + VecZnxBigAddSmallInplace<B>
257            + VecZnxBigNormalize<B>
258            + VecZnxBigAutomorphismInplace<B>
259            + VecZnxBigSubSmallBInplace<B>,
260        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
261    {
262        unsafe {
263            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
264            self.automorphism_sub_ba(module, &*self_ptr, rhs, scratch);
265        }
266    }
267}
poulpy_core/automorphism/glwe_ct.rs

poulpy_core/automorphism/
glwe_ct.rs
