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, VecZnxBigSubImpl, VecZnxBigSubInplaceImpl, VecZnxBigSubNegateInplaceImpl,
14 VecZnxBigSubSmallAImpl, VecZnxBigSubSmallBImpl, VecZnxBigSubSmallInplaceImpl, VecZnxBigSubSmallNegateInplaceImpl,
15 },
16 reference::{
17 fft64::vec_znx_big::vec_znx_big_normalize,
18 vec_znx::{vec_znx_add_normal_ref, vec_znx_normalize_tmp_bytes},
19 znx::{znx_copy_ref, znx_zero_ref},
20 },
21 source::Source,
22};
23
24unsafe impl VecZnxBigAllocBytesImpl<Self> for FFT64Spqlios {
25 fn vec_znx_big_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
26 Self::layout_big_word_count() * n * cols * size * size_of::<f64>()
27 }
28}
29
30unsafe impl VecZnxBigAllocImpl<Self> for FFT64Spqlios {
31 fn vec_znx_big_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxBigOwned<Self> {
32 VecZnxBig::alloc(n, cols, size)
33 }
34}
35
36unsafe impl VecZnxBigFromBytesImpl<Self> for FFT64Spqlios {
37 fn vec_znx_big_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<Self> {
38 VecZnxBig::from_bytes(n, cols, size, bytes)
39 }
40}
41
42unsafe impl VecZnxBigFromSmallImpl<Self> for FFT64Spqlios {
43 fn vec_znx_big_from_small_impl<R, A>(res: &mut R, res_col: usize, a: &A, a_col: usize)
44 where
45 R: VecZnxBigToMut<Self>,
46 A: VecZnxToRef,
47 {
48 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
49 let a: VecZnx<&[u8]> = a.to_ref();
50
51 #[cfg(debug_assertions)]
52 {
53 assert_eq!(res.n(), a.n());
54 }
55
56 let res_size: usize = res.size();
57 let a_size: usize = a.size();
58
59 let min_size: usize = res_size.min(a_size);
60
61 for j in 0..min_size {
62 znx_copy_ref(res.at_mut(res_col, j), a.at(a_col, j));
63 }
64
65 for j in min_size..res_size {
66 znx_zero_ref(res.at_mut(res_col, j));
67 }
68 }
69}
70
71unsafe impl VecZnxBigAddNormalImpl<Self> for FFT64Spqlios {
72 fn add_normal_impl<R: VecZnxBigToMut<Self>>(
73 _module: &Module<Self>,
74 base2k: usize,
75 res: &mut R,
76 res_col: usize,
77 k: usize,
78 source: &mut Source,
79 sigma: f64,
80 bound: f64,
81 ) {
82 let res: VecZnxBig<&mut [u8], FFT64Spqlios> = res.to_mut();
83
84 let mut res_znx: VecZnx<&mut [u8]> = VecZnx {
85 data: res.data,
86 n: res.n,
87 cols: res.cols,
88 size: res.size,
89 max_size: res.max_size,
90 };
91
92 vec_znx_add_normal_ref(base2k, &mut res_znx, res_col, k, sigma, bound, source);
93 }
94}
95
96unsafe impl VecZnxBigAddImpl<Self> for FFT64Spqlios {
97 /// Adds `a` to `b` and stores the result on `c`.
98 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)
99 where
100 R: VecZnxBigToMut<Self>,
101 A: VecZnxBigToRef<Self>,
102 B: VecZnxBigToRef<Self>,
103 {
104 let a: VecZnxBig<&[u8], Self> = a.to_ref();
105 let b: VecZnxBig<&[u8], Self> = b.to_ref();
106 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
107
108 #[cfg(debug_assertions)]
109 {
110 assert_eq!(a.n(), res.n());
111 assert_eq!(b.n(), res.n());
112 assert_ne!(a.as_ptr(), b.as_ptr());
113 }
114 unsafe {
115 vec_znx::vec_znx_add(
116 module.ptr(),
117 res.at_mut_ptr(res_col, 0),
118 res.size() as u64,
119 res.sl() as u64,
120 a.at_ptr(a_col, 0),
121 a.size() as u64,
122 a.sl() as u64,
123 b.at_ptr(b_col, 0),
124 b.size() as u64,
125 b.sl() as u64,
126 )
127 }
128 }
129}
130
131unsafe impl VecZnxBigAddInplaceImpl<Self> for FFT64Spqlios {
132 /// Adds `a` to `b` and stores the result on `b`.
133 fn vec_znx_big_add_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
134 where
135 R: VecZnxBigToMut<Self>,
136 A: VecZnxBigToRef<Self>,
137 {
138 let a: VecZnxBig<&[u8], Self> = a.to_ref();
139 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
140
141 #[cfg(debug_assertions)]
142 {
143 assert_eq!(a.n(), res.n());
144 }
145 unsafe {
146 vec_znx::vec_znx_add(
147 module.ptr(),
148 res.at_mut_ptr(res_col, 0),
149 res.size() as u64,
150 res.sl() as u64,
151 a.at_ptr(a_col, 0),
152 a.size() as u64,
153 a.sl() as u64,
154 res.at_ptr(res_col, 0),
155 res.size() as u64,
156 res.sl() as u64,
157 )
158 }
159 }
160}
161
162unsafe impl VecZnxBigAddSmallImpl<Self> for FFT64Spqlios {
163 /// Adds `a` to `b` and stores the result on `c`.
164 fn vec_znx_big_add_small_impl<R, A, B>(
165 module: &Module<Self>,
166 res: &mut R,
167 res_col: usize,
168 a: &A,
169 a_col: usize,
170 b: &B,
171 b_col: usize,
172 ) where
173 R: VecZnxBigToMut<Self>,
174 A: VecZnxBigToRef<Self>,
175 B: VecZnxToRef,
176 {
177 let a: VecZnxBig<&[u8], Self> = a.to_ref();
178 let b: VecZnx<&[u8]> = 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 VecZnxBigAddSmallInplaceImpl<Self> for FFT64Spqlios {
205 /// Adds `a` to `b` and stores the result on `b`.
206 fn vec_znx_big_add_small_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: VecZnxToRef,
210 {
211 let a: VecZnx<&[u8]> = 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 res.at_ptr(res_col, 0),
225 res.size() as u64,
226 res.sl() as u64,
227 a.at_ptr(a_col, 0),
228 a.size() as u64,
229 a.sl() as u64,
230 )
231 }
232 }
233}
234
235unsafe impl VecZnxBigSubImpl<Self> for FFT64Spqlios {
236 /// Subtracts `a` to `b` and stores the result on `c`.
237 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)
238 where
239 R: VecZnxBigToMut<Self>,
240 A: VecZnxBigToRef<Self>,
241 B: VecZnxBigToRef<Self>,
242 {
243 let a: VecZnxBig<&[u8], Self> = a.to_ref();
244 let b: VecZnxBig<&[u8], Self> = b.to_ref();
245 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
246
247 #[cfg(debug_assertions)]
248 {
249 assert_eq!(a.n(), res.n());
250 assert_eq!(b.n(), res.n());
251 assert_ne!(a.as_ptr(), b.as_ptr());
252 }
253 unsafe {
254 vec_znx::vec_znx_sub(
255 module.ptr(),
256 res.at_mut_ptr(res_col, 0),
257 res.size() as u64,
258 res.sl() as u64,
259 a.at_ptr(a_col, 0),
260 a.size() as u64,
261 a.sl() as u64,
262 b.at_ptr(b_col, 0),
263 b.size() as u64,
264 b.sl() as u64,
265 )
266 }
267 }
268}
269
270unsafe impl VecZnxBigSubInplaceImpl<Self> for FFT64Spqlios {
271 /// Subtracts `a` from `b` and stores the result on `b`.
272 fn vec_znx_big_sub_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
273 where
274 R: VecZnxBigToMut<Self>,
275 A: VecZnxBigToRef<Self>,
276 {
277 let a: VecZnxBig<&[u8], Self> = a.to_ref();
278 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
279
280 #[cfg(debug_assertions)]
281 {
282 assert_eq!(a.n(), res.n());
283 }
284 unsafe {
285 vec_znx::vec_znx_sub(
286 module.ptr(),
287 res.at_mut_ptr(res_col, 0),
288 res.size() as u64,
289 res.sl() as u64,
290 res.at_ptr(res_col, 0),
291 res.size() as u64,
292 res.sl() as u64,
293 a.at_ptr(a_col, 0),
294 a.size() as u64,
295 a.sl() as u64,
296 )
297 }
298 }
299}
300
301unsafe impl VecZnxBigSubNegateInplaceImpl<Self> for FFT64Spqlios {
302 /// Subtracts `b` from `a` and stores the result on `b`.
303 fn vec_znx_big_sub_negate_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
304 where
305 R: VecZnxBigToMut<Self>,
306 A: VecZnxBigToRef<Self>,
307 {
308 let a: VecZnxBig<&[u8], Self> = a.to_ref();
309 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
310
311 #[cfg(debug_assertions)]
312 {
313 assert_eq!(a.n(), res.n());
314 }
315 unsafe {
316 vec_znx::vec_znx_sub(
317 module.ptr(),
318 res.at_mut_ptr(res_col, 0),
319 res.size() as u64,
320 res.sl() as u64,
321 a.at_ptr(a_col, 0),
322 a.size() as u64,
323 a.sl() as u64,
324 res.at_ptr(res_col, 0),
325 res.size() as u64,
326 res.sl() as u64,
327 )
328 }
329 }
330}
331
332unsafe impl VecZnxBigSubSmallAImpl<Self> for FFT64Spqlios {
333 /// Subtracts `b` from `a` and stores the result on `c`.
334 fn vec_znx_big_sub_small_a_impl<R, A, B>(
335 module: &Module<Self>,
336 res: &mut R,
337 res_col: usize,
338 a: &A,
339 a_col: usize,
340 b: &B,
341 b_col: usize,
342 ) where
343 R: VecZnxBigToMut<Self>,
344 A: VecZnxToRef,
345 B: VecZnxBigToRef<Self>,
346 {
347 let a: VecZnx<&[u8]> = a.to_ref();
348 let b: VecZnxBig<&[u8], Self> = b.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 assert_eq!(b.n(), res.n());
355 assert_ne!(a.as_ptr(), b.as_ptr());
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 a.at_ptr(a_col, 0),
364 a.size() as u64,
365 a.sl() as u64,
366 b.at_ptr(b_col, 0),
367 b.size() as u64,
368 b.sl() as u64,
369 )
370 }
371 }
372}
373
374unsafe impl VecZnxBigSubSmallInplaceImpl<Self> for FFT64Spqlios {
375 /// Subtracts `a` from `res` and stores the result on `res`.
376 fn vec_znx_big_sub_small_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: VecZnxToRef,
380 {
381 let a: VecZnx<&[u8]> = 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 res.at_ptr(res_col, 0),
395 res.size() as u64,
396 res.sl() as u64,
397 a.at_ptr(a_col, 0),
398 a.size() as u64,
399 a.sl() as u64,
400 )
401 }
402 }
403}
404
405unsafe impl VecZnxBigSubSmallBImpl<Self> for FFT64Spqlios {
406 /// Subtracts `b` from `a` and stores the result on `c`.
407 fn vec_znx_big_sub_small_b_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: VecZnxBigToRef<Self>,
418 B: VecZnxToRef,
419 {
420 let a: VecZnxBig<&[u8], Self> = a.to_ref();
421 let b: VecZnx<&[u8]> = 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 VecZnxBigSubSmallNegateInplaceImpl<Self> for FFT64Spqlios {
448 /// Subtracts `res` from `a` and stores the result on `res`.
449 fn vec_znx_big_sub_small_negate_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 a.at_ptr(a_col, 0),
468 a.size() as u64,
469 a.sl() as u64,
470 res.at_ptr(res_col, 0),
471 res.size() as u64,
472 res.sl() as u64,
473 )
474 }
475 }
476}
477
478unsafe impl VecZnxBigNegateImpl<Self> for FFT64Spqlios {
479 fn vec_znx_big_negate_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
480 where
481 R: VecZnxBigToMut<Self>,
482 A: VecZnxBigToRef<Self>,
483 {
484 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
485 let a: VecZnxBig<&[u8], Self> = a.to_ref();
486 unsafe {
487 vec_znx::vec_znx_negate(
488 module.ptr(),
489 res.at_mut_ptr(res_col, 0),
490 res.size() as u64,
491 res.sl() as u64,
492 a.at_ptr(a_col, 0),
493 a.size() as u64,
494 a.sl() as u64,
495 )
496 }
497 }
498}
499
500unsafe impl VecZnxBigNegateInplaceImpl<Self> for FFT64Spqlios {
501 fn vec_znx_big_negate_inplace_impl<A>(module: &Module<Self>, a: &mut A, a_col: usize)
502 where
503 A: VecZnxBigToMut<Self>,
504 {
505 let mut a: VecZnxBig<&mut [u8], Self> = a.to_mut();
506 unsafe {
507 vec_znx::vec_znx_negate(
508 module.ptr(),
509 a.at_mut_ptr(a_col, 0),
510 a.size() as u64,
511 a.sl() as u64,
512 a.at_ptr(a_col, 0),
513 a.size() as u64,
514 a.sl() as u64,
515 )
516 }
517 }
518}
519
520unsafe impl VecZnxBigNormalizeTmpBytesImpl<Self> for FFT64Spqlios {
521 fn vec_znx_big_normalize_tmp_bytes_impl(module: &Module<Self>) -> usize {
522 vec_znx_normalize_tmp_bytes(module.n())
523 }
524}
525
526unsafe impl VecZnxBigNormalizeImpl<Self> for FFT64Spqlios
527where
528 Self: TakeSliceImpl<Self>,
529{
530 fn vec_znx_big_normalize_impl<R, A>(
531 module: &Module<Self>,
532 res_basek: usize,
533 res: &mut R,
534 res_col: usize,
535 a_basek: usize,
536 a: &A,
537 a_col: usize,
538 scratch: &mut Scratch<Self>,
539 ) where
540 R: VecZnxToMut,
541 A: VecZnxBigToRef<Self>,
542 {
543 let (carry, _) = scratch.take_slice(module.vec_znx_big_normalize_tmp_bytes() / size_of::<i64>());
544 // unsafe {
545 // vec_znx::vec_znx_normalize_base2k(
546 // module.ptr(),
547 // base2k as u64,
548 // res.at_mut_ptr(res_col, 0),
549 // res.size() as u64,
550 // res.sl() as u64,
551 // a.at_ptr(a_col, 0),
552 // a.size() as u64,
553 // a.sl() as u64,
554 // tmp_bytes.as_mut_ptr(),
555 // );
556 // }
557 vec_znx_big_normalize(res_basek, res, res_col, a_basek, a, a_col, carry);
558 }
559}
560
561unsafe impl VecZnxBigAutomorphismImpl<Self> for FFT64Spqlios {
562 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
563 fn vec_znx_big_automorphism_impl<R, A>(module: &Module<Self>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
564 where
565 R: VecZnxBigToMut<Self>,
566 A: VecZnxBigToRef<Self>,
567 {
568 let a: VecZnxBig<&[u8], Self> = a.to_ref();
569 let mut res: VecZnxBig<&mut [u8], Self> = res.to_mut();
570
571 #[cfg(debug_assertions)]
572 {
573 assert_eq!(a.n(), res.n());
574 }
575 unsafe {
576 vec_znx::vec_znx_automorphism(
577 module.ptr(),
578 k,
579 res.at_mut_ptr(res_col, 0),
580 res.size() as u64,
581 res.sl() as u64,
582 a.at_ptr(a_col, 0),
583 a.size() as u64,
584 a.sl() as u64,
585 )
586 }
587 }
588}
589
590unsafe impl VecZnxBigAutomorphismInplaceTmpBytesImpl<Self> for FFT64Spqlios {
591 fn vec_znx_big_automorphism_inplace_tmp_bytes_impl(_module: &Module<Self>) -> usize {
592 0
593 }
594}
595
596unsafe impl VecZnxBigAutomorphismInplaceImpl<Self> for FFT64Spqlios {
597 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
598 fn vec_znx_big_automorphism_inplace_impl<A>(
599 module: &Module<Self>,
600 k: i64,
601 a: &mut A,
602 a_col: usize,
603 _scratch: &mut Scratch<Self>,
604 ) where
605 A: VecZnxBigToMut<Self>,
606 {
607 let mut a: VecZnxBig<&mut [u8], Self> = a.to_mut();
608 unsafe {
609 vec_znx::vec_znx_automorphism(
610 module.ptr(),
611 k,
612 a.at_mut_ptr(a_col, 0),
613 a.size() as u64,
614 a.sl() as u64,
615 a.at_ptr(a_col, 0),
616 a.size() as u64,
617 a.sl() as u64,
618 )
619 }
620 }
621}