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