poulpy/poulpy-hal/src/reference/fft64/vec_znx_big.rs
poulpy-hal/src/reference/fft64/vec_znx_big.rsRaw
1use std::f64::consts::SQRT_2;
2
3use crate::{
4 api::VecZnxBigAddNormal,
5 layouts::{
6 Backend, Module, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxBigToRef, VecZnxToMut, VecZnxToRef, ZnxView, ZnxViewMut,
7 },
8 oep::VecZnxBigAllocBytesImpl,
9 reference::{
10 vec_znx::{
11 vec_znx_add, vec_znx_add_inplace, vec_znx_automorphism, vec_znx_automorphism_inplace, vec_znx_negate,
12 vec_znx_negate_inplace, vec_znx_normalize, vec_znx_sub, vec_znx_sub_inplace, vec_znx_sub_negate_inplace,
13 },
14 znx::{
15 ZnxAdd, ZnxAddInplace, ZnxAutomorphism, ZnxCopy, ZnxExtractDigitAddMul, ZnxMulPowerOfTwoInplace, ZnxNegate,
16 ZnxNegateInplace, ZnxNormalizeDigit, ZnxNormalizeFinalStep, ZnxNormalizeFirstStep, ZnxNormalizeFirstStepCarryOnly,
17 ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly, ZnxSub, ZnxSubInplace, ZnxSubNegateInplace, ZnxZero,
18 znx_add_normal_f64_ref,
19 },
20 },
21 source::Source,
22};
23
24pub fn vec_znx_big_add<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
25where
26 BE: Backend<ScalarBig = i64> + ZnxAdd + ZnxCopy + ZnxZero,
27 R: VecZnxBigToMut<BE>,
28 A: VecZnxBigToRef<BE>,
29 B: VecZnxBigToRef<BE>,
30{
31 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
32 let a: VecZnxBig<&[u8], BE> = a.to_ref();
33 let b: VecZnxBig<&[u8], BE> = b.to_ref();
34
35 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
36 data: res.data,
37 n: res.n,
38 cols: res.cols,
39 size: res.size,
40 max_size: res.max_size,
41 };
42
43 let a_vznx: VecZnx<&[u8]> = VecZnx {
44 data: a.data,
45 n: a.n,
46 cols: a.cols,
47 size: a.size,
48 max_size: a.max_size,
49 };
50
51 let b_vznx: VecZnx<&[u8]> = VecZnx {
52 data: b.data,
53 n: b.n,
54 cols: b.cols,
55 size: b.size,
56 max_size: b.max_size,
57 };
58
59 vec_znx_add::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, &b_vznx, b_col);
60}
61
62pub fn vec_znx_big_add_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
63where
64 BE: Backend<ScalarBig = i64> + ZnxAddInplace,
65 R: VecZnxBigToMut<BE>,
66 A: VecZnxBigToRef<BE>,
67{
68 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
69 let a: VecZnxBig<&[u8], BE> = a.to_ref();
70
71 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
72 data: res.data,
73 n: res.n,
74 cols: res.cols,
75 size: res.size,
76 max_size: res.max_size,
77 };
78
79 let a_vznx: VecZnx<&[u8]> = VecZnx {
80 data: a.data,
81 n: a.n,
82 cols: a.cols,
83 size: a.size,
84 max_size: a.max_size,
85 };
86
87 vec_znx_add_inplace::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
88}
89
90pub fn vec_znx_big_add_small<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
91where
92 BE: Backend<ScalarBig = i64> + ZnxAdd + ZnxCopy + ZnxZero,
93 R: VecZnxBigToMut<BE>,
94 A: VecZnxBigToRef<BE>,
95 B: VecZnxToRef,
96{
97 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
98 let a: VecZnxBig<&[u8], BE> = a.to_ref();
99
100 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
101 data: res.data,
102 n: res.n,
103 cols: res.cols,
104 size: res.size,
105 max_size: res.max_size,
106 };
107
108 let a_vznx: VecZnx<&[u8]> = VecZnx {
109 data: a.data,
110 n: a.n,
111 cols: a.cols,
112 size: a.size,
113 max_size: a.max_size,
114 };
115
116 vec_znx_add::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, b, b_col);
117}
118
119pub fn vec_znx_big_add_small_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
120where
121 BE: Backend<ScalarBig = i64> + ZnxAddInplace,
122 R: VecZnxBigToMut<BE>,
123 A: VecZnxToRef,
124{
125 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
126 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
127 data: res.data,
128 n: res.n,
129 cols: res.cols,
130 size: res.size,
131 max_size: res.max_size,
132 };
133
134 vec_znx_add_inplace::<_, _, BE>(&mut res_vznx, res_col, a, a_col);
135}
136
137pub fn vec_znx_big_automorphism_inplace_tmp_bytes(n: usize) -> usize {
138 n * size_of::<i64>()
139}
140
141pub fn vec_znx_big_automorphism<R, A, BE>(p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
142where
143 BE: Backend<ScalarBig = i64> + ZnxAutomorphism + ZnxZero,
144 R: VecZnxBigToMut<BE>,
145 A: VecZnxBigToRef<BE>,
146{
147 let res: VecZnxBig<&mut [u8], _> = res.to_mut();
148 let a: VecZnxBig<&[u8], _> = a.to_ref();
149
150 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
151 data: res.data,
152 n: res.n,
153 cols: res.cols,
154 size: res.size,
155 max_size: res.max_size,
156 };
157
158 let a_vznx: VecZnx<&[u8]> = VecZnx {
159 data: a.data,
160 n: a.n,
161 cols: a.cols,
162 size: a.size,
163 max_size: a.max_size,
164 };
165
166 vec_znx_automorphism::<_, _, BE>(p, &mut res_vznx, res_col, &a_vznx, a_col);
167}
168
169pub fn vec_znx_big_automorphism_inplace<R, BE>(p: i64, res: &mut R, res_col: usize, tmp: &mut [i64])
170where
171 BE: Backend<ScalarBig = i64> + ZnxAutomorphism + ZnxCopy,
172 R: VecZnxBigToMut<BE>,
173{
174 let res: VecZnxBig<&mut [u8], _> = res.to_mut();
175
176 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
177 data: res.data,
178 n: res.n,
179 cols: res.cols,
180 size: res.size,
181 max_size: res.max_size,
182 };
183
184 vec_znx_automorphism_inplace::<_, BE>(p, &mut res_vznx, res_col, tmp);
185}
186
187pub fn vec_znx_big_negate<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
188where
189 BE: Backend<ScalarBig = i64> + ZnxNegate + ZnxZero,
190 R: VecZnxBigToMut<BE>,
191 A: VecZnxBigToRef<BE>,
192{
193 let res: VecZnxBig<&mut [u8], _> = res.to_mut();
194 let a: VecZnxBig<&[u8], _> = a.to_ref();
195
196 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
197 data: res.data,
198 n: res.n,
199 cols: res.cols,
200 size: res.size,
201 max_size: res.max_size,
202 };
203
204 let a_vznx: VecZnx<&[u8]> = VecZnx {
205 data: a.data,
206 n: a.n,
207 cols: a.cols,
208 size: a.size,
209 max_size: a.max_size,
210 };
211
212 vec_znx_negate::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
213}
214
215pub fn vec_znx_big_negate_inplace<R, BE>(res: &mut R, res_col: usize)
216where
217 BE: Backend<ScalarBig = i64> + ZnxNegateInplace,
218 R: VecZnxBigToMut<BE>,
219{
220 let res: VecZnxBig<&mut [u8], _> = res.to_mut();
221
222 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
223 data: res.data,
224 n: res.n,
225 cols: res.cols,
226 size: res.size,
227 max_size: res.max_size,
228 };
229
230 vec_znx_negate_inplace::<_, BE>(&mut res_vznx, res_col);
231}
232
233pub fn vec_znx_big_normalize_tmp_bytes(n: usize) -> usize {
234 2 * n * size_of::<i64>()
235}
236
237pub fn vec_znx_big_normalize<R, A, BE>(
238 res_basek: usize,
239 res: &mut R,
240 res_col: usize,
241 a_basek: usize,
242 a: &A,
243 a_col: usize,
244 carry: &mut [i64],
245) where
246 R: VecZnxToMut,
247 A: VecZnxBigToRef<BE>,
248 BE: Backend<ScalarBig = i64>
249 + ZnxZero
250 + ZnxCopy
251 + ZnxAddInplace
252 + ZnxMulPowerOfTwoInplace
253 + ZnxNormalizeFirstStepCarryOnly
254 + ZnxNormalizeMiddleStepCarryOnly
255 + ZnxNormalizeMiddleStep
256 + ZnxNormalizeFinalStep
257 + ZnxNormalizeFirstStep
258 + ZnxExtractDigitAddMul
259 + ZnxNormalizeDigit,
260{
261 let a: VecZnxBig<&[u8], _> = a.to_ref();
262 let a_vznx: VecZnx<&[u8]> = VecZnx {
263 data: a.data,
264 n: a.n,
265 cols: a.cols,
266 size: a.size,
267 max_size: a.max_size,
268 };
269
270 vec_znx_normalize::<_, _, BE>(res_basek, res, res_col, a_basek, &a_vznx, a_col, carry);
271}
272
273pub fn vec_znx_big_add_normal_ref<R, B: Backend<ScalarBig = i64>>(
274 base2k: usize,
275 res: &mut R,
276 res_col: usize,
277 k: usize,
278 sigma: f64,
279 bound: f64,
280 source: &mut Source,
281) where
282 R: VecZnxBigToMut<B>,
283{
284 let mut res: VecZnxBig<&mut [u8], B> = res.to_mut();
285 assert!(
286 (bound.log2().ceil() as i64) < 64,
287 "invalid bound: ceil(log2(bound))={} > 63",
288 (bound.log2().ceil() as i64)
289 );
290
291 let limb: usize = k.div_ceil(base2k) - 1;
292 let scale: f64 = (1 << ((limb + 1) * base2k - k)) as f64;
293 znx_add_normal_f64_ref(
294 res.at_mut(res_col, limb),
295 sigma * scale,
296 bound * scale,
297 source,
298 )
299}
300
301pub fn test_vec_znx_big_add_normal<B>(module: &Module<B>)
302where
303 Module<B>: VecZnxBigAddNormal<B>,
304 B: Backend<ScalarBig = i64> + VecZnxBigAllocBytesImpl<B>,
305{
306 let n: usize = module.n();
307 let base2k: usize = 17;
308 let k: usize = 2 * 17;
309 let size: usize = 5;
310 let sigma: f64 = 3.2;
311 let bound: f64 = 6.0 * sigma;
312 let mut source: Source = Source::new([0u8; 32]);
313 let cols: usize = 2;
314 let zero: Vec<i64> = vec![0; n];
315 let k_f64: f64 = (1u64 << k as u64) as f64;
316 let sqrt2: f64 = SQRT_2;
317 (0..cols).for_each(|col_i| {
318 let mut a: VecZnxBig<Vec<u8>, B> = VecZnxBig::alloc(n, cols, size);
319 module.vec_znx_big_add_normal(base2k, &mut a, col_i, k, &mut source, sigma, bound);
320 module.vec_znx_big_add_normal(base2k, &mut a, col_i, k, &mut source, sigma, bound);
321 (0..cols).for_each(|col_j| {
322 if col_j != col_i {
323 (0..size).for_each(|limb_i| {
324 assert_eq!(a.at(col_j, limb_i), zero);
325 })
326 } else {
327 let std: f64 = a.stats(base2k, col_i).std() * k_f64;
328 assert!(
329 (std - sigma * sqrt2).abs() < 0.1,
330 "std={} ~!= {}",
331 std,
332 sigma * sqrt2
333 );
334 }
335 })
336 });
337}
338
339/// R <- A - B
340pub fn vec_znx_big_sub<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
341where
342 BE: Backend<ScalarBig = i64> + ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
343 R: VecZnxBigToMut<BE>,
344 A: VecZnxBigToRef<BE>,
345 B: VecZnxBigToRef<BE>,
346{
347 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
348 let a: VecZnxBig<&[u8], BE> = a.to_ref();
349 let b: VecZnxBig<&[u8], BE> = b.to_ref();
350
351 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
352 data: res.data,
353 n: res.n,
354 cols: res.cols,
355 size: res.size,
356 max_size: res.max_size,
357 };
358
359 let a_vznx: VecZnx<&[u8]> = VecZnx {
360 data: a.data,
361 n: a.n,
362 cols: a.cols,
363 size: a.size,
364 max_size: a.max_size,
365 };
366
367 let b_vznx: VecZnx<&[u8]> = VecZnx {
368 data: b.data,
369 n: b.n,
370 cols: b.cols,
371 size: b.size,
372 max_size: b.max_size,
373 };
374
375 vec_znx_sub::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, &b_vznx, b_col);
376}
377
378/// R <- A - B
379pub fn vec_znx_big_sub_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
380where
381 BE: Backend<ScalarBig = i64> + ZnxSubInplace,
382 R: VecZnxBigToMut<BE>,
383 A: VecZnxBigToRef<BE>,
384{
385 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
386 let a: VecZnxBig<&[u8], BE> = a.to_ref();
387
388 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
389 data: res.data,
390 n: res.n,
391 cols: res.cols,
392 size: res.size,
393 max_size: res.max_size,
394 };
395
396 let a_vznx: VecZnx<&[u8]> = VecZnx {
397 data: a.data,
398 n: a.n,
399 cols: a.cols,
400 size: a.size,
401 max_size: a.max_size,
402 };
403
404 vec_znx_sub_inplace::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
405}
406
407/// R <- B - A
408pub fn vec_znx_big_sub_negate_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
409where
410 BE: Backend<ScalarBig = i64> + ZnxSubNegateInplace + ZnxNegateInplace,
411 R: VecZnxBigToMut<BE>,
412 A: VecZnxBigToRef<BE>,
413{
414 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
415 let a: VecZnxBig<&[u8], BE> = a.to_ref();
416
417 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
418 data: res.data,
419 n: res.n,
420 cols: res.cols,
421 size: res.size,
422 max_size: res.max_size,
423 };
424
425 let a_vznx: VecZnx<&[u8]> = VecZnx {
426 data: a.data,
427 n: a.n,
428 cols: a.cols,
429 size: a.size,
430 max_size: a.max_size,
431 };
432
433 vec_znx_sub_negate_inplace::<_, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col);
434}
435
436/// R <- A - B
437pub fn vec_znx_big_sub_small_a<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
438where
439 BE: Backend<ScalarBig = i64> + ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
440 R: VecZnxBigToMut<BE>,
441 A: VecZnxToRef,
442 B: VecZnxBigToRef<BE>,
443{
444 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
445 let b: VecZnxBig<&[u8], BE> = b.to_ref();
446
447 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
448 data: res.data,
449 n: res.n,
450 cols: res.cols,
451 size: res.size,
452 max_size: res.max_size,
453 };
454
455 let b_vznx: VecZnx<&[u8]> = VecZnx {
456 data: b.data,
457 n: b.n,
458 cols: b.cols,
459 size: b.size,
460 max_size: b.max_size,
461 };
462
463 vec_znx_sub::<_, _, _, BE>(&mut res_vznx, res_col, a, a_col, &b_vznx, b_col);
464}
465
466/// R <- A - B
467pub fn vec_znx_big_sub_small_b<R, A, B, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
468where
469 BE: Backend<ScalarBig = i64> + ZnxSub + ZnxNegate + ZnxZero + ZnxCopy,
470 R: VecZnxBigToMut<BE>,
471 A: VecZnxBigToRef<BE>,
472 B: VecZnxToRef,
473{
474 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
475 let a: VecZnxBig<&[u8], BE> = a.to_ref();
476
477 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
478 data: res.data,
479 n: res.n,
480 cols: res.cols,
481 size: res.size,
482 max_size: res.max_size,
483 };
484
485 let a_vznx: VecZnx<&[u8]> = VecZnx {
486 data: a.data,
487 n: a.n,
488 cols: a.cols,
489 size: a.size,
490 max_size: a.max_size,
491 };
492
493 vec_znx_sub::<_, _, _, BE>(&mut res_vznx, res_col, &a_vznx, a_col, b, b_col);
494}
495
496/// R <- R - A
497pub fn vec_znx_big_sub_small_a_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
498where
499 BE: Backend<ScalarBig = i64> + ZnxSubInplace,
500 R: VecZnxBigToMut<BE>,
501 A: VecZnxToRef,
502{
503 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
504
505 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
506 data: res.data,
507 n: res.n,
508 cols: res.cols,
509 size: res.size,
510 max_size: res.max_size,
511 };
512
513 vec_znx_sub_inplace::<_, _, BE>(&mut res_vznx, res_col, a, a_col);
514}
515
516/// R <- A - R
517pub fn vec_znx_big_sub_small_b_inplace<R, A, BE>(res: &mut R, res_col: usize, a: &A, a_col: usize)
518where
519 BE: Backend<ScalarBig = i64> + ZnxSubNegateInplace + ZnxNegateInplace,
520 R: VecZnxBigToMut<BE>,
521 A: VecZnxToRef,
522{
523 let res: VecZnxBig<&mut [u8], BE> = res.to_mut();
524
525 let mut res_vznx: VecZnx<&mut [u8]> = VecZnx {
526 data: res.data,
527 n: res.n,
528 cols: res.cols,
529 size: res.size,
530 max_size: res.max_size,
531 };
532
533 vec_znx_sub_negate_inplace::<_, _, BE>(&mut res_vznx, res_col, a, a_col);
534}