poulpy/poulpy-backend/src/cpu_spqlios/fft64/vec_znx_big.rs
poulpy-backend/src/cpu_spqlios/fft64/vec_znx_big.rsRaw
1use rand_distr::{Distribution, Normal};
2
3use crate::cpu_spqlios::{FFT64, ffi::vec_znx};
4use poulpy_hal::{
5 api::{
6 TakeSlice, VecZnxBigAddDistF64, VecZnxBigFillDistF64, VecZnxBigNormalizeTmpBytes, ZnxInfos, ZnxSliceSize, ZnxView,
7 ZnxViewMut,
8 },
9 layouts::{
10 Backend, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxBigToMut, VecZnxBigToRef, VecZnxToMut, VecZnxToRef,
11 },
12 oep::{
13 TakeSliceImpl, VecZnxBigAddDistF64Impl, VecZnxBigAddImpl, VecZnxBigAddInplaceImpl, VecZnxBigAddNormalImpl,
14 VecZnxBigAddSmallImpl, VecZnxBigAddSmallInplaceImpl, VecZnxBigAllocBytesImpl, VecZnxBigAllocImpl,
15 VecZnxBigAutomorphismImpl, VecZnxBigAutomorphismInplaceImpl, VecZnxBigFillDistF64Impl, VecZnxBigFillNormalImpl,
16 VecZnxBigFromBytesImpl, VecZnxBigNegateInplaceImpl, VecZnxBigNormalizeImpl, VecZnxBigNormalizeTmpBytesImpl,
17 VecZnxBigSubABInplaceImpl, VecZnxBigSubBAInplaceImpl, VecZnxBigSubImpl, VecZnxBigSubSmallAImpl,
18 VecZnxBigSubSmallAInplaceImpl, VecZnxBigSubSmallBImpl, VecZnxBigSubSmallBInplaceImpl,
19 },
20 source::Source,
21};
22
23unsafe impl VecZnxBigAllocBytesImpl<Self> for FFT64 {
24 fn vec_znx_big_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
25 Self::layout_big_word_count() * n * cols * size * size_of::<f64>()
26 }
27}
28
29unsafe impl VecZnxBigAllocImpl<Self> for FFT64 {
30 fn vec_znx_big_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxBigOwned<Self> {
31 VecZnxBig::alloc(n, cols, size)
32 }
33}
34
35unsafe impl VecZnxBigFromBytesImpl<Self> for FFT64 {
36 fn vec_znx_big_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<Self> {
37 VecZnxBig::from_bytes(n, cols, size, bytes)
38 }
39}
40
41unsafe impl VecZnxBigAddDistF64Impl<Self> for FFT64 {
42 fn add_dist_f64_impl<R: VecZnxBigToMut<Self>, D: Distribution<f64>>(
43 _module: &Module<Self>,
44 basek: usize,
45 res: &mut R,
46 res_col: usize,
47 k: usize,
48 source: &mut Source,
49 dist: D,
50 bound: f64,
51 ) {
52 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
53 assert!(
54 (bound.log2().ceil() as i64) < 64,
55 "invalid bound: ceil(log2(bound))={} > 63",
56 (bound.log2().ceil() as i64)
57 );
58
59 let limb: usize = k.div_ceil(basek) - 1;
60 let basek_rem: usize = (limb + 1) * basek - k;
61
62 if basek_rem != 0 {
63 res.at_mut(res_col, limb).iter_mut().for_each(|x| {
64 let mut dist_f64: f64 = dist.sample(source);
65 while dist_f64.abs() > bound {
66 dist_f64 = dist.sample(source)
67 }
68 *x += (dist_f64.round() as i64) << basek_rem;
69 });
70 } else {
71 res.at_mut(res_col, limb).iter_mut().for_each(|x| {
72 let mut dist_f64: f64 = dist.sample(source);
73 while dist_f64.abs() > bound {
74 dist_f64 = dist.sample(source)
75 }
76 *x += dist_f64.round() as i64
77 });
78 }
79 }
80}
81
82unsafe impl VecZnxBigAddNormalImpl<Self> for FFT64 {
83 fn add_normal_impl<R: VecZnxBigToMut<Self>>(
84 module: &Module<Self>,
85 basek: usize,
86 res: &mut R,
87 res_col: usize,
88 k: usize,
89 source: &mut Source,
90 sigma: f64,
91 bound: f64,
92 ) {
93 module.vec_znx_big_add_dist_f64(
94 basek,
95 res,
96 res_col,
97 k,
98 source,
99 Normal::new(0.0, sigma).unwrap(),
100 bound,
101 );
102 }
103}
104
105unsafe impl VecZnxBigFillDistF64Impl<Self> for FFT64 {
106 fn fill_dist_f64_impl<R: VecZnxBigToMut<Self>, D: Distribution<f64>>(
107 _module: &Module<Self>,
108 basek: usize,
109 res: &mut R,
110 res_col: usize,
111 k: usize,
112 source: &mut Source,
113 dist: D,
114 bound: f64,
115 ) {
116 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
117 assert!(
118 (bound.log2().ceil() as i64) < 64,
119 "invalid bound: ceil(log2(bound))={} > 63",
120 (bound.log2().ceil() as i64)
121 );
122
123 let limb: usize = k.div_ceil(basek) - 1;
124 let basek_rem: usize = (limb + 1) * basek - k;
125
126 if basek_rem != 0 {
127 res.at_mut(res_col, limb).iter_mut().for_each(|x| {
128 let mut dist_f64: f64 = dist.sample(source);
129 while dist_f64.abs() > bound {
130 dist_f64 = dist.sample(source)
131 }
132 *x = (dist_f64.round() as i64) << basek_rem;
133 });
134 } else {
135 res.at_mut(res_col, limb).iter_mut().for_each(|x| {
136 let mut dist_f64: f64 = dist.sample(source);
137 while dist_f64.abs() > bound {
138 dist_f64 = dist.sample(source)
139 }
140 *x = dist_f64.round() as i64
141 });
142 }
143 }
144}
145
146unsafe impl VecZnxBigFillNormalImpl<Self> for FFT64 {
147 fn fill_normal_impl<R: VecZnxBigToMut<Self>>(
148 module: &Module<Self>,
149 basek: usize,
150 res: &mut R,
151 res_col: usize,
152 k: usize,
153 source: &mut Source,
154 sigma: f64,
155 bound: f64,
156 ) {
157 module.vec_znx_big_fill_dist_f64(
158 basek,
159 res,
160 res_col,
161 k,
162 source,
163 Normal::new(0.0, sigma).unwrap(),
164 bound,
165 );
166 }
167}
168
169unsafe impl VecZnxBigAddImpl<Self> for FFT64 {
170 /// Adds `a` to `b` and stores the result on `c`.
171 fn vec_znx_big_add_impl<R, A, B>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
172 where
173 R: VecZnxBigToMut<Self>,
174 A: VecZnxBigToRef<Self>,
175 B: VecZnxBigToRef<Self>,
176 {
177 let a: VecZnxBig<&[u8], Self> = a.to_ref();
178 let b: VecZnxBig<&[u8], Self> = b.to_ref();
179 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
180
181 #[cfg(debug_assertions)]
182 {
183 assert_eq!(a.n(), res.n());
184 assert_eq!(b.n(), res.n());
185 assert_ne!(a.as_ptr(), b.as_ptr());
186 }
187 unsafe {
188 vec_znx::vec_znx_add(
189 module.ptr(),
190 res.at_mut_ptr(res_col, 0),
191 res.size() as u64,
192 res.sl() as u64,
193 a.at_ptr(a_col, 0),
194 a.size() as u64,
195 a.sl() as u64,
196 b.at_ptr(b_col, 0),
197 b.size() as u64,
198 b.sl() as u64,
199 )
200 }
201 }
202}
203
204unsafe impl VecZnxBigAddInplaceImpl<Self> for FFT64 {
205 /// Adds `a` to `b` and stores the result on `b`.
206 fn vec_znx_big_add_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
207 where
208 R: VecZnxBigToMut<Self>,
209 A: VecZnxBigToRef<Self>,
210 {
211 let a: VecZnxBig<&[u8], Self> = a.to_ref();
212 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
213
214 #[cfg(debug_assertions)]
215 {
216 assert_eq!(a.n(), res.n());
217 }
218 unsafe {
219 vec_znx::vec_znx_add(
220 module.ptr(),
221 res.at_mut_ptr(res_col, 0),
222 res.size() as u64,
223 res.sl() as u64,
224 a.at_ptr(a_col, 0),
225 a.size() as u64,
226 a.sl() as u64,
227 res.at_ptr(res_col, 0),
228 res.size() as u64,
229 res.sl() as u64,
230 )
231 }
232 }
233}
234
235unsafe impl VecZnxBigAddSmallImpl<Self> for FFT64 {
236 /// Adds `a` to `b` and stores the result on `c`.
237 fn vec_znx_big_add_small_impl<R, A, B>(
238 module: &Module<Self>,
239 res: &mut R,
240 res_col: usize,
241 a: &A,
242 a_col: usize,
243 b: &B,
244 b_col: usize,
245 ) where
246 R: VecZnxBigToMut<Self>,
247 A: VecZnxBigToRef<Self>,
248 B: VecZnxToRef,
249 {
250 let a: VecZnxBig<&[u8], Self> = a.to_ref();
251 let b: VecZnx<&[u8]> = b.to_ref();
252 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
253
254 #[cfg(debug_assertions)]
255 {
256 assert_eq!(a.n(), res.n());
257 assert_eq!(b.n(), res.n());
258 assert_ne!(a.as_ptr(), b.as_ptr());
259 }
260 unsafe {
261 vec_znx::vec_znx_add(
262 module.ptr(),
263 res.at_mut_ptr(res_col, 0),
264 res.size() as u64,
265 res.sl() as u64,
266 a.at_ptr(a_col, 0),
267 a.size() as u64,
268 a.sl() as u64,
269 b.at_ptr(b_col, 0),
270 b.size() as u64,
271 b.sl() as u64,
272 )
273 }
274 }
275}
276
277unsafe impl VecZnxBigAddSmallInplaceImpl<Self> for FFT64 {
278 /// Adds `a` to `b` and stores the result on `b`.
279 fn vec_znx_big_add_small_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
280 where
281 R: VecZnxBigToMut<Self>,
282 A: VecZnxToRef,
283 {
284 let a: VecZnx<&[u8]> = a.to_ref();
285 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
286
287 #[cfg(debug_assertions)]
288 {
289 assert_eq!(a.n(), res.n());
290 }
291 unsafe {
292 vec_znx::vec_znx_add(
293 module.ptr(),
294 res.at_mut_ptr(res_col, 0),
295 res.size() as u64,
296 res.sl() as u64,
297 res.at_ptr(res_col, 0),
298 res.size() as u64,
299 res.sl() as u64,
300 a.at_ptr(a_col, 0),
301 a.size() as u64,
302 a.sl() as u64,
303 )
304 }
305 }
306}
307
308unsafe impl VecZnxBigSubImpl<Self> for FFT64 {
309 /// Subtracts `a` to `b` and stores the result on `c`.
310 fn vec_znx_big_sub_impl<R, A, B>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
311 where
312 R: VecZnxBigToMut<Self>,
313 A: VecZnxBigToRef<Self>,
314 B: VecZnxBigToRef<Self>,
315 {
316 let a: VecZnxBig<&[u8], Self> = a.to_ref();
317 let b: VecZnxBig<&[u8], Self> = b.to_ref();
318 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
319
320 #[cfg(debug_assertions)]
321 {
322 assert_eq!(a.n(), res.n());
323 assert_eq!(b.n(), res.n());
324 assert_ne!(a.as_ptr(), b.as_ptr());
325 }
326 unsafe {
327 vec_znx::vec_znx_sub(
328 module.ptr(),
329 res.at_mut_ptr(res_col, 0),
330 res.size() as u64,
331 res.sl() as u64,
332 a.at_ptr(a_col, 0),
333 a.size() as u64,
334 a.sl() as u64,
335 b.at_ptr(b_col, 0),
336 b.size() as u64,
337 b.sl() as u64,
338 )
339 }
340 }
341}
342
343unsafe impl VecZnxBigSubABInplaceImpl<Self> for FFT64 {
344 /// Subtracts `a` from `b` and stores the result on `b`.
345 fn vec_znx_big_sub_ab_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
346 where
347 R: VecZnxBigToMut<Self>,
348 A: VecZnxBigToRef<Self>,
349 {
350 let a: VecZnxBig<&[u8], Self> = a.to_ref();
351 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
352
353 #[cfg(debug_assertions)]
354 {
355 assert_eq!(a.n(), res.n());
356 }
357 unsafe {
358 vec_znx::vec_znx_sub(
359 module.ptr(),
360 res.at_mut_ptr(res_col, 0),
361 res.size() as u64,
362 res.sl() as u64,
363 res.at_ptr(res_col, 0),
364 res.size() as u64,
365 res.sl() as u64,
366 a.at_ptr(a_col, 0),
367 a.size() as u64,
368 a.sl() as u64,
369 )
370 }
371 }
372}
373
374unsafe impl VecZnxBigSubBAInplaceImpl<Self> for FFT64 {
375 /// Subtracts `b` from `a` and stores the result on `b`.
376 fn vec_znx_big_sub_ba_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
377 where
378 R: VecZnxBigToMut<Self>,
379 A: VecZnxBigToRef<Self>,
380 {
381 let a: VecZnxBig<&[u8], Self> = a.to_ref();
382 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
383
384 #[cfg(debug_assertions)]
385 {
386 assert_eq!(a.n(), res.n());
387 }
388 unsafe {
389 vec_znx::vec_znx_sub(
390 module.ptr(),
391 res.at_mut_ptr(res_col, 0),
392 res.size() as u64,
393 res.sl() as u64,
394 a.at_ptr(a_col, 0),
395 a.size() as u64,
396 a.sl() as u64,
397 res.at_ptr(res_col, 0),
398 res.size() as u64,
399 res.sl() as u64,
400 )
401 }
402 }
403}
404
405unsafe impl VecZnxBigSubSmallAImpl<Self> for FFT64 {
406 /// Subtracts `b` from `a` and stores the result on `c`.
407 fn vec_znx_big_sub_small_a_impl<R, A, B>(
408 module: &Module<Self>,
409 res: &mut R,
410 res_col: usize,
411 a: &A,
412 a_col: usize,
413 b: &B,
414 b_col: usize,
415 ) where
416 R: VecZnxBigToMut<Self>,
417 A: VecZnxToRef,
418 B: VecZnxBigToRef<Self>,
419 {
420 let a: VecZnx<&[u8]> = a.to_ref();
421 let b: VecZnxBig<&[u8], Self> = b.to_ref();
422 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
423
424 #[cfg(debug_assertions)]
425 {
426 assert_eq!(a.n(), res.n());
427 assert_eq!(b.n(), res.n());
428 assert_ne!(a.as_ptr(), b.as_ptr());
429 }
430 unsafe {
431 vec_znx::vec_znx_sub(
432 module.ptr(),
433 res.at_mut_ptr(res_col, 0),
434 res.size() as u64,
435 res.sl() as u64,
436 a.at_ptr(a_col, 0),
437 a.size() as u64,
438 a.sl() as u64,
439 b.at_ptr(b_col, 0),
440 b.size() as u64,
441 b.sl() as u64,
442 )
443 }
444 }
445}
446
447unsafe impl VecZnxBigSubSmallAInplaceImpl<Self> for FFT64 {
448 /// Subtracts `a` from `res` and stores the result on `res`.
449 fn vec_znx_big_sub_small_a_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
450 where
451 R: VecZnxBigToMut<Self>,
452 A: VecZnxToRef,
453 {
454 let a: VecZnx<&[u8]> = a.to_ref();
455 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
456
457 #[cfg(debug_assertions)]
458 {
459 assert_eq!(a.n(), res.n());
460 }
461 unsafe {
462 vec_znx::vec_znx_sub(
463 module.ptr(),
464 res.at_mut_ptr(res_col, 0),
465 res.size() as u64,
466 res.sl() as u64,
467 res.at_ptr(res_col, 0),
468 res.size() as u64,
469 res.sl() as u64,
470 a.at_ptr(a_col, 0),
471 a.size() as u64,
472 a.sl() as u64,
473 )
474 }
475 }
476}
477
478unsafe impl VecZnxBigSubSmallBImpl<Self> for FFT64 {
479 /// Subtracts `b` from `a` and stores the result on `c`.
480 fn vec_znx_big_sub_small_b_impl<R, A, B>(
481 module: &Module<Self>,
482 res: &mut R,
483 res_col: usize,
484 a: &A,
485 a_col: usize,
486 b: &B,
487 b_col: usize,
488 ) where
489 R: VecZnxBigToMut<Self>,
490 A: VecZnxBigToRef<Self>,
491 B: VecZnxToRef,
492 {
493 let a: VecZnxBig<&[u8], Self> = a.to_ref();
494 let b: VecZnx<&[u8]> = b.to_ref();
495 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
496
497 #[cfg(debug_assertions)]
498 {
499 assert_eq!(a.n(), res.n());
500 assert_eq!(b.n(), res.n());
501 assert_ne!(a.as_ptr(), b.as_ptr());
502 }
503 unsafe {
504 vec_znx::vec_znx_sub(
505 module.ptr(),
506 res.at_mut_ptr(res_col, 0),
507 res.size() as u64,
508 res.sl() as u64,
509 a.at_ptr(a_col, 0),
510 a.size() as u64,
511 a.sl() as u64,
512 b.at_ptr(b_col, 0),
513 b.size() as u64,
514 b.sl() as u64,
515 )
516 }
517 }
518}
519
520unsafe impl VecZnxBigSubSmallBInplaceImpl<Self> for FFT64 {
521 /// Subtracts `res` from `a` and stores the result on `res`.
522 fn vec_znx_big_sub_small_b_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
523 where
524 R: VecZnxBigToMut<Self>,
525 A: VecZnxToRef,
526 {
527 let a: VecZnx<&[u8]> = a.to_ref();
528 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
529
530 #[cfg(debug_assertions)]
531 {
532 assert_eq!(a.n(), res.n());
533 }
534 unsafe {
535 vec_znx::vec_znx_sub(
536 module.ptr(),
537 res.at_mut_ptr(res_col, 0),
538 res.size() as u64,
539 res.sl() as u64,
540 a.at_ptr(a_col, 0),
541 a.size() as u64,
542 a.sl() as u64,
543 res.at_ptr(res_col, 0),
544 res.size() as u64,
545 res.sl() as u64,
546 )
547 }
548 }
549}
550
551unsafe impl VecZnxBigNegateInplaceImpl<Self> for FFT64 {
552 fn vec_znx_big_negate_inplace_impl<A>(module: &Module<Self>, a: &mut A, a_col: usize)
553 where
554 A: VecZnxBigToMut<Self>,
555 {
556 let mut a: VecZnxBig<&mut [u8], Self> = a.to_mut();
557 unsafe {
558 vec_znx::vec_znx_negate(
559 module.ptr(),
560 a.at_mut_ptr(a_col, 0),
561 a.size() as u64,
562 a.sl() as u64,
563 a.at_ptr(a_col, 0),
564 a.size() as u64,
565 a.sl() as u64,
566 )
567 }
568 }
569}
570
571unsafe impl VecZnxBigNormalizeTmpBytesImpl<Self> for FFT64 {
572 fn vec_znx_big_normalize_tmp_bytes_impl(module: &Module<Self>) -> usize {
573 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(module.ptr()) as usize }
574 }
575}
576
577unsafe impl VecZnxBigNormalizeImpl<Self> for FFT64
578where
579 Self: TakeSliceImpl<Self>,
580{
581 fn vec_znx_big_normalize_impl<R, A>(
582 module: &Module<Self>,
583 basek: usize,
584 res: &mut R,
585 res_col: usize,
586 a: &A,
587 a_col: usize,
588 scratch: &mut Scratch<Self>,
589 ) where
590 R: VecZnxToMut,
591 A: VecZnxBigToRef<Self>,
592 {
593 let a: VecZnxBig<&[u8], Self> = a.to_ref();
594 let mut res: VecZnx<&mut [u8]> = res.to_mut();
595
596 #[cfg(debug_assertions)]
597 {
598 assert_eq!(res.n(), a.n());
599 }
600
601 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_big_normalize_tmp_bytes());
602 unsafe {
603 vec_znx::vec_znx_normalize_base2k(
604 module.ptr(),
605 basek as u64,
606 res.at_mut_ptr(res_col, 0),
607 res.size() as u64,
608 res.sl() as u64,
609 a.at_ptr(a_col, 0),
610 a.size() as u64,
611 a.sl() as u64,
612 tmp_bytes.as_mut_ptr(),
613 );
614 }
615 }
616}
617
618unsafe impl VecZnxBigAutomorphismImpl<Self> for FFT64 {
619 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
620 fn vec_znx_big_automorphism_impl<R, A>(module: &Module<Self>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
621 where
622 R: VecZnxBigToMut<Self>,
623 A: VecZnxBigToRef<Self>,
624 {
625 let a: VecZnxBig<&[u8], Self> = a.to_ref();
626 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
627
628 #[cfg(debug_assertions)]
629 {
630 assert_eq!(a.n(), res.n());
631 }
632 unsafe {
633 vec_znx::vec_znx_automorphism(
634 module.ptr(),
635 k,
636 res.at_mut_ptr(res_col, 0),
637 res.size() as u64,
638 res.sl() as u64,
639 a.at_ptr(a_col, 0),
640 a.size() as u64,
641 a.sl() as u64,
642 )
643 }
644 }
645}
646
647unsafe impl VecZnxBigAutomorphismInplaceImpl<Self> for FFT64 {
648 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
649 fn vec_znx_big_automorphism_inplace_impl<A>(module: &Module<Self>, k: i64, a: &mut A, a_col: usize)
650 where
651 A: VecZnxBigToMut<Self>,
652 {
653 let mut a: VecZnxBig<&mut [u8], Self> = a.to_mut();
654 unsafe {
655 vec_znx::vec_znx_automorphism(
656 module.ptr(),
657 k,
658 a.at_mut_ptr(a_col, 0),
659 a.size() as u64,
660 a.sl() as u64,
661 a.at_ptr(a_col, 0),
662 a.size() as u64,
663 a.sl() as u64,
664 )
665 }
666 }
667}