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 cols_out: usize = rhs.rank_out() + 1;
178
179 // Space fr normalized VMP result outside of DFT domain
180 let (res_idft_data, scratch1) = scratch.tmp_vec_znx(module, cols_out, lhs.size());
181
182 let mut res_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
183 data: res_idft_data,
184 basek: self.basek,
185 k: self.k,
186 };
187
188 res_idft.keyswitch_from_fourier(module, self, rhs, scratch1);
189
190 (0..cols_out).for_each(|i| {
191 module.vec_znx_dft(self, i, &res_idft, i);
192 });
193 }
194
195 pub fn keyswitch_inplace<DataRhs>(
196 &mut self,
197 module: &Module<FFT64>,
198 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
199 scratch: &mut Scratch,
200 ) where
201 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
202 {
203 unsafe {
204 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
205 self.keyswitch(&module, &*self_ptr, rhs, scratch);
206 }
207 }
208
209 pub fn external_product<DataLhs, DataRhs>(
210 &mut self,
211 module: &Module<FFT64>,
212 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
213 rhs: &GGSWCiphertext<DataRhs, FFT64>,
214 scratch: &mut Scratch,
215 ) where
216 VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
217 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
218 {
219 let basek: usize = self.basek();
220
221 #[cfg(debug_assertions)]
222 {
223 assert_eq!(rhs.rank(), lhs.rank());
224 assert_eq!(rhs.rank(), self.rank());
225 assert_eq!(self.basek(), basek);
226 assert_eq!(lhs.basek(), basek);
227 assert_eq!(rhs.n(), module.n());
228 assert_eq!(self.n(), module.n());
229 assert_eq!(lhs.n(), module.n());
230 }
231
232 let cols: usize = rhs.rank() + 1;
233
234 // Space for VMP result in DFT domain and high precision
235 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
236
237 {
238 module.vmp_apply(&mut res_dft, lhs, rhs, scratch1);
239 }
240
241 // VMP result in high precision
242 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
243
244 // Space for VMP result normalized
245 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
246 (0..cols).for_each(|i| {
247 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
248 module.vec_znx_dft(self, i, &res_small, i);
249 });
250 }
251
252 pub fn external_product_inplace<DataRhs>(
253 &mut self,
254 module: &Module<FFT64>,
255 rhs: &GGSWCiphertext<DataRhs, FFT64>,
256 scratch: &mut Scratch,
257 ) where
258 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
259 {
260 unsafe {
261 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
262 self.external_product(&module, &*self_ptr, rhs, scratch);
263 }
264 }
265}
266
267impl<DataSelf> GLWECiphertextFourier<DataSelf, FFT64>
268where
269 VecZnxDft<DataSelf, FFT64>: VecZnxDftToRef<FFT64>,
270{
271 pub fn decrypt<DataPt, DataSk>(
272 &self,
273 module: &Module<FFT64>,
274 pt: &mut GLWEPlaintext<DataPt>,
275 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
276 scratch: &mut Scratch,
277 ) where
278 VecZnx<DataPt>: VecZnxToMut + VecZnxToRef,
279 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
280 {
281 #[cfg(debug_assertions)]
282 {
283 assert_eq!(self.rank(), sk_dft.rank());
284 assert_eq!(self.n(), module.n());
285 assert_eq!(pt.n(), module.n());
286 assert_eq!(sk_dft.n(), module.n());
287 }
288
289 let cols = self.rank() + 1;
290
291 let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
292 pt_big.zero();
293
294 {
295 (1..cols).for_each(|i| {
296 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
297 module.svp_apply(&mut ci_dft, 0, sk_dft, i - 1, self, i);
298 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
299 module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
300 });
301 }
302
303 {
304 let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
305 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
306 module.vec_znx_idft(&mut c0_big, 0, self, 0, scratch_2);
307 module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
308 }
309
310 // pt = norm(BIG(m + e))
311 module.vec_znx_big_normalize(self.basek(), pt, 0, &mut pt_big, 0, scratch_1);
312
313 pt.basek = self.basek();
314 pt.k = pt.k().min(self.k());
315 }
316
317 #[allow(dead_code)]
318 pub(crate) fn idft<DataRes>(&self, module: &Module<FFT64>, res: &mut GLWECiphertext<DataRes>, scratch: &mut Scratch)
319 where
320 GLWECiphertext<DataRes>: VecZnxToMut,
321 {
322 #[cfg(debug_assertions)]
323 {
324 assert_eq!(self.rank(), res.rank());
325 assert_eq!(self.basek(), res.basek())
326 }
327
328 let min_size: usize = self.size().min(res.size());
329
330 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
331
332 (0..self.rank() + 1).for_each(|i| {
333 module.vec_znx_idft(&mut res_big, 0, self, i, scratch1);
334 module.vec_znx_big_normalize(self.basek(), res, i, &res_big, 0, scratch1);
335 });
336 }
337}