poulpy/rlwe/src/elem_rgsw.rs
rlwe/src/elem_rgsw.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftScratch, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx,
3 ScalarZnxDft, ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigOps, VecZnxBigScratch,
4 VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, VecZnxToMut, VecZnxToRef, ZnxInfos,
5 ZnxZero,
6};
7use sampling::source::Source;
8
9use crate::{
10 elem::Infos,
11 elem_grlwe::GRLWECt,
12 elem_rlwe::{RLWECt, RLWECtDft, RLWEPt, encrypt_rlwe_sk},
13 keys::SecretKeyDft,
14 utils::derive_size,
15};
16
17pub struct RGSWCt<C, B: Backend> {
18 pub data: MatZnxDft<C, B>,
19 pub log_base2k: usize,
20 pub log_k: usize,
21}
22
23impl<B: Backend> RGSWCt<Vec<u8>, B> {
24 pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
25 Self {
26 data: module.new_mat_znx_dft(rows, 2, 2, derive_size(log_base2k, log_k)),
27 log_base2k: log_base2k,
28 log_k: log_k,
29 }
30 }
31}
32
33impl<T, B: Backend> Infos for RGSWCt<T, B> {
34 type Inner = MatZnxDft<T, B>;
35
36 fn inner(&self) -> &Self::Inner {
37 &self.data
38 }
39
40 fn log_base2k(&self) -> usize {
41 self.log_base2k
42 }
43
44 fn log_k(&self) -> usize {
45 self.log_k
46 }
47}
48
49impl<C, B: Backend> MatZnxDftToMut<B> for RGSWCt<C, B>
50where
51 MatZnxDft<C, B>: MatZnxDftToMut<B>,
52{
53 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
54 self.data.to_mut()
55 }
56}
57
58impl<C, B: Backend> MatZnxDftToRef<B> for RGSWCt<C, B>
59where
60 MatZnxDft<C, B>: MatZnxDftToRef<B>,
61{
62 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
63 self.data.to_ref()
64 }
65}
66
67impl RGSWCt<Vec<u8>, FFT64> {
68 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
69 RLWECt::encrypt_sk_scratch_space(module, size)
70 + module.bytes_of_vec_znx(2, size)
71 + module.bytes_of_vec_znx(1, size)
72 + module.bytes_of_vec_znx_dft(2, size)
73 }
74
75 pub fn mul_rlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, rgsw_size: usize) -> usize {
76 module.bytes_of_vec_znx_dft(2, rgsw_size)
77 + ((module.bytes_of_vec_znx_dft(2, a_size) + module.vmp_apply_tmp_bytes(res_size, a_size, a_size, 2, 2, rgsw_size))
78 | module.vec_znx_big_normalize_tmp_bytes())
79 }
80
81 pub fn mul_rlwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rgsw_size: usize) -> usize {
82 Self::mul_rlwe_scratch_space(module, res_size, res_size, rgsw_size)
83 }
84}
85
86impl<C> RGSWCt<C, FFT64>
87where
88 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
89{
90 pub fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut RLWECtDft<R, FFT64>)
91 where
92 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
93 {
94 module.vmp_extract_row(res, self, row_i, col_j);
95 }
96}
97
98impl<C> RGSWCt<C, FFT64>
99where
100 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
101{
102 pub fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &RLWECtDft<R, FFT64>)
103 where
104 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
105 {
106 module.vmp_prepare_row(self, row_i, col_j, a);
107 }
108}
109
110pub fn encrypt_rgsw_sk<C, P, S>(
111 module: &Module<FFT64>,
112 ct: &mut RGSWCt<C, FFT64>,
113 pt: &ScalarZnx<P>,
114 sk_dft: &SecretKeyDft<S, FFT64>,
115 source_xa: &mut Source,
116 source_xe: &mut Source,
117 sigma: f64,
118 bound: f64,
119 scratch: &mut Scratch,
120) where
121 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
122 ScalarZnx<P>: ScalarZnxToRef,
123 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
124{
125 let size: usize = ct.size();
126 let log_base2k: usize = ct.log_base2k();
127
128 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
129 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, 2, size);
130
131 let mut vec_znx_pt: RLWEPt<&mut [u8]> = RLWEPt {
132 data: tmp_znx_pt,
133 log_base2k: log_base2k,
134 log_k: ct.log_k(),
135 };
136
137 let mut vec_znx_ct: RLWECt<&mut [u8]> = RLWECt {
138 data: tmp_znx_ct,
139 log_base2k: log_base2k,
140 log_k: ct.log_k(),
141 };
142
143 (0..ct.rows()).for_each(|row_j| {
144 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
145 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
146 module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scrach_2);
147
148 (0..ct.cols()).for_each(|col_i| {
149 // rlwe encrypt of vec_znx_pt into vec_znx_ct
150 encrypt_rlwe_sk(
151 module,
152 &mut vec_znx_ct,
153 Some((&vec_znx_pt, col_i)),
154 sk_dft,
155 source_xa,
156 source_xe,
157 sigma,
158 bound,
159 scrach_2,
160 );
161
162 // Switch vec_znx_ct into DFT domain
163 {
164 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, 2, size);
165 module.vec_znx_dft(&mut vec_znx_dft_ct, 0, &vec_znx_ct, 0);
166 module.vec_znx_dft(&mut vec_znx_dft_ct, 1, &vec_znx_ct, 1);
167 module.vmp_prepare_row(ct, row_j, col_i, &vec_znx_dft_ct);
168 }
169 });
170
171 vec_znx_pt.data.zero(); // zeroes for next iteration
172 });
173}
174
175impl<C> RGSWCt<C, FFT64> {
176 pub fn encrypt_sk<P, S>(
177 &mut self,
178 module: &Module<FFT64>,
179 pt: &ScalarZnx<P>,
180 sk_dft: &SecretKeyDft<S, FFT64>,
181 source_xa: &mut Source,
182 source_xe: &mut Source,
183 sigma: f64,
184 bound: f64,
185 scratch: &mut Scratch,
186 ) where
187 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
188 ScalarZnx<P>: ScalarZnxToRef,
189 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
190 {
191 encrypt_rgsw_sk(
192 module, self, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
193 )
194 }
195
196 pub fn mul_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, a: &RLWECt<A>, scratch: &mut Scratch)
197 where
198 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
199 VecZnx<R>: VecZnxToMut,
200 VecZnx<A>: VecZnxToRef,
201 {
202 let log_base2k: usize = self.log_base2k();
203
204 #[cfg(debug_assertions)]
205 {
206 assert_eq!(res.log_base2k(), log_base2k);
207 assert_eq!(a.log_base2k(), log_base2k);
208 assert_eq!(self.n(), module.n());
209 assert_eq!(res.n(), module.n());
210 assert_eq!(a.n(), module.n());
211 }
212
213 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, 2, self.size()); // Todo optimise
214
215 {
216 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, 2, a.size());
217 module.vec_znx_dft(&mut a_dft, 0, a, 0);
218 module.vec_znx_dft(&mut a_dft, 1, a, 1);
219 module.vmp_apply(&mut res_dft, &a_dft, self, scratch2);
220 }
221
222 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
223
224 module.vec_znx_big_normalize(log_base2k, res, 0, &res_big, 0, scratch1);
225 module.vec_znx_big_normalize(log_base2k, res, 1, &res_big, 1, scratch1);
226 }
227
228 pub fn mul_rlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, scratch: &mut Scratch)
229 where
230 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
231 VecZnx<R>: VecZnxToMut + VecZnxToRef,
232 {
233 unsafe {
234 let res_ptr: *mut RLWECt<R> = res as *mut RLWECt<R>; // This is ok because [Self::mul_rlwe] only updates res at the end.
235 self.mul_rlwe(&module, &mut *res_ptr, &*res_ptr, scratch);
236 }
237 }
238
239 pub fn mul_rlwe_dft<R, A>(
240 &self,
241 module: &Module<FFT64>,
242 res: &mut RLWECtDft<R, FFT64>,
243 a: &RLWECtDft<A, FFT64>,
244 scratch: &mut Scratch,
245 ) where
246 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
247 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
248 VecZnxDft<A, FFT64>: VecZnxDftToRef<FFT64> + ZnxInfos,
249 {
250 let log_base2k: usize = self.log_base2k();
251
252 #[cfg(debug_assertions)]
253 {
254 assert_eq!(res.log_base2k(), log_base2k);
255 assert_eq!(self.n(), module.n());
256 assert_eq!(res.n(), module.n());
257 }
258
259 let (a_data, scratch_1) = scratch.tmp_vec_znx(module, 2, a.size());
260
261 let mut a_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
262 data: a_data,
263 log_base2k: a.log_base2k(),
264 log_k: a.log_k(),
265 };
266
267 a.idft(module, &mut a_idft, scratch_1);
268
269 let (res_data, scratch_2) = scratch_1.tmp_vec_znx(module, 2, res.size());
270
271 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
272 data: res_data,
273 log_base2k: res.log_base2k(),
274 log_k: res.log_k(),
275 };
276
277 self.mul_rlwe(module, &mut res_idft, &a_idft, scratch_2);
278
279 module.vec_znx_dft(res, 0, &res_idft, 0);
280 module.vec_znx_dft(res, 1, &res_idft, 1);
281 }
282
283 pub fn mul_rlwe_dft_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECtDft<R, FFT64>, scratch: &mut Scratch)
284 where
285 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
286 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
287 {
288 let log_base2k: usize = self.log_base2k();
289
290 #[cfg(debug_assertions)]
291 {
292 assert_eq!(res.log_base2k(), log_base2k);
293 assert_eq!(self.n(), module.n());
294 assert_eq!(res.n(), module.n());
295 }
296
297 let (res_data, scratch_1) = scratch.tmp_vec_znx(module, 2, res.size());
298
299 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
300 data: res_data,
301 log_base2k: res.log_base2k(),
302 log_k: res.log_k(),
303 };
304
305 res.idft(module, &mut res_idft, scratch_1);
306
307 self.mul_rlwe_inplace(module, &mut res_idft, scratch_1);
308
309 module.vec_znx_dft(res, 0, &res_idft, 0);
310 module.vec_znx_dft(res, 1, &res_idft, 1);
311 }
312
313 pub fn mul_grlwe<R, A>(
314 &self,
315 module: &Module<FFT64>,
316 res: &mut GRLWECt<R, FFT64>,
317 a: &GRLWECt<A, FFT64>,
318 scratch: &mut Scratch,
319 ) where
320 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
321 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
322 MatZnxDft<A, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
323 {
324 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
325
326 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
327 data: tmp_row_data,
328 log_base2k: a.log_base2k(),
329 log_k: a.log_k(),
330 };
331
332 let min_rows: usize = res.rows().min(a.rows());
333
334 (0..min_rows).for_each(|row_i| {
335 a.get_row(module, row_i, &mut tmp_row);
336 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
337 res.set_row(module, row_i, &tmp_row);
338 });
339
340 tmp_row.data.zero();
341
342 (min_rows..res.rows()).for_each(|row_i| {
343 res.set_row(module, row_i, &tmp_row);
344 })
345 }
346
347 pub fn mul_grlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut GRLWECt<R, FFT64>, scratch: &mut Scratch)
348 where
349 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
350 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
351 {
352 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
353
354 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
355 data: tmp_row_data,
356 log_base2k: res.log_base2k(),
357 log_k: res.log_k(),
358 };
359
360 (0..res.rows()).for_each(|row_i| {
361 res.get_row(module, row_i, &mut tmp_row);
362 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
363 res.set_row(module, row_i, &tmp_row);
364 });
365 }
366
367 pub fn mul_rgsw<R, A>(&self, module: &Module<FFT64>, res: &mut RGSWCt<R, FFT64>, a: &RGSWCt<A, FFT64>, scratch: &mut Scratch)
368 where
369 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
370 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
371 MatZnxDft<A, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
372 {
373 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
374
375 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
376 data: tmp_row_data,
377 log_base2k: a.log_base2k(),
378 log_k: a.log_k(),
379 };
380
381 let min_rows: usize = res.rows().min(a.rows());
382
383 (0..min_rows).for_each(|row_i| {
384 a.get_row(module, row_i, 0, &mut tmp_row);
385 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
386 res.set_row(module, row_i, 0, &tmp_row);
387 });
388
389 (0..min_rows).for_each(|row_i| {
390 a.get_row(module, row_i, 1, &mut tmp_row);
391 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
392 res.set_row(module, row_i, 1, &tmp_row);
393 });
394
395 tmp_row.data.zero();
396
397 (min_rows..res.rows()).for_each(|row_i| {
398 res.set_row(module, row_i, 0, &tmp_row);
399 res.set_row(module, row_i, 1, &tmp_row);
400 })
401 }
402
403 pub fn mul_rgsw_inplace<R>(&self, module: &Module<FFT64>, res: &mut RGSWCt<R, FFT64>, scratch: &mut Scratch)
404 where
405 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
406 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
407 {
408 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
409
410 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
411 data: tmp_row_data,
412 log_base2k: res.log_base2k(),
413 log_k: res.log_k(),
414 };
415
416 (0..res.rows()).for_each(|row_i| {
417 res.get_row(module, row_i, 0, &mut tmp_row);
418 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
419 res.set_row(module, row_i, 0, &tmp_row);
420 });
421
422 (0..res.rows()).for_each(|row_i| {
423 res.get_row(module, row_i, 1, &mut tmp_row);
424 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
425 res.set_row(module, row_i, 1, &tmp_row);
426 });
427 }
428}