poulpy/base2k/src/vec_znx_big_ops.rs
base2k/src/vec_znx_big_ops.rsRaw
1use crate::ffi::vec_znx;
2use crate::znx_base::{ZnxInfos, ZnxView, ZnxViewMut};
3use crate::{
4 Backend, FFT64, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxBigToMut, VecZnxBigToRef, VecZnxScratch,
5 VecZnxToMut, VecZnxToRef, ZnxSliceSize, bytes_of_vec_znx_big,
6};
7
8pub trait VecZnxBigAlloc<B: Backend> {
9 /// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
10 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBigOwned<B>;
11
12 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
13 ///
14 /// Behavior: takes ownership of the backing array.
15 ///
16 /// # Arguments
17 ///
18 /// * `cols`: the number of polynomials..
19 /// * `size`: the number of polynomials per column.
20 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
21 ///
22 /// # Panics
23 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
24 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<B>;
25
26 // /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
27 // ///
28 // /// Behavior: the backing array is only borrowed.
29 // ///
30 // /// # Arguments
31 // ///
32 // /// * `cols`: the number of polynomials..
33 // /// * `size`: the number of polynomials per column.
34 // /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
35 // ///
36 // /// # Panics
37 // /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
38 // fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<B>;
39
40 /// Returns the minimum number of bytes necessary to allocate
41 /// a new [VecZnxBig] through [VecZnxBig::from_bytes].
42 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize;
43}
44
45pub trait VecZnxBigOps<BACKEND: Backend> {
46 /// Adds `a` to `b` and stores the result on `c`.
47 fn vec_znx_big_add<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
48 where
49 R: VecZnxBigToMut<BACKEND>,
50 A: VecZnxBigToRef<BACKEND>,
51 B: VecZnxBigToRef<BACKEND>;
52
53 /// Adds `a` to `b` and stores the result on `b`.
54 fn vec_znx_big_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
55 where
56 R: VecZnxBigToMut<BACKEND>,
57 A: VecZnxBigToRef<BACKEND>;
58
59 /// Adds `a` to `b` and stores the result on `c`.
60 fn vec_znx_big_add_small<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
61 where
62 R: VecZnxBigToMut<BACKEND>,
63 A: VecZnxBigToRef<BACKEND>,
64 B: VecZnxToRef;
65
66 /// Adds `a` to `b` and stores the result on `b`.
67 fn vec_znx_big_add_small_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
68 where
69 R: VecZnxBigToMut<BACKEND>,
70 A: VecZnxToRef;
71
72 /// Subtracts `a` to `b` and stores the result on `c`.
73 fn vec_znx_big_sub<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
74 where
75 R: VecZnxBigToMut<BACKEND>,
76 A: VecZnxBigToRef<BACKEND>,
77 B: VecZnxBigToRef<BACKEND>;
78
79 /// Subtracts `a` from `b` and stores the result on `b`.
80 fn vec_znx_big_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
81 where
82 R: VecZnxBigToMut<BACKEND>,
83 A: VecZnxBigToRef<BACKEND>;
84
85 /// Subtracts `b` from `a` and stores the result on `b`.
86 fn vec_znx_big_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
87 where
88 R: VecZnxBigToMut<BACKEND>,
89 A: VecZnxBigToRef<BACKEND>;
90
91 /// Subtracts `b` from `a` and stores the result on `c`.
92 fn vec_znx_big_sub_small_a<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
93 where
94 R: VecZnxBigToMut<BACKEND>,
95 A: VecZnxToRef,
96 B: VecZnxBigToRef<BACKEND>;
97
98 /// Subtracts `a` from `res` and stores the result on `res`.
99 fn vec_znx_big_sub_small_a_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
100 where
101 R: VecZnxBigToMut<BACKEND>,
102 A: VecZnxToRef;
103
104 /// Subtracts `b` from `a` and stores the result on `c`.
105 fn vec_znx_big_sub_small_b<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
106 where
107 R: VecZnxBigToMut<BACKEND>,
108 A: VecZnxBigToRef<BACKEND>,
109 B: VecZnxToRef;
110
111 /// Subtracts `res` from `a` and stores the result on `res`.
112 fn vec_znx_big_sub_small_b_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
113 where
114 R: VecZnxBigToMut<BACKEND>,
115 A: VecZnxToRef;
116
117 /// Normalizes `a` and stores the result on `b`.
118 ///
119 /// # Arguments
120 ///
121 /// * `log_base2k`: normalization basis.
122 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
123 fn vec_znx_big_normalize<R, A>(
124 &self,
125 log_base2k: usize,
126 res: &mut R,
127 res_col: usize,
128 a: &A,
129 a_col: usize,
130 scratch: &mut Scratch,
131 ) where
132 R: VecZnxToMut,
133 A: VecZnxBigToRef<BACKEND>;
134
135 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
136 fn vec_znx_big_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
137 where
138 R: VecZnxBigToMut<BACKEND>,
139 A: VecZnxBigToRef<BACKEND>;
140
141 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
142 fn vec_znx_big_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
143 where
144 A: VecZnxBigToMut<BACKEND>;
145}
146
147pub trait VecZnxBigScratch {
148 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
149 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
150}
151
152impl<B: Backend> VecZnxBigAlloc<B> for Module<B> {
153 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBigOwned<B> {
154 VecZnxBig::new(self, cols, size)
155 }
156
157 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<B> {
158 VecZnxBig::new_from_bytes(self, cols, size, bytes)
159 }
160
161 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
162 bytes_of_vec_znx_big(self, cols, size)
163 }
164}
165
166impl VecZnxBigOps<FFT64> for Module<FFT64> {
167 fn vec_znx_big_add<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
168 where
169 R: VecZnxBigToMut<FFT64>,
170 A: VecZnxBigToRef<FFT64>,
171 B: VecZnxBigToRef<FFT64>,
172 {
173 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
174 let b: VecZnxBig<&[u8], FFT64> = b.to_ref();
175 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
176
177 #[cfg(debug_assertions)]
178 {
179 assert_eq!(a.n(), self.n());
180 assert_eq!(b.n(), self.n());
181 assert_eq!(res.n(), self.n());
182 assert_ne!(a.as_ptr(), b.as_ptr());
183 }
184 unsafe {
185 vec_znx::vec_znx_add(
186 self.ptr,
187 res.at_mut_ptr(res_col, 0),
188 res.size() as u64,
189 res.sl() as u64,
190 a.at_ptr(a_col, 0),
191 a.size() as u64,
192 a.sl() as u64,
193 b.at_ptr(b_col, 0),
194 b.size() as u64,
195 b.sl() as u64,
196 )
197 }
198 }
199
200 fn vec_znx_big_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
201 where
202 R: VecZnxBigToMut<FFT64>,
203 A: VecZnxBigToRef<FFT64>,
204 {
205 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
206 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
207
208 #[cfg(debug_assertions)]
209 {
210 assert_eq!(a.n(), self.n());
211 assert_eq!(res.n(), self.n());
212 }
213 unsafe {
214 vec_znx::vec_znx_add(
215 self.ptr,
216 res.at_mut_ptr(res_col, 0),
217 res.size() as u64,
218 res.sl() as u64,
219 a.at_ptr(a_col, 0),
220 a.size() as u64,
221 a.sl() as u64,
222 res.at_ptr(res_col, 0),
223 res.size() as u64,
224 res.sl() as u64,
225 )
226 }
227 }
228
229 fn vec_znx_big_sub<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
230 where
231 R: VecZnxBigToMut<FFT64>,
232 A: VecZnxBigToRef<FFT64>,
233 B: VecZnxBigToRef<FFT64>,
234 {
235 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
236 let b: VecZnxBig<&[u8], FFT64> = b.to_ref();
237 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
238
239 #[cfg(debug_assertions)]
240 {
241 assert_eq!(a.n(), self.n());
242 assert_eq!(b.n(), self.n());
243 assert_eq!(res.n(), self.n());
244 assert_ne!(a.as_ptr(), b.as_ptr());
245 }
246 unsafe {
247 vec_znx::vec_znx_sub(
248 self.ptr,
249 res.at_mut_ptr(res_col, 0),
250 res.size() as u64,
251 res.sl() as u64,
252 a.at_ptr(a_col, 0),
253 a.size() as u64,
254 a.sl() as u64,
255 b.at_ptr(b_col, 0),
256 b.size() as u64,
257 b.sl() as u64,
258 )
259 }
260 }
261
262 fn vec_znx_big_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
263 where
264 R: VecZnxBigToMut<FFT64>,
265 A: VecZnxBigToRef<FFT64>,
266 {
267 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
268 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
269
270 #[cfg(debug_assertions)]
271 {
272 assert_eq!(a.n(), self.n());
273 assert_eq!(res.n(), self.n());
274 }
275 unsafe {
276 vec_znx::vec_znx_sub(
277 self.ptr,
278 res.at_mut_ptr(res_col, 0),
279 res.size() as u64,
280 res.sl() as u64,
281 res.at_ptr(res_col, 0),
282 res.size() as u64,
283 res.sl() as u64,
284 a.at_ptr(a_col, 0),
285 a.size() as u64,
286 a.sl() as u64,
287 )
288 }
289 }
290
291 fn vec_znx_big_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
292 where
293 R: VecZnxBigToMut<FFT64>,
294 A: VecZnxBigToRef<FFT64>,
295 {
296 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
297 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
298
299 #[cfg(debug_assertions)]
300 {
301 assert_eq!(a.n(), self.n());
302 assert_eq!(res.n(), self.n());
303 }
304 unsafe {
305 vec_znx::vec_znx_sub(
306 self.ptr,
307 res.at_mut_ptr(res_col, 0),
308 res.size() as u64,
309 res.sl() as u64,
310 a.at_ptr(a_col, 0),
311 a.size() as u64,
312 a.sl() as u64,
313 res.at_ptr(res_col, 0),
314 res.size() as u64,
315 res.sl() as u64,
316 )
317 }
318 }
319
320 fn vec_znx_big_sub_small_b<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
321 where
322 R: VecZnxBigToMut<FFT64>,
323 A: VecZnxBigToRef<FFT64>,
324 B: VecZnxToRef,
325 {
326 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
327 let b: VecZnx<&[u8]> = b.to_ref();
328 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
329
330 #[cfg(debug_assertions)]
331 {
332 assert_eq!(a.n(), self.n());
333 assert_eq!(b.n(), self.n());
334 assert_eq!(res.n(), self.n());
335 assert_ne!(a.as_ptr(), b.as_ptr());
336 }
337 unsafe {
338 vec_znx::vec_znx_sub(
339 self.ptr,
340 res.at_mut_ptr(res_col, 0),
341 res.size() as u64,
342 res.sl() as u64,
343 a.at_ptr(a_col, 0),
344 a.size() as u64,
345 a.sl() as u64,
346 b.at_ptr(b_col, 0),
347 b.size() as u64,
348 b.sl() as u64,
349 )
350 }
351 }
352
353 fn vec_znx_big_sub_small_b_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
354 where
355 R: VecZnxBigToMut<FFT64>,
356 A: VecZnxToRef,
357 {
358 let a: VecZnx<&[u8]> = a.to_ref();
359 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
360
361 #[cfg(debug_assertions)]
362 {
363 assert_eq!(a.n(), self.n());
364 assert_eq!(res.n(), self.n());
365 }
366 unsafe {
367 vec_znx::vec_znx_sub(
368 self.ptr,
369 res.at_mut_ptr(res_col, 0),
370 res.size() as u64,
371 res.sl() as u64,
372 a.at_ptr(a_col, 0),
373 a.size() as u64,
374 a.sl() as u64,
375 res.at_ptr(res_col, 0),
376 res.size() as u64,
377 res.sl() as u64,
378 )
379 }
380 }
381
382 fn vec_znx_big_sub_small_a<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
383 where
384 R: VecZnxBigToMut<FFT64>,
385 A: VecZnxToRef,
386 B: VecZnxBigToRef<FFT64>,
387 {
388 let a: VecZnx<&[u8]> = a.to_ref();
389 let b: VecZnxBig<&[u8], FFT64> = b.to_ref();
390 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
391
392 #[cfg(debug_assertions)]
393 {
394 assert_eq!(a.n(), self.n());
395 assert_eq!(b.n(), self.n());
396 assert_eq!(res.n(), self.n());
397 assert_ne!(a.as_ptr(), b.as_ptr());
398 }
399 unsafe {
400 vec_znx::vec_znx_sub(
401 self.ptr,
402 res.at_mut_ptr(res_col, 0),
403 res.size() as u64,
404 res.sl() as u64,
405 a.at_ptr(a_col, 0),
406 a.size() as u64,
407 a.sl() as u64,
408 b.at_ptr(b_col, 0),
409 b.size() as u64,
410 b.sl() as u64,
411 )
412 }
413 }
414
415 fn vec_znx_big_sub_small_a_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
416 where
417 R: VecZnxBigToMut<FFT64>,
418 A: VecZnxToRef,
419 {
420 let a: VecZnx<&[u8]> = a.to_ref();
421 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
422
423 #[cfg(debug_assertions)]
424 {
425 assert_eq!(a.n(), self.n());
426 assert_eq!(res.n(), self.n());
427 }
428 unsafe {
429 vec_znx::vec_znx_sub(
430 self.ptr,
431 res.at_mut_ptr(res_col, 0),
432 res.size() as u64,
433 res.sl() as u64,
434 res.at_ptr(res_col, 0),
435 res.size() as u64,
436 res.sl() as u64,
437 a.at_ptr(a_col, 0),
438 a.size() as u64,
439 a.sl() as u64,
440 )
441 }
442 }
443
444 fn vec_znx_big_add_small<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
445 where
446 R: VecZnxBigToMut<FFT64>,
447 A: VecZnxBigToRef<FFT64>,
448 B: VecZnxToRef,
449 {
450 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
451 let b: VecZnx<&[u8]> = b.to_ref();
452 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
453
454 #[cfg(debug_assertions)]
455 {
456 assert_eq!(a.n(), self.n());
457 assert_eq!(b.n(), self.n());
458 assert_eq!(res.n(), self.n());
459 assert_ne!(a.as_ptr(), b.as_ptr());
460 }
461 unsafe {
462 vec_znx::vec_znx_add(
463 self.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 b.at_ptr(b_col, 0),
471 b.size() as u64,
472 b.sl() as u64,
473 )
474 }
475 }
476
477 fn vec_znx_big_add_small_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
478 where
479 R: VecZnxBigToMut<FFT64>,
480 A: VecZnxToRef,
481 {
482 let a: VecZnx<&[u8]> = a.to_ref();
483 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
484
485 #[cfg(debug_assertions)]
486 {
487 assert_eq!(a.n(), self.n());
488 assert_eq!(res.n(), self.n());
489 }
490 unsafe {
491 vec_znx::vec_znx_add(
492 self.ptr,
493 res.at_mut_ptr(res_col, 0),
494 res.size() as u64,
495 res.sl() as u64,
496 res.at_ptr(res_col, 0),
497 res.size() as u64,
498 res.sl() as u64,
499 a.at_ptr(a_col, 0),
500 a.size() as u64,
501 a.sl() as u64,
502 )
503 }
504 }
505
506 fn vec_znx_big_normalize<R, A>(
507 &self,
508 log_base2k: usize,
509 res: &mut R,
510 res_col: usize,
511 a: &A,
512 a_col: usize,
513 scratch: &mut Scratch,
514 ) where
515 R: VecZnxToMut,
516 A: VecZnxBigToRef<FFT64>,
517 {
518 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
519 let mut res: VecZnx<&mut [u8]> = res.to_mut();
520
521 #[cfg(debug_assertions)]
522 {
523 assert_eq!(a.n(), self.n());
524 assert_eq!(res.n(), self.n());
525 //(Jay)Note: This is calling VezZnxOps::vec_znx_normalize_tmp_bytes and not VecZnxBigOps::vec_znx_big_normalize_tmp_bytes.
526 // In the FFT backend the tmp sizes are same but will be different in the NTT backend
527 // assert!(tmp_bytes.len() >= <Self as VecZnxOps<&mut [u8], & [u8]>>::vec_znx_normalize_tmp_bytes(&self));
528 // assert_alignement(tmp_bytes.as_ptr());
529 }
530
531 let (tmp_bytes, _) = scratch.tmp_slice(<Self as VecZnxBigScratch>::vec_znx_big_normalize_tmp_bytes(
532 &self,
533 ));
534 unsafe {
535 vec_znx::vec_znx_normalize_base2k(
536 self.ptr,
537 log_base2k as u64,
538 res.at_mut_ptr(res_col, 0),
539 res.size() as u64,
540 res.sl() as u64,
541 a.at_ptr(a_col, 0),
542 a.size() as u64,
543 a.sl() as u64,
544 tmp_bytes.as_mut_ptr(),
545 );
546 }
547 }
548
549 fn vec_znx_big_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
550 where
551 R: VecZnxBigToMut<FFT64>,
552 A: VecZnxBigToRef<FFT64>,
553 {
554 let a: VecZnxBig<&[u8], FFT64> = a.to_ref();
555 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
556
557 #[cfg(debug_assertions)]
558 {
559 assert_eq!(a.n(), self.n());
560 assert_eq!(res.n(), self.n());
561 }
562 unsafe {
563 vec_znx::vec_znx_automorphism(
564 self.ptr,
565 k,
566 res.at_mut_ptr(res_col, 0),
567 res.size() as u64,
568 res.sl() as u64,
569 a.at_ptr(a_col, 0),
570 a.size() as u64,
571 a.sl() as u64,
572 )
573 }
574 }
575
576 fn vec_znx_big_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
577 where
578 A: VecZnxBigToMut<FFT64>,
579 {
580 let mut a: VecZnxBig<&mut [u8], FFT64> = a.to_mut();
581
582 #[cfg(debug_assertions)]
583 {
584 assert_eq!(a.n(), self.n());
585 }
586 unsafe {
587 vec_znx::vec_znx_automorphism(
588 self.ptr,
589 k,
590 a.at_mut_ptr(a_col, 0),
591 a.size() as u64,
592 a.sl() as u64,
593 a.at_ptr(a_col, 0),
594 a.size() as u64,
595 a.sl() as u64,
596 )
597 }
598 }
599}
600
601impl<B: Backend> VecZnxBigScratch for Module<B> {
602 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
603 <Self as VecZnxScratch>::vec_znx_normalize_tmp_bytes(self)
604 }
605}