poulpy/core/src/glwe_ciphertext_fourier.rs
core/src/glwe_ciphertext_fourier.rsRaw
1use backend::{
2 Backend, FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc,
3 VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{GGSWCiphertext, GLWECiphertext, GLWEPlaintext, GLWESecret, GLWESwitchingKey, Infos, ScratchCore, div_ceil};
8
9pub struct GLWECiphertextFourier<C, B: Backend> {
10 pub data: VecZnxDft<C, B>,
11 pub basek: usize,
12 pub k: usize,
13}
14
15impl<B: Backend> GLWECiphertextFourier<Vec<u8>, B> {
16 pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
17 Self {
18 data: module.new_vec_znx_dft(rank + 1, div_ceil(k, basek)),
19 basek: basek,
20 k: k,
21 }
22 }
23
24 pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
25 module.bytes_of_vec_znx_dft(rank + 1, div_ceil(k, basek))
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 GLWECiphertextFourier<Vec<u8>, FFT64> {
52 #[allow(dead_code)]
53 pub(crate) fn idft_scratch_space(module: &Module<FFT64>, basek: usize, k: usize) -> usize {
54 module.bytes_of_vec_znx(1, div_ceil(k, basek))
55 + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())
56 }
57
58 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
59 module.bytes_of_vec_znx(rank + 1, div_ceil(k, basek)) + GLWECiphertext::encrypt_sk_scratch_space(module, basek, k)
60 }
61
62 pub fn decrypt_scratch_space(module: &Module<FFT64>, basek: usize, k: usize) -> usize {
63 let size: usize = div_ceil(k, basek);
64 (module.vec_znx_big_normalize_tmp_bytes()
65 | module.bytes_of_vec_znx_dft(1, size)
66 | (module.bytes_of_vec_znx_big(1, size) + module.vec_znx_idft_tmp_bytes()))
67 + module.bytes_of_vec_znx_big(1, size)
68 }
69
70 pub fn keyswitch_scratch_space(
71 module: &Module<FFT64>,
72 basek: usize,
73 k_out: usize,
74 k_in: usize,
75 k_ksk: usize,
76 digits: usize,
77 rank_in: usize,
78 rank_out: usize,
79 ) -> usize {
80 GLWECiphertext::bytes_of(module, basek, k_out, rank_out)
81 + GLWECiphertext::keyswitch_from_fourier_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank_in, rank_out)
82 }
83
84 pub fn keyswitch_inplace_scratch_space(
85 module: &Module<FFT64>,
86 basek: usize,
87 k_out: usize,
88 k_ksk: usize,
89 digits: usize,
90 rank: usize,
91 ) -> usize {
92 Self::keyswitch_scratch_space(module, basek, k_out, k_out, k_ksk, digits, rank, rank)
93 }
94
95 // WARNING TODO: UPDATE
96 pub fn external_product_scratch_space(
97 module: &Module<FFT64>,
98 basek: usize,
99 _k_out: usize,
100 k_in: usize,
101 k_ggsw: usize,
102 digits: usize,
103 rank: usize,
104 ) -> usize {
105 let ggsw_size: usize = div_ceil(k_ggsw, basek);
106 let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, ggsw_size);
107 let in_size: usize = div_ceil(div_ceil(k_in, basek), digits);
108 let ggsw_size: usize = div_ceil(k_ggsw, basek);
109 let vmp: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size)
110 + module.vmp_apply_tmp_bytes(ggsw_size, in_size, in_size, rank + 1, rank + 1, ggsw_size);
111 let res_small: usize = module.bytes_of_vec_znx(rank + 1, ggsw_size);
112 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
113 res_dft + (vmp | (res_small + normalize))
114 }
115
116 pub fn external_product_inplace_scratch_space(
117 module: &Module<FFT64>,
118 basek: usize,
119 k_out: usize,
120 k_ggsw: usize,
121 digits: usize,
122 rank: usize,
123 ) -> usize {
124 Self::external_product_scratch_space(module, basek, k_out, k_out, k_ggsw, digits, rank)
125 }
126}
127
128impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
129 pub fn encrypt_zero_sk<DataSk: AsRef<[u8]>>(
130 &mut self,
131 module: &Module<FFT64>,
132 sk: &GLWESecret<DataSk, FFT64>,
133 source_xa: &mut Source,
134 source_xe: &mut Source,
135 sigma: f64,
136 scratch: &mut Scratch,
137 ) {
138 let (mut tmp_ct, scratch1) = scratch.tmp_glwe_ct(module, self.basek(), self.k(), self.rank());
139 tmp_ct.encrypt_zero_sk(module, sk, source_xa, source_xe, sigma, scratch1);
140 tmp_ct.dft(module, self);
141 }
142
143 pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
144 &mut self,
145 module: &Module<FFT64>,
146 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
147 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
148 scratch: &mut Scratch,
149 ) {
150 let (mut tmp_ct, scratch1) = scratch.tmp_glwe_ct(module, self.basek(), self.k(), self.rank());
151 tmp_ct.keyswitch_from_fourier(module, lhs, rhs, scratch1);
152 tmp_ct.dft(module, self);
153 }
154
155 pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
156 &mut self,
157 module: &Module<FFT64>,
158 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
159 scratch: &mut Scratch,
160 ) {
161 unsafe {
162 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
163 self.keyswitch(&module, &*self_ptr, rhs, scratch);
164 }
165 }
166
167 pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
168 &mut self,
169 module: &Module<FFT64>,
170 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
171 rhs: &GGSWCiphertext<DataRhs, FFT64>,
172 scratch: &mut Scratch,
173 ) {
174 let basek: usize = self.basek();
175
176 #[cfg(debug_assertions)]
177 {
178 assert_eq!(rhs.rank(), lhs.rank());
179 assert_eq!(rhs.rank(), self.rank());
180 assert_eq!(self.basek(), basek);
181 assert_eq!(lhs.basek(), basek);
182 assert_eq!(rhs.n(), module.n());
183 assert_eq!(self.n(), module.n());
184 assert_eq!(lhs.n(), module.n());
185 assert!(
186 scratch.available()
187 >= GLWECiphertextFourier::external_product_scratch_space(
188 module,
189 self.basek(),
190 self.k(),
191 lhs.k(),
192 rhs.k(),
193 rhs.digits(),
194 rhs.rank(),
195 )
196 );
197 }
198
199 let cols: usize = rhs.rank() + 1;
200 let digits = rhs.digits();
201
202 // Space for VMP result in DFT domain and high precision
203 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
204 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, (lhs.size() + digits - 1) / digits);
205
206 {
207 (0..digits).for_each(|di| {
208 a_dft.set_size((lhs.size() + di) / digits);
209
210 // Small optimization for digits > 2
211 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
212 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
213 // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
214 // It is possible to further ignore the last digits-1 limbs, but this introduce
215 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
216 // noise is kept with respect to the ideal functionality.
217 res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
218
219 (0..cols).for_each(|col_i| {
220 module.vec_znx_dft_copy(digits, digits - 1 - di, &mut a_dft, col_i, &lhs.data, col_i);
221 });
222
223 if di == 0 {
224 module.vmp_apply(&mut res_dft, &a_dft, &rhs.data, scratch2);
225 } else {
226 module.vmp_apply_add(&mut res_dft, &a_dft, &rhs.data, di, scratch2);
227 }
228 });
229 }
230
231 // VMP result in high precision
232 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
233
234 // Space for VMP result normalized
235 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
236 (0..cols).for_each(|i| {
237 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
238 module.vec_znx_dft(1, 0, &mut self.data, i, &res_small, i);
239 });
240 }
241
242 pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
243 &mut self,
244 module: &Module<FFT64>,
245 rhs: &GGSWCiphertext<DataRhs, FFT64>,
246 scratch: &mut Scratch,
247 ) {
248 unsafe {
249 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
250 self.external_product(&module, &*self_ptr, rhs, scratch);
251 }
252 }
253}
254
255impl<DataSelf: AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
256 pub fn decrypt<DataPt: AsRef<[u8]> + AsMut<[u8]>, DataSk: AsRef<[u8]>>(
257 &self,
258 module: &Module<FFT64>,
259 pt: &mut GLWEPlaintext<DataPt>,
260 sk: &GLWESecret<DataSk, FFT64>,
261 scratch: &mut Scratch,
262 ) {
263 #[cfg(debug_assertions)]
264 {
265 assert_eq!(self.rank(), sk.rank());
266 assert_eq!(self.n(), module.n());
267 assert_eq!(pt.n(), module.n());
268 assert_eq!(sk.n(), module.n());
269 }
270
271 let cols = self.rank() + 1;
272
273 let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
274 pt_big.zero();
275
276 {
277 (1..cols).for_each(|i| {
278 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
279 module.svp_apply(&mut ci_dft, 0, &sk.data_fourier, i - 1, &self.data, i);
280 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
281 module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
282 });
283 }
284
285 {
286 let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
287 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
288 module.vec_znx_idft(&mut c0_big, 0, &self.data, 0, scratch_2);
289 module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
290 }
291
292 // pt = norm(BIG(m + e))
293 module.vec_znx_big_normalize(self.basek(), &mut pt.data, 0, &mut pt_big, 0, scratch_1);
294
295 pt.basek = self.basek();
296 pt.k = pt.k().min(self.k());
297 }
298
299 #[allow(dead_code)]
300 pub(crate) fn idft<DataRes: AsRef<[u8]> + AsMut<[u8]>>(
301 &self,
302 module: &Module<FFT64>,
303 res: &mut GLWECiphertext<DataRes>,
304 scratch: &mut Scratch,
305 ) {
306 #[cfg(debug_assertions)]
307 {
308 assert_eq!(self.rank(), res.rank());
309 assert_eq!(self.basek(), res.basek())
310 }
311
312 let min_size: usize = self.size().min(res.size());
313
314 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
315
316 (0..self.rank() + 1).for_each(|i| {
317 module.vec_znx_idft(&mut res_big, 0, &self.data, i, scratch1);
318 module.vec_znx_big_normalize(self.basek(), &mut res.data, i, &res_big, 0, scratch1);
319 });
320 }
321}