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