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::{
8 elem::Infos, ggsw_ciphertext::GGSWCiphertext, glwe_ciphertext::GLWECiphertext, glwe_plaintext::GLWEPlaintext,
9 keys::SecretKeyFourier, keyswitch_key::GLWESwitchingKey, utils::derive_size,
10};
11
12pub struct GLWECiphertextFourier<C, B: Backend> {
13 pub data: VecZnxDft<C, B>,
14 pub basek: usize,
15 pub k: usize,
16}
17
18impl<B: Backend> GLWECiphertextFourier<Vec<u8>, B> {
19 pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
20 Self {
21 data: module.new_vec_znx_dft(rank + 1, derive_size(basek, k)),
22 basek: basek,
23 k: k,
24 }
25 }
26
27 pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
28 module.bytes_of_vec_znx_dft(rank + 1, derive_size(basek, k))
29 }
30}
31
32impl<T, B: Backend> Infos for GLWECiphertextFourier<T, B> {
33 type Inner = VecZnxDft<T, B>;
34
35 fn inner(&self) -> &Self::Inner {
36 &self.data
37 }
38
39 fn basek(&self) -> usize {
40 self.basek
41 }
42
43 fn k(&self) -> usize {
44 self.k
45 }
46}
47
48impl<T, B: Backend> GLWECiphertextFourier<T, B> {
49 pub fn rank(&self) -> usize {
50 self.cols() - 1
51 }
52}
53
54impl GLWECiphertextFourier<Vec<u8>, FFT64> {
55 #[allow(dead_code)]
56 pub(crate) fn idft_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
57 module.bytes_of_vec_znx(1, size) + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())
58 }
59
60 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, ct_size: usize) -> usize {
61 module.bytes_of_vec_znx(rank + 1, ct_size) + GLWECiphertext::encrypt_sk_scratch_space(module, ct_size)
62 }
63
64 pub fn decrypt_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
65 (module.vec_znx_big_normalize_tmp_bytes()
66 | module.bytes_of_vec_znx_dft(1, ct_size)
67 | (module.bytes_of_vec_znx_big(1, ct_size) + module.vec_znx_idft_tmp_bytes()))
68 + module.bytes_of_vec_znx_big(1, ct_size)
69 }
70
71 pub fn keyswitch_scratch_space(
72 module: &Module<FFT64>,
73 out_size: usize,
74 out_rank: usize,
75 in_size: usize,
76 in_rank: usize,
77 ksk_size: usize,
78 ) -> usize {
79 module.bytes_of_vec_znx(out_rank + 1, out_size)
80 + GLWECiphertext::keyswitch_from_fourier_scratch_space(module, out_size, out_rank, in_size, in_rank, ksk_size)
81 }
82
83 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
84 Self::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, ksk_size)
85 }
86
87 pub fn external_product_scratch_space(
88 module: &Module<FFT64>,
89 out_size: usize,
90 in_size: usize,
91 ggsw_size: usize,
92 rank: usize,
93 ) -> usize {
94 let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
95 let vmp: usize = module.vmp_apply_tmp_bytes(out_size, in_size, in_size, rank + 1, rank + 1, ggsw_size);
96 let res_small: usize = module.bytes_of_vec_znx(rank + 1, out_size);
97 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
98
99 res_dft + (vmp | (res_small + normalize))
100 }
101
102 pub fn external_product_inplace_scratch_space(
103 module: &Module<FFT64>,
104 out_size: usize,
105 ggsw_size: usize,
106 rank: usize,
107 ) -> usize {
108 Self::external_product_scratch_space(module, out_size, out_size, ggsw_size, rank)
109 }
110}
111
112impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
113 pub fn encrypt_zero_sk<DataSk: AsRef<[u8]>>(
114 &mut self,
115 module: &Module<FFT64>,
116 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
117 source_xa: &mut Source,
118 source_xe: &mut Source,
119 sigma: f64,
120 scratch: &mut Scratch,
121 ) {
122 let (vec_znx_tmp, scratch_1) = scratch.tmp_vec_znx(module, self.rank() + 1, self.size());
123 let mut ct_idft = GLWECiphertext {
124 data: vec_znx_tmp,
125 basek: self.basek,
126 k: self.k,
127 };
128 ct_idft.encrypt_zero_sk(module, sk_dft, source_xa, source_xe, sigma, scratch_1);
129
130 ct_idft.dft(module, self);
131 }
132
133 pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
134 &mut self,
135 module: &Module<FFT64>,
136 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
137 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
138 scratch: &mut Scratch,
139 ) {
140 let cols_out: usize = rhs.rank_out() + 1;
141
142 // Space fr normalized VMP result outside of DFT domain
143 let (res_idft_data, scratch1) = scratch.tmp_vec_znx(module, cols_out, lhs.size());
144
145 let mut res_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
146 data: res_idft_data,
147 basek: lhs.basek,
148 k: lhs.k,
149 };
150
151 res_idft.keyswitch_from_fourier(module, lhs, rhs, scratch1);
152
153 (0..cols_out).for_each(|i| {
154 module.vec_znx_dft(&mut self.data, i, &res_idft.data, i);
155 });
156 }
157
158 pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
159 &mut self,
160 module: &Module<FFT64>,
161 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
162 scratch: &mut Scratch,
163 ) {
164 unsafe {
165 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
166 self.keyswitch(&module, &*self_ptr, rhs, scratch);
167 }
168 }
169
170 pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
171 &mut self,
172 module: &Module<FFT64>,
173 lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
174 rhs: &GGSWCiphertext<DataRhs, FFT64>,
175 scratch: &mut Scratch,
176 ) {
177 let basek: usize = self.basek();
178
179 #[cfg(debug_assertions)]
180 {
181 assert_eq!(rhs.rank(), lhs.rank());
182 assert_eq!(rhs.rank(), self.rank());
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 }
189
190 let cols: usize = rhs.rank() + 1;
191
192 // Space for VMP result in DFT domain and high precision
193 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size());
194
195 {
196 module.vmp_apply(&mut res_dft, &lhs.data, &rhs.data, scratch1);
197 }
198
199 // VMP result in high precision
200 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
201
202 // Space for VMP result normalized
203 let (mut res_small, scratch2) = scratch1.tmp_vec_znx(module, cols, rhs.size());
204 (0..cols).for_each(|i| {
205 module.vec_znx_big_normalize(basek, &mut res_small, i, &res_big, i, scratch2);
206 module.vec_znx_dft(&mut self.data, i, &res_small, i);
207 });
208 }
209
210 pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
211 &mut self,
212 module: &Module<FFT64>,
213 rhs: &GGSWCiphertext<DataRhs, FFT64>,
214 scratch: &mut Scratch,
215 ) {
216 unsafe {
217 let self_ptr: *mut GLWECiphertextFourier<DataSelf, FFT64> = self as *mut GLWECiphertextFourier<DataSelf, FFT64>;
218 self.external_product(&module, &*self_ptr, rhs, scratch);
219 }
220 }
221}
222
223impl<DataSelf: AsRef<[u8]>> GLWECiphertextFourier<DataSelf, FFT64> {
224 pub fn decrypt<DataPt: AsRef<[u8]> + AsMut<[u8]>, DataSk: AsRef<[u8]>>(
225 &self,
226 module: &Module<FFT64>,
227 pt: &mut GLWEPlaintext<DataPt>,
228 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
229 scratch: &mut Scratch,
230 ) {
231 #[cfg(debug_assertions)]
232 {
233 assert_eq!(self.rank(), sk_dft.rank());
234 assert_eq!(self.n(), module.n());
235 assert_eq!(pt.n(), module.n());
236 assert_eq!(sk_dft.n(), module.n());
237 }
238
239 let cols = self.rank() + 1;
240
241 let (mut pt_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
242 pt_big.zero();
243
244 {
245 (1..cols).for_each(|i| {
246 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
247 module.svp_apply(&mut ci_dft, 0, &sk_dft.data, i - 1, &self.data, i);
248 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
249 module.vec_znx_big_add_inplace(&mut pt_big, 0, &ci_big, 0);
250 });
251 }
252
253 {
254 let (mut c0_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, self.size());
255 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
256 module.vec_znx_idft(&mut c0_big, 0, &self.data, 0, scratch_2);
257 module.vec_znx_big_add_inplace(&mut pt_big, 0, &c0_big, 0);
258 }
259
260 // pt = norm(BIG(m + e))
261 module.vec_znx_big_normalize(self.basek(), &mut pt.data, 0, &mut pt_big, 0, scratch_1);
262
263 pt.basek = self.basek();
264 pt.k = pt.k().min(self.k());
265 }
266
267 #[allow(dead_code)]
268 pub(crate) fn idft<DataRes: AsRef<[u8]> + AsMut<[u8]>>(
269 &self,
270 module: &Module<FFT64>,
271 res: &mut GLWECiphertext<DataRes>,
272 scratch: &mut Scratch,
273 ) {
274 #[cfg(debug_assertions)]
275 {
276 assert_eq!(self.rank(), res.rank());
277 assert_eq!(self.basek(), res.basek())
278 }
279
280 let min_size: usize = self.size().min(res.size());
281
282 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 1, min_size);
283
284 (0..self.rank() + 1).for_each(|i| {
285 module.vec_znx_idft(&mut res_big, 0, &self.data, i, scratch1);
286 module.vec_znx_big_normalize(self.basek(), &mut res.data, i, &res_big, 0, scratch1);
287 });
288 }
289}