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