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