poulpy/poulpy-core/src/operations/glwe.rs
poulpy-core/src/operations/glwe.rsRaw
1use poulpy_hal::{
2 api::{
3 CnvPVecBytesOf, Convolution, ModuleN, ScratchTakeBasic, VecZnxAdd, VecZnxAddInplace, VecZnxBigAddSmallInplace,
4 VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftApply, VecZnxDftBytesOf,
5 VecZnxIdftApplyConsume, VecZnxMulXpMinusOne, VecZnxMulXpMinusOneInplace, VecZnxNegate, VecZnxNormalize,
6 VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotate, VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub,
7 VecZnxSubInplace, VecZnxSubNegateInplace, VecZnxZero,
8 },
9 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, VecZnxBig},
10 reference::vec_znx::vec_znx_rotate_inplace_tmp_bytes,
11};
12
13use crate::{
14 GGLWEProduct, ScratchTakeCore,
15 layouts::{
16 Base2K, GGLWEInfos, GLWE, GLWEInfos, GLWEPlaintext, GLWETensor, GLWETensorKeyPrepared, GLWEToMut, GLWEToRef, LWEInfos,
17 TorusPrecision,
18 },
19};
20
21pub trait GLWEMulConst<BE: Backend> {
22 fn glwe_mul_const_tmp_bytes<R, A>(&self, res: &R, res_offset: usize, a: &A, b_size: usize) -> usize
23 where
24 R: GLWEInfos,
25 A: GLWEInfos;
26
27 fn glwe_mul_const<R, A>(&self, res: &mut GLWE<R>, res_offset: usize, a: &GLWE<A>, b: &[i64], scratch: &mut Scratch<BE>)
28 where
29 R: DataMut,
30 A: DataRef;
31
32 fn glwe_mul_const_inplace<R>(&self, res: &mut GLWE<R>, res_offset: usize, b: &[i64], scratch: &mut Scratch<BE>)
33 where
34 R: DataMut;
35}
36
37impl<BE: Backend> GLWEMulConst<BE> for Module<BE>
38where
39 Self: Convolution<BE> + VecZnxBigBytesOf + VecZnxBigNormalize<BE> + VecZnxBigNormalizeTmpBytes,
40 Scratch<BE>: ScratchTakeCore<BE>,
41{
42 fn glwe_mul_const_tmp_bytes<R, A>(&self, res: &R, res_offset: usize, a: &A, b_size: usize) -> usize
43 where
44 R: GLWEInfos,
45 A: GLWEInfos,
46 {
47 let a_base2k: usize = a.base2k().as_usize();
48 let res_base2k: usize = res.base2k().as_usize();
49
50 let res_size: usize = (res.size() * res_base2k).div_ceil(a_base2k);
51 let res_big: usize = self.bytes_of_vec_znx_big(1, res_size);
52 let cnv: usize = self.cnv_by_const_apply_tmp_bytes(res_size, res_offset, a.size(), b_size);
53 let normalize: usize = self.vec_znx_big_normalize_tmp_bytes();
54
55 res_big + cnv.max(normalize)
56 }
57
58 fn glwe_mul_const<R, A>(&self, res: &mut GLWE<R>, res_offset: usize, a: &GLWE<A>, b: &[i64], scratch: &mut Scratch<BE>)
59 where
60 R: DataMut,
61 A: DataRef,
62 {
63 assert_eq!(res.rank(), a.rank());
64
65 let cols: usize = res.rank().as_usize() + 1;
66 let a_base2k: usize = a.base2k().as_usize();
67 let res_base2k: usize = res.base2k().as_usize();
68
69 let (res_offset_hi, res_offset_lo) = if res_offset < a_base2k {
70 (0, -((a_base2k - (res_offset % a_base2k)) as i64))
71 } else {
72 ((res_offset / a_base2k).saturating_sub(1), (res_offset % a_base2k) as i64)
73 };
74
75 let res_dft_size = res
76 .k()
77 .as_usize()
78 .div_ceil(a.base2k().as_usize())
79 .min(a.size() + b.len() - res_offset_hi);
80
81 let (mut res_big, scratch_1) = scratch.take_vec_znx_big(self, 1, res_dft_size);
82 for i in 0..cols {
83 self.cnv_by_const_apply(&mut res_big, res_offset_hi, 0, a.data(), i, b, scratch_1);
84 self.vec_znx_big_normalize(res.data_mut(), res_base2k, res_offset_lo, i, &res_big, a_base2k, 0, scratch_1);
85 }
86 }
87
88 fn glwe_mul_const_inplace<R>(&self, res: &mut GLWE<R>, res_offset: usize, b: &[i64], scratch: &mut Scratch<BE>)
89 where
90 R: DataMut,
91 {
92 let cols: usize = res.rank().as_usize() + 1;
93 let res_base2k: usize = res.base2k().as_usize();
94
95 let (res_offset_hi, res_offset_lo) = if res_offset < res_base2k {
96 (0, -((res_base2k - (res_offset % res_base2k)) as i64))
97 } else {
98 ((res_offset / res_base2k).saturating_sub(1), (res_offset % res_base2k) as i64)
99 };
100
101 let (mut res_big, scratch_1) = scratch.take_vec_znx_big(self, 1, res.size());
102 for i in 0..cols {
103 self.cnv_by_const_apply(&mut res_big, res_offset_hi, 0, res.data(), i, b, scratch_1);
104 self.vec_znx_big_normalize(
105 res.data_mut(),
106 res_base2k,
107 res_offset_lo,
108 i,
109 &res_big,
110 res_base2k,
111 0,
112 scratch_1,
113 );
114 }
115 }
116}
117
118impl<BE: Backend> GLWEMulPlain<BE> for Module<BE>
119where
120 Self: Sized
121 + ModuleN
122 + CnvPVecBytesOf
123 + VecZnxDftBytesOf
124 + VecZnxIdftApplyConsume<BE>
125 + VecZnxBigNormalize<BE>
126 + Convolution<BE>
127 + VecZnxBigNormalizeTmpBytes,
128 Scratch<BE>: ScratchTakeCore<BE>,
129{
130 fn glwe_mul_plain_tmp_bytes<R, A, B>(&self, res: &R, res_offset: usize, a: &A, b: &B) -> usize
131 where
132 R: GLWEInfos,
133 A: GLWEInfos,
134 B: GLWEInfos,
135 {
136 let ab_base2k: Base2K = a.base2k();
137 assert_eq!(b.base2k(), ab_base2k);
138
139 let cols: usize = res.rank().as_usize() + 1;
140
141 let a_size: usize = a.size();
142 let b_size: usize = b.size();
143 let res_size: usize = res.size();
144
145 let cnv_pvec: usize = self.bytes_of_cnv_pvec_left(cols, a_size) + self.bytes_of_cnv_pvec_right(1, b_size);
146 let cnv_prep: usize = self
147 .cnv_prepare_left_tmp_bytes(a_size, a_size)
148 .max(self.cnv_prepare_right_tmp_bytes(a_size, a_size));
149 let cnv_apply: usize = self.cnv_apply_dft_tmp_bytes(res_size, res_offset, a_size, b_size);
150
151 let res_dft_size = res
152 .k()
153 .as_usize()
154 .div_ceil(a.base2k().as_usize())
155 .min(a_size + b_size - res_offset / ab_base2k.as_usize());
156
157 let res_dft: usize = self.bytes_of_vec_znx_dft(1, res_dft_size);
158 let norm: usize = self.vec_znx_big_normalize_tmp_bytes();
159
160 cnv_pvec + cnv_prep + res_dft + cnv_apply.max(norm)
161 }
162
163 fn glwe_mul_plain<R, A, B>(
164 &self,
165 res: &mut GLWE<R>,
166 res_offset: usize,
167 a: &GLWE<A>,
168 b: &GLWEPlaintext<B>,
169 scratch: &mut Scratch<BE>,
170 ) where
171 R: DataMut,
172 A: DataRef,
173 B: DataRef,
174 {
175 assert_eq!(res.rank(), a.rank());
176
177 let a_base2k: usize = a.base2k().as_usize();
178 assert_eq!(b.base2k().as_usize(), a_base2k);
179 let res_base2k: usize = res.base2k().as_usize();
180
181 let cols: usize = res.rank().as_usize() + 1;
182
183 let (mut a_prep, scratch_1) = scratch.take_cnv_pvec_left(self, cols, a.size());
184 let (mut b_prep, scratch_2) = scratch_1.take_cnv_pvec_right(self, 1, b.size());
185
186 self.cnv_prepare_left(&mut a_prep, a.data(), scratch_2);
187 self.cnv_prepare_right(&mut b_prep, b.data(), scratch_2);
188
189 let (res_offset_hi, res_offset_lo) = if res_offset < a_base2k {
190 (0, -((a_base2k - (res_offset % a_base2k)) as i64))
191 } else {
192 ((res_offset / a_base2k).saturating_sub(1), (res_offset % a_base2k) as i64)
193 };
194
195 let res_dft_size = res
196 .k()
197 .as_usize()
198 .div_ceil(a.base2k().as_usize())
199 .min(a.size() + b.size() - res_offset_hi);
200
201 for i in 0..cols {
202 let (mut res_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, res_dft_size);
203 self.cnv_apply_dft(&mut res_dft, res_offset_hi, 0, &a_prep, i, &b_prep, 0, scratch_3);
204 let res_big = self.vec_znx_idft_apply_consume(res_dft);
205
206 self.vec_znx_big_normalize(res.data_mut(), res_base2k, res_offset_lo, i, &res_big, a_base2k, 0, scratch_3);
207 }
208 }
209
210 fn glwe_mul_plain_inplace<R, A>(&self, res: &mut GLWE<R>, res_offset: usize, a: &GLWEPlaintext<A>, scratch: &mut Scratch<BE>)
211 where
212 R: DataMut,
213 A: DataRef,
214 {
215 assert_eq!(res.rank(), a.rank());
216
217 let a_base2k: usize = a.base2k().as_usize();
218 let res_base2k: usize = res.base2k().as_usize();
219 assert_eq!(res_base2k, a_base2k);
220
221 let cols: usize = res.rank().as_usize() + 1;
222
223 let (mut res_prep, scratch_1) = scratch.take_cnv_pvec_left(self, cols, res.size());
224 let (mut a_prep, scratch_2) = scratch_1.take_cnv_pvec_right(self, 1, a.size());
225
226 self.cnv_prepare_left(&mut res_prep, res.data(), scratch_2);
227 self.cnv_prepare_right(&mut a_prep, a.data(), scratch_2);
228
229 let (res_offset_hi, res_offset_lo) = if res_offset < a_base2k {
230 (0, -((a_base2k - (res_offset % a_base2k)) as i64))
231 } else {
232 ((res_offset / a_base2k).saturating_sub(1), (res_offset % a_base2k) as i64)
233 };
234
235 let res_dft_size = res
236 .k()
237 .as_usize()
238 .div_ceil(a.base2k().as_usize())
239 .min(a.size() + res.size() - res_offset_hi);
240
241 for i in 0..cols {
242 let (mut res_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, res_dft_size);
243 self.cnv_apply_dft(&mut res_dft, res_offset, 0, &res_prep, i, &a_prep, 0, scratch_3);
244 let res_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(res_dft);
245 self.vec_znx_big_normalize(res.data_mut(), res_base2k, res_offset_lo, i, &res_big, a_base2k, 0, scratch_3);
246 }
247 }
248}
249
250pub trait GLWEMulPlain<BE: Backend> {
251 fn glwe_mul_plain_tmp_bytes<R, A, B>(&self, res: &R, res_offset: usize, a: &A, b: &B) -> usize
252 where
253 R: GLWEInfos,
254 A: GLWEInfos,
255 B: GLWEInfos;
256
257 fn glwe_mul_plain<R, A, B>(
258 &self,
259 res: &mut GLWE<R>,
260 res_offset: usize,
261 a: &GLWE<A>,
262 b: &GLWEPlaintext<B>,
263 scratch: &mut Scratch<BE>,
264 ) where
265 R: DataMut,
266 A: DataRef,
267 B: DataRef;
268
269 fn glwe_mul_plain_inplace<R, A>(&self, res: &mut GLWE<R>, res_offset: usize, a: &GLWEPlaintext<A>, scratch: &mut Scratch<BE>)
270 where
271 R: DataMut,
272 A: DataRef;
273}
274
275pub trait GLWETensoring<BE: Backend> {
276 fn glwe_tensor_apply_tmp_bytes<R, A, B>(&self, res: &R, res_offset: usize, a: &A, b: &B) -> usize
277 where
278 R: GLWEInfos,
279 A: GLWEInfos,
280 B: GLWEInfos;
281
282 /// res = (a (x) b) * 2^{res_offset * a_base2k}
283 ///
284 /// # Requires
285 /// * a.base2k() == b.base2k()
286 /// * a.rank() == b.rank()
287 ///
288 /// # Behavior
289 /// * res precision is truncated to res.max_k().min(a.max_k() + b.max_k() + k * a_base2k)
290 fn glwe_tensor_apply<R, A, B>(
291 &self,
292 res: &mut GLWETensor<R>,
293 res_offset: usize,
294 a: &GLWE<A>,
295 b: &GLWE<B>,
296 scratch: &mut Scratch<BE>,
297 ) where
298 R: DataMut,
299 A: DataRef,
300 B: DataRef;
301
302 fn glwe_tensor_relinearize<R, A, B>(
303 &self,
304 res: &mut GLWE<R>,
305 a: &GLWETensor<A>,
306 tsk: &GLWETensorKeyPrepared<B, BE>,
307 tsk_size: usize,
308 scratch: &mut Scratch<BE>,
309 ) where
310 R: DataMut,
311 A: DataRef,
312 B: DataRef;
313
314 fn glwe_tensor_relinearize_tmp_bytes<R, A, B>(&self, res: &R, a: &A, tsk: &B) -> usize
315 where
316 R: GLWEInfos,
317 A: GLWEInfos,
318 B: GGLWEInfos;
319}
320
321impl<BE: Backend> GLWETensoring<BE> for Module<BE>
322where
323 Self: Sized
324 + ModuleN
325 + CnvPVecBytesOf
326 + VecZnxDftBytesOf
327 + VecZnxIdftApplyConsume<BE>
328 + VecZnxBigNormalize<BE>
329 + Convolution<BE>
330 + VecZnxSubInplace
331 + VecZnxNegate
332 + VecZnxAddInplace
333 + VecZnxBigNormalizeTmpBytes
334 + VecZnxCopy
335 + VecZnxNormalize<BE>
336 + VecZnxDftApply<BE>
337 + GGLWEProduct<BE>
338 + VecZnxBigAddSmallInplace<BE>
339 + VecZnxNormalizeTmpBytes,
340 Scratch<BE>: ScratchTakeCore<BE>,
341{
342 fn glwe_tensor_apply_tmp_bytes<R, A, B>(&self, res: &R, res_offset: usize, a: &A, b: &B) -> usize
343 where
344 R: GLWEInfos,
345 A: GLWEInfos,
346 B: GLWEInfos,
347 {
348 let ab_base2k: Base2K = a.base2k();
349 assert_eq!(b.base2k(), ab_base2k);
350
351 let cols: usize = res.rank().as_usize() + 1;
352
353 let a_size: usize = a.size();
354 let b_size: usize = b.size();
355 let res_size: usize = res.size();
356
357 let cnv_pvec: usize = self.bytes_of_cnv_pvec_left(cols, a_size) + self.bytes_of_cnv_pvec_right(cols, b_size);
358 let cnv_prep: usize = self
359 .cnv_prepare_left_tmp_bytes(a_size, a_size)
360 .max(self.cnv_prepare_right_tmp_bytes(a_size, a_size));
361 let cnv_apply: usize = self
362 .cnv_apply_dft_tmp_bytes(res_size, res_offset, a_size, b_size)
363 .max(self.cnv_pairwise_apply_dft_tmp_bytes(res_size, res_offset, a_size, b_size));
364
365 let res_dft_size = res
366 .k()
367 .as_usize()
368 .div_ceil(a.base2k().as_usize())
369 .min(a_size + b_size - res_offset / ab_base2k.as_usize());
370
371 let res_dft: usize = self.bytes_of_vec_znx_dft(1, res_dft_size);
372 let tmp: usize = VecZnx::bytes_of(self.n(), 1, res.size());
373 let norm: usize = self.vec_znx_big_normalize_tmp_bytes();
374
375 cnv_pvec + cnv_prep + res_dft + cnv_apply.max(tmp + norm)
376 }
377
378 fn glwe_tensor_relinearize_tmp_bytes<R, A, B>(&self, res: &R, a: &A, tsk: &B) -> usize
379 where
380 R: GLWEInfos,
381 A: GLWEInfos,
382 B: GGLWEInfos,
383 {
384 let a_base2k: usize = a.base2k().into();
385 let key_base2k: usize = tsk.base2k().into();
386 let res_base2k: usize = res.base2k().into();
387
388 let cols: usize = tsk.rank_out().as_usize() + 1;
389 let pairs: usize = tsk.rank_in().as_usize();
390
391 let a_dft_size: usize = (a.size() * a_base2k).div_ceil(key_base2k);
392
393 let a_dft = self.bytes_of_vec_znx_dft(pairs, a_dft_size);
394
395 let a_conv: usize = if a_base2k != key_base2k || res_base2k != key_base2k {
396 VecZnx::bytes_of(self.n(), 1, a_dft_size) + self.vec_znx_normalize_tmp_bytes()
397 } else {
398 0
399 };
400
401 let res_dft: usize = self.bytes_of_vec_znx_dft(cols, tsk.size());
402
403 let gglwe_product: usize = self.gglwe_product_dft_tmp_bytes(res.size(), a_dft_size, tsk);
404
405 let big_normalize: usize = self.vec_znx_big_normalize_tmp_bytes();
406
407 a_dft.max(a_conv + big_normalize) + res_dft + gglwe_product.max(a_conv).max(big_normalize)
408 }
409
410 fn glwe_tensor_relinearize<R, A, B>(
411 &self,
412 res: &mut GLWE<R>,
413 a: &GLWETensor<A>,
414 tsk: &GLWETensorKeyPrepared<B, BE>,
415 tsk_size: usize,
416 scratch: &mut Scratch<BE>,
417 ) where
418 R: DataMut,
419 A: DataRef,
420 B: DataRef,
421 {
422 let a_base2k: usize = a.base2k().into();
423 let key_base2k: usize = tsk.base2k().into();
424 let res_base2k: usize = res.base2k().into();
425
426 assert_eq!(res.rank(), tsk.rank_out());
427 assert_eq!(a.rank(), tsk.rank_out());
428
429 let cols: usize = tsk.rank_out().as_usize() + 1;
430 let pairs: usize = tsk.rank_in().as_usize();
431
432 let a_dft_size: usize = (a.size() * a_base2k).div_ceil(key_base2k);
433
434 let (mut a_dft, scratch_1) = scratch.take_vec_znx_dft(self, pairs, a_dft_size);
435
436 if a_base2k != key_base2k {
437 let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(self.n(), 1, a_dft_size);
438 for i in 0..pairs {
439 self.vec_znx_normalize(&mut a_conv, key_base2k, 0, 0, a.data(), a_base2k, cols + i, scratch_2);
440 self.vec_znx_dft_apply(1, 0, &mut a_dft, i, &a_conv, 0);
441 }
442 } else {
443 for i in 0..pairs {
444 self.vec_znx_dft_apply(1, 0, &mut a_dft, i, a.data(), 0);
445 }
446 }
447
448 let (mut res_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, cols, tsk_size); // Todo optimise
449
450 self.gglwe_product_dft(&mut res_dft, &a_dft, &tsk.0, scratch_2);
451 let mut res_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(res_dft);
452
453 if res_base2k == key_base2k {
454 for i in 0..cols {
455 self.vec_znx_big_add_small_inplace(&mut res_big, i, a.data(), i);
456 }
457 } else {
458 let (mut a_conv, scratch_3) = scratch_2.take_vec_znx(self.n(), 1, a_dft_size);
459 for i in 0..cols {
460 self.vec_znx_normalize(&mut a_conv, key_base2k, 0, 0, a.data(), a_base2k, i, scratch_3);
461 self.vec_znx_big_add_small_inplace(&mut res_big, i, &a_conv, 0);
462 }
463 }
464
465 for i in 0..(res.rank() + 1).into() {
466 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
467 }
468 }
469
470 fn glwe_tensor_apply<R, A, B>(
471 &self,
472 res: &mut GLWETensor<R>,
473 res_offset: usize,
474 a: &GLWE<A>,
475 b: &GLWE<B>,
476 scratch: &mut Scratch<BE>,
477 ) where
478 R: DataMut,
479 A: DataRef,
480 B: DataRef,
481 {
482 let a_base2k: usize = a.base2k().as_usize();
483 assert_eq!(b.base2k().as_usize(), a_base2k);
484 let res_base2k: usize = res.base2k().as_usize();
485
486 let cols: usize = res.rank().as_usize() + 1;
487
488 let (mut a_prep, scratch_1) = scratch.take_cnv_pvec_left(self, cols, a.size());
489 let (mut b_prep, scratch_2) = scratch_1.take_cnv_pvec_right(self, cols, b.size());
490
491 self.cnv_prepare_left(&mut a_prep, a.data(), scratch_2);
492 self.cnv_prepare_right(&mut b_prep, b.data(), scratch_2);
493
494 // Example for rank=3
495 //
496 // (a0, a1, a2, a3) x (b0, b1, b2, a3)
497 // L L L L R R R R
498 //
499 // c(1) = a0 * b0 <- (L(a0) * R(b0))
500 // c(s1) = a0 * b1 + a1 * b0 <- (L(a0) + L(a1)) * (R(b0) + R(b1)) + NEG(L(a0) * R(b0)) + SUB(L(a1) * R(b1))
501 // c(s2) = a0 * b2 + a2 * b0 <- (L(a0) + L(a2)) * (R(b0) + R(b2)) + NEG(L(a0) * R(b0)) + SUB(L(a2) * R(b2))
502 // c(s3) = a0 * b3 + a3 * b0 <- (L(a0) + L(a3)) * (R(b0) + R(b3)) + NEG(L(a0) * R(b0)) + SUB(L(a3) * R(b3))
503 // c(s1^2) = a1 * b1 <- (L(a1) * R(b1))
504 // c(s1s2) = a1 * b2 + b2 * a1 <- (L(a1) + L(a2)) * (R(b1) + R(b2)) + NEG(L(a1) * R(b1)) + SUB(L(a2) * R(b2))
505 // c(s1s3) = a1 * b3 + b3 * a1 <- (L(a1) + L(a3)) * (R(b1) + R(b3)) + NEG(L(a1) * R(b1)) + SUB(L(a3) * R(b3))
506 // c(s2^2) = a2 * b2 <- (L(a2) * R(b2))
507 // c(s2s3) = a2 * b3 + a3 * b2 <- (L(a2) + L(a3)) * (R(b2) + R(b3)) + NEG(L(a2) * R(b2)) + SUB(L(a3) * R(b3))
508 // c(s3^2) = a3 * b3 <- (L(a3) * R(b3))
509
510 // Derive the offset. If res_offset < a_base2k, then we shift to a negative offset
511 // since the convolution doesn't support negative offset (yet).
512 let (res_offset_hi, res_offset_lo) = if res_offset < a_base2k {
513 (0, -((a_base2k - (res_offset % a_base2k)) as i64))
514 } else {
515 ((res_offset / a_base2k).saturating_sub(1), (res_offset % a_base2k) as i64)
516 };
517
518 let res_dft_size = res
519 .k()
520 .as_usize()
521 .div_ceil(a.base2k().as_usize())
522 .min(a.size() + b.size() - res_offset_hi);
523
524 for i in 0..cols {
525 let col_i: usize = i * cols - (i * (i + 1) / 2);
526
527 let (mut res_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, res_dft_size);
528 self.cnv_apply_dft(&mut res_dft, res_offset_hi, 0, &a_prep, i, &b_prep, i, scratch_3);
529 let res_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(res_dft);
530 let (mut tmp, scratch_4) = scratch_3.take_vec_znx(self.n(), 1, res_dft_size);
531 self.vec_znx_big_normalize(&mut tmp, res_base2k, res_offset_lo, 0, &res_big, a_base2k, 0, scratch_4);
532
533 self.vec_znx_copy(res.data_mut(), col_i + i, &tmp, 0);
534
535 // Pre-subtracts
536 // res[i!=j] = NEG(a[i] * b[i]) + SUB(a[j] * b[j])
537 for j in 0..cols {
538 if j != i {
539 if j < i {
540 let col_j = j * cols - (j * (j + 1) / 2);
541 self.vec_znx_sub_inplace(res.data_mut(), col_j + i, &tmp, 0);
542 } else {
543 self.vec_znx_negate(res.data_mut(), col_i + j, &tmp, 0);
544 }
545 }
546 }
547 }
548
549 for i in 0..cols {
550 let col_i: usize = i * cols - (i * (i + 1) / 2);
551
552 for j in i..cols {
553 if j != i {
554 // res_dft = (a[i] + a[j]) * (b[i] + b[j])
555 let (mut res_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, res.size());
556 self.cnv_pairwise_apply_dft(&mut res_dft, res_offset_hi, 0, &a_prep, &b_prep, i, j, scratch_3);
557 let res_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(res_dft);
558 let (mut tmp, scratch_3) = scratch_3.take_vec_znx(self.n(), 1, res.size());
559 self.vec_znx_big_normalize(&mut tmp, res_base2k, res_offset_lo, 0, &res_big, a_base2k, 0, scratch_3);
560
561 self.vec_znx_add_inplace(res.data_mut(), col_i + j, &tmp, 0);
562 }
563 }
564 }
565 }
566}
567
568pub trait GLWEAdd
569where
570 Self: ModuleN + VecZnxAdd + VecZnxCopy + VecZnxAddInplace + VecZnxZero,
571{
572 fn glwe_add<R, A, B>(&self, res: &mut R, a: &A, b: &B)
573 where
574 R: GLWEToMut,
575 A: GLWEToRef,
576 B: GLWEToRef,
577 {
578 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
579 let a: &mut GLWE<&[u8]> = &mut a.to_ref();
580 let b: &GLWE<&[u8]> = &b.to_ref();
581
582 assert_eq!(a.n(), self.n() as u32);
583 assert_eq!(b.n(), self.n() as u32);
584 assert_eq!(res.n(), self.n() as u32);
585 assert_eq!(a.base2k(), b.base2k());
586 assert_eq!(res.base2k(), b.base2k());
587
588 if a.rank() == 0 {
589 assert_eq!(res.rank(), b.rank());
590 } else if b.rank() == 0 {
591 assert_eq!(res.rank(), a.rank());
592 } else {
593 assert_eq!(res.rank(), a.rank());
594 assert_eq!(res.rank(), b.rank());
595 }
596
597 let min_col: usize = (a.rank().min(b.rank()) + 1).into();
598 let max_col: usize = (a.rank().max(b.rank() + 1)).into();
599 let self_col: usize = (res.rank() + 1).into();
600
601 for i in 0..min_col {
602 self.vec_znx_add(res.data_mut(), i, a.data(), i, b.data(), i);
603 }
604
605 if a.rank() > b.rank() {
606 for i in min_col..max_col {
607 self.vec_znx_copy(res.data_mut(), i, a.data(), i);
608 }
609 } else {
610 for i in min_col..max_col {
611 self.vec_znx_copy(res.data_mut(), i, b.data(), i);
612 }
613 }
614
615 for i in max_col..self_col {
616 self.vec_znx_zero(res.data_mut(), i);
617 }
618 }
619
620 fn glwe_add_inplace<R, A>(&self, res: &mut R, a: &A)
621 where
622 R: GLWEToMut,
623 A: GLWEToRef,
624 {
625 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
626 let a: &GLWE<&[u8]> = &a.to_ref();
627
628 assert_eq!(res.n(), self.n() as u32);
629 assert_eq!(a.n(), self.n() as u32);
630 assert_eq!(res.base2k(), a.base2k());
631 assert!(res.rank() >= a.rank());
632
633 for i in 0..(a.rank() + 1).into() {
634 self.vec_znx_add_inplace(res.data_mut(), i, a.data(), i);
635 }
636 }
637}
638
639impl<BE: Backend> GLWEAdd for Module<BE> where Self: ModuleN + VecZnxAdd + VecZnxCopy + VecZnxAddInplace + VecZnxZero {}
640
641impl<BE: Backend> GLWESub for Module<BE> where
642 Self: ModuleN + VecZnxSub + VecZnxCopy + VecZnxNegate + VecZnxZero + VecZnxSubInplace + VecZnxSubNegateInplace
643{
644}
645
646pub trait GLWESub
647where
648 Self: ModuleN + VecZnxSub + VecZnxCopy + VecZnxNegate + VecZnxZero + VecZnxSubInplace + VecZnxSubNegateInplace,
649{
650 fn glwe_sub<R, A, B>(&self, res: &mut R, a: &A, b: &B)
651 where
652 R: GLWEToMut,
653 A: GLWEToRef,
654 B: GLWEToRef,
655 {
656 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
657 let a: &GLWE<&[u8]> = &a.to_ref();
658 let b: &GLWE<&[u8]> = &b.to_ref();
659
660 assert_eq!(a.n(), self.n() as u32);
661 assert_eq!(b.n(), self.n() as u32);
662 assert_eq!(res.n(), self.n() as u32);
663 assert_eq!(a.base2k(), res.base2k());
664 assert_eq!(b.base2k(), res.base2k());
665
666 if a.rank() == 0 {
667 assert_eq!(res.rank(), b.rank());
668 } else if b.rank() == 0 {
669 assert_eq!(res.rank(), a.rank());
670 } else {
671 assert_eq!(res.rank(), a.rank());
672 assert_eq!(res.rank(), b.rank());
673 }
674
675 let min_col: usize = (a.rank().min(b.rank()) + 1).into();
676 let max_col: usize = (a.rank().max(b.rank() + 1)).into();
677 let self_col: usize = (res.rank() + 1).into();
678
679 for i in 0..min_col {
680 self.vec_znx_sub(res.data_mut(), i, a.data(), i, b.data(), i);
681 }
682
683 if a.rank() > b.rank() {
684 for i in min_col..max_col {
685 self.vec_znx_copy(res.data_mut(), i, a.data(), i);
686 }
687 } else {
688 for i in min_col..max_col {
689 self.vec_znx_negate(res.data_mut(), i, b.data(), i);
690 }
691 }
692
693 for i in max_col..self_col {
694 self.vec_znx_zero(res.data_mut(), i);
695 }
696 }
697
698 fn glwe_sub_inplace<R, A>(&self, res: &mut R, a: &A)
699 where
700 R: GLWEToMut,
701 A: GLWEToRef,
702 {
703 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
704 let a: &GLWE<&[u8]> = &a.to_ref();
705
706 assert_eq!(res.n(), self.n() as u32);
707 assert_eq!(a.n(), self.n() as u32);
708 assert_eq!(res.base2k(), a.base2k());
709 assert!(res.rank() == a.rank() || a.rank() == 0);
710
711 for i in 0..(a.rank() + 1).into() {
712 self.vec_znx_sub_inplace(res.data_mut(), i, a.data(), i);
713 }
714 }
715
716 fn glwe_sub_negate_inplace<R, A>(&self, res: &mut R, a: &A)
717 where
718 R: GLWEToMut,
719 A: GLWEToRef,
720 {
721 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
722 let a: &GLWE<&[u8]> = &a.to_ref();
723
724 assert_eq!(res.n(), self.n() as u32);
725 assert_eq!(a.n(), self.n() as u32);
726 assert_eq!(res.base2k(), a.base2k());
727 assert!(res.rank() == a.rank() || a.rank() == 0);
728
729 for i in 0..(a.rank() + 1).into() {
730 self.vec_znx_sub_negate_inplace(res.data_mut(), i, a.data(), i);
731 }
732 }
733}
734
735impl<BE: Backend> GLWERotate<BE> for Module<BE> where Self: ModuleN + VecZnxRotate + VecZnxRotateInplace<BE> + VecZnxZero {}
736
737pub trait GLWERotate<BE: Backend>
738where
739 Self: ModuleN + VecZnxRotate + VecZnxRotateInplace<BE> + VecZnxZero,
740{
741 fn glwe_rotate_tmp_bytes(&self) -> usize {
742 vec_znx_rotate_inplace_tmp_bytes(self.n())
743 }
744
745 fn glwe_rotate<R, A>(&self, k: i64, res: &mut R, a: &A)
746 where
747 R: GLWEToMut,
748 A: GLWEToRef,
749 {
750 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
751 let a: &GLWE<&[u8]> = &a.to_ref();
752
753 assert_eq!(a.n(), self.n() as u32);
754 assert_eq!(res.n(), self.n() as u32);
755 assert!(res.rank() == a.rank() || a.rank() == 0);
756
757 let res_cols = (res.rank() + 1).into();
758 let a_cols = (a.rank() + 1).into();
759
760 for i in 0..a_cols {
761 self.vec_znx_rotate(k, res.data_mut(), i, a.data(), i);
762 }
763 for i in a_cols..res_cols {
764 self.vec_znx_zero(res.data_mut(), i);
765 }
766 }
767
768 fn glwe_rotate_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
769 where
770 R: GLWEToMut,
771 Scratch<BE>: ScratchTakeCore<BE>,
772 {
773 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
774
775 for i in 0..(res.rank() + 1).into() {
776 self.vec_znx_rotate_inplace(k, res.data_mut(), i, scratch);
777 }
778 }
779}
780
781impl<BE: Backend> GLWEMulXpMinusOne<BE> for Module<BE> where Self: ModuleN + VecZnxMulXpMinusOne + VecZnxMulXpMinusOneInplace<BE> {}
782
783pub trait GLWEMulXpMinusOne<BE: Backend>
784where
785 Self: ModuleN + VecZnxMulXpMinusOne + VecZnxMulXpMinusOneInplace<BE>,
786{
787 fn glwe_mul_xp_minus_one<R, A>(&self, k: i64, res: &mut R, a: &A)
788 where
789 R: GLWEToMut,
790 A: GLWEToRef,
791 {
792 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
793 let a: &GLWE<&[u8]> = &a.to_ref();
794
795 assert_eq!(res.n(), self.n() as u32);
796 assert_eq!(a.n(), self.n() as u32);
797 assert_eq!(res.rank(), a.rank());
798
799 for i in 0..res.rank().as_usize() + 1 {
800 self.vec_znx_mul_xp_minus_one(k, res.data_mut(), i, a.data(), i);
801 }
802 }
803
804 fn glwe_mul_xp_minus_one_inplace<R>(&self, k: i64, res: &mut R, scratch: &mut Scratch<BE>)
805 where
806 R: GLWEToMut,
807 {
808 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
809
810 assert_eq!(res.n(), self.n() as u32);
811
812 for i in 0..res.rank().as_usize() + 1 {
813 self.vec_znx_mul_xp_minus_one_inplace(k, res.data_mut(), i, scratch);
814 }
815 }
816}
817
818impl<BE: Backend> GLWECopy for Module<BE> where Self: ModuleN + VecZnxCopy + VecZnxZero {}
819
820pub trait GLWECopy
821where
822 Self: ModuleN + VecZnxCopy + VecZnxZero,
823{
824 fn glwe_copy<R, A>(&self, res: &mut R, a: &A)
825 where
826 R: GLWEToMut,
827 A: GLWEToRef,
828 {
829 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
830 let a: &GLWE<&[u8]> = &a.to_ref();
831
832 assert_eq!(res.n(), self.n() as u32);
833 assert_eq!(a.n(), self.n() as u32);
834 assert!(res.rank() == a.rank() || a.rank() == 0);
835
836 let min_rank: usize = res.rank().min(a.rank()).as_usize() + 1;
837
838 for i in 0..min_rank {
839 self.vec_znx_copy(res.data_mut(), i, a.data(), i);
840 }
841
842 for i in min_rank..(res.rank() + 1).into() {
843 self.vec_znx_zero(res.data_mut(), i);
844 }
845 }
846}
847
848impl<BE: Backend> GLWEShift<BE> for Module<BE> where Self: ModuleN + VecZnxRshInplace<BE> {}
849
850pub trait GLWEShift<BE: Backend>
851where
852 Self: ModuleN + VecZnxRshInplace<BE>,
853{
854 fn glwe_rsh_tmp_byte(&self) -> usize {
855 VecZnx::rsh_tmp_bytes(self.n())
856 }
857
858 fn glwe_rsh<R>(&self, k: usize, res: &mut R, scratch: &mut Scratch<BE>)
859 where
860 R: GLWEToMut,
861 Scratch<BE>: ScratchTakeCore<BE>,
862 {
863 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
864 let base2k: usize = res.base2k().into();
865 for i in 0..res.rank().as_usize() + 1 {
866 self.vec_znx_rsh_inplace(base2k, k, res.data_mut(), i, scratch);
867 }
868 }
869}
870
871impl GLWE<Vec<u8>> {
872 pub fn rsh_tmp_bytes<M, BE: Backend>(module: &M) -> usize
873 where
874 M: GLWEShift<BE>,
875 {
876 module.glwe_rsh_tmp_byte()
877 }
878}
879
880impl<BE: Backend> GLWENormalize<BE> for Module<BE> where
881 Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> + VecZnxNormalizeTmpBytes
882{
883}
884
885pub trait GLWENormalize<BE: Backend>
886where
887 Self: ModuleN + VecZnxNormalize<BE> + VecZnxNormalizeInplace<BE> + VecZnxNormalizeTmpBytes,
888{
889 fn glwe_normalize_tmp_bytes(&self) -> usize {
890 self.vec_znx_normalize_tmp_bytes()
891 }
892
893 /// Usage:
894 /// let mut tmp_b: Option<GLWE<&mut [u8]>> = None;
895 /// let (b_conv, scratch_1) = glwe_maybe_convert_in_place(self, b, res.base2k().as_u32(), &mut tmp_b, scratch);
896 fn glwe_maybe_cross_normalize_to_ref<'a, A>(
897 &self,
898 glwe: &'a A,
899 target_base2k: usize,
900 tmp_slot: &'a mut Option<GLWE<&'a mut [u8]>>, // caller-owned scratch-backed temp
901 scratch: &'a mut Scratch<BE>,
902 ) -> (GLWE<&'a [u8]>, &'a mut Scratch<BE>)
903 where
904 A: GLWEToRef + GLWEInfos,
905 Scratch<BE>: ScratchTakeCore<BE>,
906 {
907 // No conversion: just use the original GLWE
908 if glwe.base2k().as_usize() == target_base2k {
909 // Drop any previous temp; it's stale for this base
910 tmp_slot.take();
911 return (glwe.to_ref(), scratch);
912 }
913
914 // Conversion: allocate a temporary GLWE in scratch
915 let mut layout = glwe.glwe_layout();
916 layout.base2k = target_base2k.into();
917
918 let (tmp, scratch2) = scratch.take_glwe(&layout);
919 *tmp_slot = Some(tmp);
920
921 // Get a mutable handle to the temp and normalize into it
922 let tmp_ref: &mut GLWE<&mut [u8]> = tmp_slot.as_mut().expect("tmp_slot just set to Some, but found None");
923
924 self.glwe_normalize(tmp_ref, glwe, scratch2);
925
926 // Return a trait-object view of the temp
927 (tmp_ref.to_ref(), scratch2)
928 }
929
930 /// Usage:
931 /// let mut tmp_b: Option<GLWE<&mut [u8]>> = None;
932 /// let (b_conv, scratch_1) = glwe_maybe_convert_in_place(self, b, res.base2k().as_u32(), &mut tmp_b, scratch);
933 fn glwe_maybe_cross_normalize_to_mut<'a, A>(
934 &self,
935 glwe: &'a mut A,
936 target_base2k: usize,
937 tmp_slot: &'a mut Option<GLWE<&'a mut [u8]>>, // caller-owned scratch-backed temp
938 scratch: &'a mut Scratch<BE>,
939 ) -> (GLWE<&'a mut [u8]>, &'a mut Scratch<BE>)
940 where
941 A: GLWEToMut + GLWEInfos,
942 Scratch<BE>: ScratchTakeCore<BE>,
943 {
944 // No conversion: just use the original GLWE
945 if glwe.base2k().as_usize() == target_base2k {
946 // Drop any previous temp; it's stale for this base
947 tmp_slot.take();
948 return (glwe.to_mut(), scratch);
949 }
950
951 // Conversion: allocate a temporary GLWE in scratch
952 let mut layout = glwe.glwe_layout();
953 layout.base2k = target_base2k.into();
954
955 let (tmp, scratch2) = scratch.take_glwe(&layout);
956 *tmp_slot = Some(tmp);
957
958 // Get a mutable handle to the temp and normalize into it
959 let tmp_ref: &mut GLWE<&mut [u8]> = tmp_slot.as_mut().expect("tmp_slot just set to Some, but found None");
960
961 self.glwe_normalize(tmp_ref, glwe, scratch2);
962
963 // Return a trait-object view of the temp
964 (tmp_ref.to_mut(), scratch2)
965 }
966
967 fn glwe_normalize<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
968 where
969 R: GLWEToMut,
970 A: GLWEToRef,
971 Scratch<BE>: ScratchTakeCore<BE>,
972 {
973 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
974 let a: &GLWE<&[u8]> = &a.to_ref();
975
976 assert_eq!(res.n(), self.n() as u32);
977 assert_eq!(a.n(), self.n() as u32);
978 assert_eq!(res.rank(), a.rank());
979
980 let res_base2k = res.base2k().into();
981
982 for i in 0..res.rank().as_usize() + 1 {
983 self.vec_znx_normalize(res.data_mut(), res_base2k, 0, i, a.data(), a.base2k().into(), i, scratch);
984 }
985 }
986
987 fn glwe_normalize_inplace<R>(&self, res: &mut R, scratch: &mut Scratch<BE>)
988 where
989 R: GLWEToMut,
990 Scratch<BE>: ScratchTakeCore<BE>,
991 {
992 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
993 for i in 0..res.rank().as_usize() + 1 {
994 self.vec_znx_normalize_inplace(res.base2k().into(), res.data_mut(), i, scratch);
995 }
996 }
997}
998
999#[allow(dead_code)]
1000// c = op(a, b)
1001fn set_k_binary(c: &impl GLWEInfos, a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
1002 // If either operands is a ciphertext
1003 if a.rank() != 0 || b.rank() != 0 {
1004 // If a is a plaintext (but b ciphertext)
1005 let k = if a.rank() == 0 {
1006 b.k()
1007 // If b is a plaintext (but a ciphertext)
1008 } else if b.rank() == 0 {
1009 a.k()
1010 // If a & b are both ciphertexts
1011 } else {
1012 a.k().min(b.k())
1013 };
1014 k.min(c.k())
1015 // If a & b are both plaintexts
1016 } else {
1017 c.k()
1018 }
1019}
1020
1021#[allow(dead_code)]
1022// a = op(a, b)
1023fn set_k_unary(a: &impl GLWEInfos, b: &impl GLWEInfos) -> TorusPrecision {
1024 if a.rank() != 0 || b.rank() != 0 {
1025 a.k().min(b.k())
1026 } else {
1027 a.k()
1028 }
1029}