poulpy/core/src/glwe_ciphertext_fourier.rs
core/src/glwe_ciphertext_fourier.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftScratch, MatZnxDftToRef, Module, ScalarZnxDft, ScalarZnxDftOps,
3 ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft,
4 VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxToMut, VecZnxToRef, ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::Infos, ggsw_ciphertext::GGSWCiphertext, glwe_ciphertext::GLWECiphertext, glwe_plaintext::GLWEPlaintext,
10 keys::SecretKeyFourier, keyswitch_key::GLWESwitchingKey, utils::derive_size,
11};
12
13pub struct GLWECiphertextFourier<C, B: Backend> {
14 pub data: VecZnxDft<C, B>,
15 pub basek: usize,
16 pub k: usize,
17}
18
19impl<B: Backend> GLWECiphertextFourier<Vec<u8>, B> {
20 pub fn new(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
21 Self {
22 data: module.new_vec_znx_dft(rank + 1, derive_size(basek, k)),
23 basek: basek,
24 k: k,
25 }
26 }
27}
28
29impl<T, B: Backend> Infos for GLWECiphertextFourier<T, B> {
30 type Inner = VecZnxDft<T, B>;
31
32 fn inner(&self) -> &Self::Inner {
33 &self.data
34 }
35
36 fn basek(&self) -> usize {
37 self.basek
38 }
39
40 fn k(&self) -> usize {
41 self.k
42 }
43}
44
45impl<T, B: Backend> GLWECiphertextFourier<T, B> {
46 pub fn rank(&self) -> usize {
47 self.cols() - 1
48 }
49}
50
51impl<C, B: Backend> VecZnxDftToMut<B> for GLWECiphertextFourier<C, B>
52where
53 VecZnxDft<C, B>: VecZnxDftToMut<B>,
54{
55 fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
56 self.data.to_mut()
57 }
58}
59
60impl<C, B: Backend> VecZnxDftToRef<B> for GLWECiphertextFourier<C, B>
61where
62 VecZnxDft<C, B>: VecZnxDftToRef<B>,
63{
64 fn to_ref(&self) -> VecZnxDft<&[u8], B> {
65 self.data.to_ref()
66 }
67}
68
69impl GLWECiphertextFourier<Vec<u8>, FFT64> {
70 #[allow(dead_code)]
71 pub(crate) fn idft_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
72 module.bytes_of_vec_znx(1, size) + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())
73 }
74
75 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, ct_size: usize) -> usize {
76 module.bytes_of_vec_znx(rank + 1, ct_size) + GLWECiphertext::encrypt_sk_scratch_space(module, ct_size)
77 }
78
79 pub fn decrypt_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
80 (module.vec_znx_big_normalize_tmp_bytes()
81 | module.bytes_of_vec_znx_dft(1, ct_size)
82 | (module.bytes_of_vec_znx_big(1, ct_size) + module.vec_znx_idft_tmp_bytes()))
83 + module.bytes_of_vec_znx_big(1, ct_size)
84 }
85
86 pub fn keyswitch_scratch_space(
87 module: &Module<FFT64>,
88 out_size: usize,
89 out_rank: usize,
90 in_size: usize,
91 in_rank: usize,
92 ksk_size: usize,
93 ) -> usize {
94 let res_dft: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
95
96 let vmp = module.bytes_of_vec_znx_dft(in_rank, in_size)
97 + module.vmp_apply_tmp_bytes(
98 out_size,
99 in_size,
100 in_size,
101 in_rank + 1,
102 out_rank + 1,
103 ksk_size,
104 );
105 let res_small: usize = module.bytes_of_vec_znx(out_rank + 1, out_size);
106 let add_a0: usize = module.bytes_of_vec_znx_big(1, in_size) + module.vec_znx_idft_tmp_bytes();
107 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
108
109 res_dft + (vmp | add_a0 | (res_small + normalize))
110 }
111
112 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
113 Self::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, ksk_size)
114 }
115
116 pub fn external_product_scratch_space(
117 module: &Module<FFT64>,
118 out_size: usize,
119 in_size: usize,
120 ggsw_size: usize,
121 rank: usize,
122 ) -> usize {
123 let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
124 let vmp: usize = module.vmp_apply_tmp_bytes(out_size, in_size, in_size, rank + 1, rank + 1, ggsw_size);
125 let res_small: usize = module.bytes_of_vec_znx(rank + 1, out_size);
126 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
127
128 res_dft + (vmp | (res_small + normalize))
129 }
130
131 pub fn external_product_inplace_scratch_space(
132 module: &Module<FFT64>,
133 out_size: usize,
134 ggsw_size: usize,
135 rank: usize,
136 ) -> usize {
137 Self::external_product_scratch_space(module, out_size, out_size, ggsw_size, rank)
138 }
139}
140
141impl<DataSelf> GLWECiphertextFourier<DataSelf, FFT64>
142where
143 VecZnxDft<DataSelf, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
144{
145 pub fn encrypt_zero_sk<DataSk>(
146 &mut self,
147 module: &Module<FFT64>,
148 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
149 source_xa: &mut Source,
150 source_xe: &mut Source,
151 sigma: f64,
152 scratch: &mut Scratch,
153 ) where
154 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
155 {
156 let (vec_znx_tmp, scratch_1) = scratch.tmp_vec_znx(module, self.rank() + 1, self.size());
157 let mut ct_idft = GLWECiphertext {
158 data: vec_znx_tmp,
159 basek: self.basek,
160 k: self.k,
161 };
162 ct_idft.encrypt_zero_sk(module, sk_dft, source_xa, source_xe, sigma, scratch_1);
163
164 ct_idft.dft(module, self);
165 }
166
167 pub fn keyswitch<DataLhs, DataRhs>(
168 &mut self,
169 module: &Module<FFT64>,
170 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
171 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
172 scratch: &mut Scratch,
173 ) where
174 VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
175 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
176 {
177 let basek: usize = self.basek();
178
179 #[cfg(debug_assertions)]
180 {
181 assert_eq!(lhs.rank(), rhs.rank_in());
182 assert_eq!(self.rank(), rhs.rank_out());
183 assert_eq!(self.basek(), basek);
184 assert_eq!(lhs.basek(), basek);
185 assert_eq!(rhs.n(), module.n());
186 assert_eq!(self.n(), module.n());
187 assert_eq!(lhs.n(), module.n());
188 assert!(
189 scratch.available()
190 >= GLWECiphertextFourier::keyswitch_scratch_space(
191 module,
192 self.size(),
193 self.rank(),
194 lhs.size(),
195 lhs.rank(),
196 rhs.size(),
197 )
198 );
199 }
200
201 let cols_in: usize = rhs.rank_in();
202 let cols_out: usize = rhs.rank_out() + 1;
203
204 // Buffer of the result of VMP in DFT
205 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
206
207 {
208 // Applies VMP
209 let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, lhs.size());
210 (0..cols_in).for_each(|col_i| {
211 module.vec_znx_dft_copy(&mut ai_dft, col_i, lhs, col_i + 1);
212 });
213 module.vmp_apply(&mut res_dft, &ai_dft, rhs, scratch2);
214 }
215
216 // Switches result of VMP outside of DFT
217 let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
218
219 {
220 // Switches lhs 0-th outside of DFT domain and adds on
221 let (mut a0_big, scratch2) = scratch1.tmp_vec_znx_big(module, 1, lhs.size());
222 module.vec_znx_idft(&mut a0_big, 0, lhs, 0, scratch2);
223 module.vec_znx_big_add_inplace(&mut res_big, 0, &a0_big, 0);
224 }
225
226 // Space fr normalized VMP result outside of DFT domain
227 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols_out, lhs.size());
228 (0..cols_out).for_each(|i| {
229 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
230 module.vec_znx_dft(self, i, &res_small, i);
231 });
232 }
233
234 pub fn keyswitch_inplace<DataRhs>(
235 &mut self,
236 module: &Module<FFT64>,
237 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
238 scratch: &mut Scratch,
239 ) where
240 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
241 {
242 unsafe {
243 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
244 self.keyswitch(&module, &*self_ptr, rhs, scratch);
245 }
246 }
247
248 pub fn external_product<DataLhs, DataRhs>(
249 &mut self,
250 module: &Module<FFT64>,
251 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
252 rhs: &GGSWCiphertext<DataRhs, FFT64>,
253 scratch: &mut Scratch,
254 ) where
255 VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
256 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
257 {
258 let basek: usize = self.basek();
259
260 #[cfg(debug_assertions)]
261 {
262 assert_eq!(rhs.rank(), lhs.rank());
263 assert_eq!(rhs.rank(), self.rank());
264 assert_eq!(self.basek(), basek);
265 assert_eq!(lhs.basek(), basek);
266 assert_eq!(rhs.n(), module.n());
267 assert_eq!(self.n(), module.n());
268 assert_eq!(lhs.n(), module.n());
269 }
270
271 let cols: usize = rhs.rank() + 1;
272
273 // Space for VMP result in DFT domain and high precision
274 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
275
276 {
277 module.vmp_apply(&mut res_dft, lhs, rhs, scratch1);
278 }
279
280 // VMP result in high precision
281 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
282
283 // Space for VMP result normalized
284 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
285 (0..cols).for_each(|i| {
286 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
287 module.vec_znx_dft(self, i, &res_small, i);
288 });
289 }
290
291 pub fn external_product_inplace<DataRhs>(
292 &mut self,
293 module: &Module<FFT64>,
294 rhs: &GGSWCiphertext<DataRhs, FFT64>,
295 scratch: &mut Scratch,
296 ) where
297 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
298 {
299 unsafe {
300 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
301 self.external_product(&module, &*self_ptr, rhs, scratch);
302 }
303 }
304}
305
306impl<DataSelf> GLWECiphertextFourier<DataSelf, FFT64>
307where
308 VecZnxDft<DataSelf, FFT64>: VecZnxDftToRef<FFT64>,
309{
310 pub fn decrypt<DataPt, DataSk>(
311 &self,
312 module: &Module<FFT64>,
313 pt: &mut GLWEPlaintext<DataPt>,
314 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
315 scratch: &mut Scratch,
316 ) where
317 VecZnx<DataPt>: VecZnxToMut + VecZnxToRef,
318 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
319 {
320 #[cfg(debug_assertions)]
321 {
322 assert_eq!(self.rank(), sk_dft.rank());
323 assert_eq!(self.n(), module.n());
324 assert_eq!(pt.n(), module.n());
325 assert_eq!(sk_dft.n(), module.n());
326 }
327
328 let cols = self.rank() + 1;
329
330 let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
331 pt_big.zero();
332
333 {
334 (1..cols).for_each(|i| {
335 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
336 module.svp_apply(&mut ci_dft, 0, sk_dft, i - 1, self, i);
337 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
338 module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
339 });
340 }
341
342 {
343 let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
344 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
345 module.vec_znx_idft(&mut c0_big, 0, self, 0, scratch_2);
346 module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
347 }
348
349 // pt = norm(BIG(m + e))
350 module.vec_znx_big_normalize(self.basek(), pt, 0, &mut pt_big, 0, scratch_1);
351
352 pt.basek = self.basek();
353 pt.k = pt.k().min(self.k());
354 }
355
356 #[allow(dead_code)]
357 pub(crate) fn idft<DataRes>(&self, module: &Module<FFT64>, res: &mut GLWECiphertext<DataRes>, scratch: &mut Scratch)
358 where
359 GLWECiphertext<DataRes>: VecZnxToMut,
360 {
361 #[cfg(debug_assertions)]
362 {
363 assert_eq!(self.rank(), res.rank());
364 assert_eq!(self.basek(), res.basek())
365 }
366
367 let min_size: usize = self.size().min(res.size());
368
369 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
370
371 (0..self.rank() + 1).for_each(|i| {
372 module.vec_znx_idft(&mut res_big, 0, self, i, scratch1);
373 module.vec_znx_big_normalize(self.basek(), res, i, &res_big, 0, scratch1);
374 });
375 }
376}