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