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
203 // Space for VMP result in DFT domain and high precision
204 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
205 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, (lhs.size() + digits - 1) / digits);
206
207 {
208 (0..digits).for_each(|di| {
209
210 a_dft.set_size((lhs.size() + di) / digits);
211 res_dft.set_size(rhs.size() - (digits - di - 1));
212
213 (0..cols).for_each(|col_i| {
214 module.vec_znx_dft_copy(digits, digits - 1 - di, &mut a_dft, col_i, &lhs.data, col_i);
215 });
216
217 if di == 0 {
218 module.vmp_apply(&mut res_dft, &a_dft, &rhs.data, scratch2);
219 } else {
220 module.vmp_apply_add(&mut res_dft, &a_dft, &rhs.data, di, scratch2);
221 }
222 });
223 }
224
225 // VMP result in high precision
226 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
227
228 // Space for VMP result normalized
229 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
230 (0..cols).for_each(|i| {
231 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
232 module.vec_znx_dft(1, 0, &mut self.data, i, &res_small, i);
233 });
234 }
235
236 pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
237 &mut self,
238 module: &Module<FFT64>,
239 rhs: &GGSWCiphertext<DataRhs, FFT64>,
240 scratch: &mut Scratch,
241 ) {
242 unsafe {
243 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
244 self.external_product(&module, &*self_ptr, rhs, scratch);
245 }
246 }
247}
248
249impl<DataSelf: AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
250 pub fn decrypt<DataPt: AsRef<[u8]> + AsMut<[u8]>, DataSk: AsRef<[u8]>>(
251 &self,
252 module: &Module<FFT64>,
253 pt: &mut GLWEPlaintext<DataPt>,
254 sk: &GLWESecret<DataSk, FFT64>,
255 scratch: &mut Scratch,
256 ) {
257 #[cfg(debug_assertions)]
258 {
259 assert_eq!(self.rank(), sk.rank());
260 assert_eq!(self.n(), module.n());
261 assert_eq!(pt.n(), module.n());
262 assert_eq!(sk.n(), module.n());
263 }
264
265 let cols = self.rank() + 1;
266
267 let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
268 pt_big.zero();
269
270 {
271 (1..cols).for_each(|i| {
272 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
273 module.svp_apply(&mut ci_dft, 0, &sk.data_fourier, i - 1, &self.data, i);
274 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
275 module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
276 });
277 }
278
279 {
280 let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
281 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
282 module.vec_znx_idft(&mut c0_big, 0, &self.data, 0, scratch_2);
283 module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
284 }
285
286 // pt = norm(BIG(m + e))
287 module.vec_znx_big_normalize(self.basek(), &mut pt.data, 0, &mut pt_big, 0, scratch_1);
288
289 pt.basek = self.basek();
290 pt.k = pt.k().min(self.k());
291 }
292
293 #[allow(dead_code)]
294 pub(crate) fn idft<DataRes: AsRef<[u8]> + AsMut<[u8]>>(
295 &self,
296 module: &Module<FFT64>,
297 res: &mut GLWECiphertext<DataRes>,
298 scratch: &mut Scratch,
299 ) {
300 #[cfg(debug_assertions)]
301 {
302 assert_eq!(self.rank(), res.rank());
303 assert_eq!(self.basek(), res.basek())
304 }
305
306 let min_size: usize = self.size().min(res.size());
307
308 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
309
310 (0..self.rank() + 1).for_each(|i| {
311 module.vec_znx_idft(&mut res_big, 0, &self.data, i, scratch1);
312 module.vec_znx_big_normalize(self.basek(), &mut res.data, i, &res_big, 0, scratch1);
313 });
314 }
315}