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 pub fn mul_rlwe_dft_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
86 (Self::mul_rlwe_scratch_space(module, res_size, a_size, grlwe_size) | module.vec_znx_idft_tmp_bytes())
87 + module.bytes_of_vec_znx(2, a_size)
88 + module.bytes_of_vec_znx(2, res_size)
89 }
90
91 pub fn mul_rlwe_dft_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, grlwe_size: usize) -> usize {
92 (Self::mul_rlwe_inplace_scratch_space(module, res_size, grlwe_size) | module.vec_znx_idft_tmp_bytes())
93 + module.bytes_of_vec_znx(2, res_size)
94 }
95
96 pub fn mul_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
97 Self::mul_rlwe_dft_inplace_scratch_space(module, res_size, grlwe_size) + module.bytes_of_vec_znx_dft(2, a_size)
98 }
99
100 pub fn mul_grlwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
101 Self::mul_rlwe_dft_inplace_scratch_space(module, res_size, grlwe_size) + module.bytes_of_vec_znx_dft(2, a_size)
102 }
103
104 pub fn mul_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
105 Self::mul_rlwe_dft_inplace_scratch_space(module, res_size, grlwe_size) + module.bytes_of_vec_znx_dft(2, a_size)
106 }
107
108 pub fn mul_rgsw_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, a_size: usize, grlwe_size: usize) -> usize {
109 Self::mul_rlwe_dft_inplace_scratch_space(module, res_size, grlwe_size) + module.bytes_of_vec_znx_dft(2, a_size)
110 }
111}
112
113impl<C> RGSWCt<C, FFT64>
114where
115 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
116{
117 pub fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut RLWECtDft<R, FFT64>)
118 where
119 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
120 {
121 module.vmp_extract_row(res, self, row_i, col_j);
122 }
123}
124
125impl<C> RGSWCt<C, FFT64>
126where
127 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
128{
129 pub fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &RLWECtDft<R, FFT64>)
130 where
131 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
132 {
133 module.vmp_prepare_row(self, row_i, col_j, a);
134 }
135}
136
137pub fn encrypt_rgsw_sk<C, P, S>(
138 module: &Module<FFT64>,
139 ct: &mut RGSWCt<C, FFT64>,
140 pt: &ScalarZnx<P>,
141 sk_dft: &SecretKeyDft<S, FFT64>,
142 source_xa: &mut Source,
143 source_xe: &mut Source,
144 sigma: f64,
145 bound: f64,
146 scratch: &mut Scratch,
147) where
148 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
149 ScalarZnx<P>: ScalarZnxToRef,
150 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
151{
152 let size: usize = ct.size();
153 let log_base2k: usize = ct.log_base2k();
154
155 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
156 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, 2, size);
157
158 let mut vec_znx_pt: RLWEPt<&mut [u8]> = RLWEPt {
159 data: tmp_znx_pt,
160 log_base2k: log_base2k,
161 log_k: ct.log_k(),
162 };
163
164 let mut vec_znx_ct: RLWECt<&mut [u8]> = RLWECt {
165 data: tmp_znx_ct,
166 log_base2k: log_base2k,
167 log_k: ct.log_k(),
168 };
169
170 (0..ct.rows()).for_each(|row_j| {
171 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
172 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
173 module.vec_znx_normalize_inplace(log_base2k, &mut vec_znx_pt, 0, scrach_2);
174
175 (0..ct.cols()).for_each(|col_i| {
176 // rlwe encrypt of vec_znx_pt into vec_znx_ct
177 encrypt_rlwe_sk(
178 module,
179 &mut vec_znx_ct,
180 Some((&vec_znx_pt, col_i)),
181 sk_dft,
182 source_xa,
183 source_xe,
184 sigma,
185 bound,
186 scrach_2,
187 );
188
189 // Switch vec_znx_ct into DFT domain
190 {
191 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, 2, size);
192 module.vec_znx_dft(&mut vec_znx_dft_ct, 0, &vec_znx_ct, 0);
193 module.vec_znx_dft(&mut vec_znx_dft_ct, 1, &vec_znx_ct, 1);
194 module.vmp_prepare_row(ct, row_j, col_i, &vec_znx_dft_ct);
195 }
196 });
197
198 vec_znx_pt.data.zero(); // zeroes for next iteration
199 });
200}
201
202impl<C> RGSWCt<C, FFT64> {
203 pub fn encrypt_sk<P, S>(
204 &mut self,
205 module: &Module<FFT64>,
206 pt: &ScalarZnx<P>,
207 sk_dft: &SecretKeyDft<S, FFT64>,
208 source_xa: &mut Source,
209 source_xe: &mut Source,
210 sigma: f64,
211 bound: f64,
212 scratch: &mut Scratch,
213 ) where
214 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
215 ScalarZnx<P>: ScalarZnxToRef,
216 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
217 {
218 encrypt_rgsw_sk(
219 module, self, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
220 )
221 }
222
223 pub fn mul_rlwe<R, A>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, a: &RLWECt<A>, scratch: &mut Scratch)
224 where
225 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
226 VecZnx<R>: VecZnxToMut,
227 VecZnx<A>: VecZnxToRef,
228 {
229 let log_base2k: usize = self.log_base2k();
230
231 #[cfg(debug_assertions)]
232 {
233 assert_eq!(res.log_base2k(), log_base2k);
234 assert_eq!(a.log_base2k(), log_base2k);
235 assert_eq!(self.n(), module.n());
236 assert_eq!(res.n(), module.n());
237 assert_eq!(a.n(), module.n());
238 }
239
240 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, 2, self.size()); // Todo optimise
241
242 {
243 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, 2, a.size());
244 module.vec_znx_dft(&mut a_dft, 0, a, 0);
245 module.vec_znx_dft(&mut a_dft, 1, a, 1);
246 module.vmp_apply(&mut res_dft, &a_dft, self, scratch2);
247 }
248
249 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
250
251 module.vec_znx_big_normalize(log_base2k, res, 0, &res_big, 0, scratch1);
252 module.vec_znx_big_normalize(log_base2k, res, 1, &res_big, 1, scratch1);
253 }
254
255 pub fn mul_rlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, scratch: &mut Scratch)
256 where
257 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
258 VecZnx<R>: VecZnxToMut + VecZnxToRef,
259 {
260 unsafe {
261 let res_ptr: *mut RLWECt<R> = res as *mut RLWECt<R>; // This is ok because [Self::mul_rlwe] only updates res at the end.
262 self.mul_rlwe(&module, &mut *res_ptr, &*res_ptr, scratch);
263 }
264 }
265
266 pub fn mul_rlwe_dft<R, A>(
267 &self,
268 module: &Module<FFT64>,
269 res: &mut RLWECtDft<R, FFT64>,
270 a: &RLWECtDft<A, FFT64>,
271 scratch: &mut Scratch,
272 ) where
273 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
274 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
275 VecZnxDft<A, FFT64>: VecZnxDftToRef<FFT64> + ZnxInfos,
276 {
277 let log_base2k: usize = self.log_base2k();
278
279 #[cfg(debug_assertions)]
280 {
281 assert_eq!(res.log_base2k(), log_base2k);
282 assert_eq!(self.n(), module.n());
283 assert_eq!(res.n(), module.n());
284 }
285
286 let (a_data, scratch_1) = scratch.tmp_vec_znx(module, 2, a.size());
287
288 let mut a_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
289 data: a_data,
290 log_base2k: a.log_base2k(),
291 log_k: a.log_k(),
292 };
293
294 a.idft(module, &mut a_idft, scratch_1);
295
296 let (res_data, scratch_2) = scratch_1.tmp_vec_znx(module, 2, res.size());
297
298 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
299 data: res_data,
300 log_base2k: res.log_base2k(),
301 log_k: res.log_k(),
302 };
303
304 self.mul_rlwe(module, &mut res_idft, &a_idft, scratch_2);
305
306 module.vec_znx_dft(res, 0, &res_idft, 0);
307 module.vec_znx_dft(res, 1, &res_idft, 1);
308 }
309
310 pub fn mul_rlwe_dft_inplace<R>(&self, module: &Module<FFT64>, res: &mut RLWECtDft<R, FFT64>, scratch: &mut Scratch)
311 where
312 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
313 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
314 {
315 let log_base2k: usize = self.log_base2k();
316
317 #[cfg(debug_assertions)]
318 {
319 assert_eq!(res.log_base2k(), log_base2k);
320 assert_eq!(self.n(), module.n());
321 assert_eq!(res.n(), module.n());
322 }
323
324 let (res_data, scratch_1) = scratch.tmp_vec_znx(module, 2, res.size());
325
326 let mut res_idft: RLWECt<&mut [u8]> = RLWECt::<&mut [u8]> {
327 data: res_data,
328 log_base2k: res.log_base2k(),
329 log_k: res.log_k(),
330 };
331
332 res.idft(module, &mut res_idft, scratch_1);
333
334 self.mul_rlwe_inplace(module, &mut res_idft, scratch_1);
335
336 module.vec_znx_dft(res, 0, &res_idft, 0);
337 module.vec_znx_dft(res, 1, &res_idft, 1);
338 }
339
340 pub fn mul_grlwe<R, A>(
341 &self,
342 module: &Module<FFT64>,
343 res: &mut GRLWECt<R, FFT64>,
344 a: &GRLWECt<A, FFT64>,
345 scratch: &mut Scratch,
346 ) where
347 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
348 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
349 MatZnxDft<A, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
350 {
351 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
352
353 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
354 data: tmp_row_data,
355 log_base2k: a.log_base2k(),
356 log_k: a.log_k(),
357 };
358
359 let min_rows: usize = res.rows().min(a.rows());
360
361 (0..min_rows).for_each(|row_i| {
362 a.get_row(module, row_i, &mut tmp_row);
363 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
364 res.set_row(module, row_i, &tmp_row);
365 });
366
367 tmp_row.data.zero();
368
369 (min_rows..res.rows()).for_each(|row_i| {
370 res.set_row(module, row_i, &tmp_row);
371 })
372 }
373
374 pub fn mul_grlwe_inplace<R>(&self, module: &Module<FFT64>, res: &mut GRLWECt<R, FFT64>, scratch: &mut Scratch)
375 where
376 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
377 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
378 {
379 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
380
381 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
382 data: tmp_row_data,
383 log_base2k: res.log_base2k(),
384 log_k: res.log_k(),
385 };
386
387 (0..res.rows()).for_each(|row_i| {
388 res.get_row(module, row_i, &mut tmp_row);
389 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
390 res.set_row(module, row_i, &tmp_row);
391 });
392 }
393
394 pub fn mul_rgsw<R, A>(&self, module: &Module<FFT64>, res: &mut RGSWCt<R, FFT64>, a: &RGSWCt<A, FFT64>, scratch: &mut Scratch)
395 where
396 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
397 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
398 MatZnxDft<A, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
399 {
400 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
401
402 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
403 data: tmp_row_data,
404 log_base2k: a.log_base2k(),
405 log_k: a.log_k(),
406 };
407
408 let min_rows: usize = res.rows().min(a.rows());
409
410 (0..min_rows).for_each(|row_i| {
411 a.get_row(module, row_i, 0, &mut tmp_row);
412 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
413 res.set_row(module, row_i, 0, &tmp_row);
414 });
415
416 (0..min_rows).for_each(|row_i| {
417 a.get_row(module, row_i, 1, &mut tmp_row);
418 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
419 res.set_row(module, row_i, 1, &tmp_row);
420 });
421
422 tmp_row.data.zero();
423
424 (min_rows..res.rows()).for_each(|row_i| {
425 res.set_row(module, row_i, 0, &tmp_row);
426 res.set_row(module, row_i, 1, &tmp_row);
427 })
428 }
429
430 pub fn mul_rgsw_inplace<R>(&self, module: &Module<FFT64>, res: &mut RGSWCt<R, FFT64>, scratch: &mut Scratch)
431 where
432 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64> + ZnxInfos,
433 MatZnxDft<R, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64> + ZnxInfos,
434 {
435 let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
436
437 let mut tmp_row: RLWECtDft<&mut [u8], FFT64> = RLWECtDft::<&mut [u8], FFT64> {
438 data: tmp_row_data,
439 log_base2k: res.log_base2k(),
440 log_k: res.log_k(),
441 };
442
443 (0..res.rows()).for_each(|row_i| {
444 res.get_row(module, row_i, 0, &mut tmp_row);
445 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
446 res.set_row(module, row_i, 0, &tmp_row);
447 });
448
449 (0..res.rows()).for_each(|row_i| {
450 res.get_row(module, row_i, 1, &mut tmp_row);
451 self.mul_rlwe_dft_inplace(module, &mut tmp_row, scratch1);
452 res.set_row(module, row_i, 1, &tmp_row);
453 });
454 }
455}