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, DataView, FFT64, Module, ScratchBorr, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxScratch, assert_alignement,
5 bytes_of_vec_znx_big,
6};
7
8pub trait VecZnxBigAlloc<B> {
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<DataMut, Data, B> {
46 /// Adds `a` to `b` and stores the result on `c`.
47 fn vec_znx_big_add(
48 &self,
49 res: &mut VecZnxBig<DataMut, B>,
50 res_col: usize,
51 a: &VecZnxBig<Data, B>,
52 a_col: usize,
53 b: &VecZnxBig<Data, B>,
54 b_col: usize,
55 );
56
57 /// Adds `a` to `b` and stores the result on `b`.
58 fn vec_znx_big_add_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnxBig<Data, B>, a_col: usize);
59
60 /// Adds `a` to `b` and stores the result on `c`.
61 fn vec_znx_big_add_small(
62 &self,
63 res: &mut VecZnxBig<DataMut, B>,
64 res_col: usize,
65 a: &VecZnxBig<Data, B>,
66 a_col: usize,
67 b: &VecZnx<Data>,
68 b_col: usize,
69 );
70
71 /// Adds `a` to `b` and stores the result on `b`.
72 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
73
74 /// Subtracts `a` to `b` and stores the result on `c`.
75 fn vec_znx_big_sub(
76 &self,
77 res: &mut VecZnxBig<DataMut, B>,
78 res_col: usize,
79 a: &VecZnxBig<Data, B>,
80 a_col: usize,
81 b: &VecZnxBig<Data, B>,
82 b_col: usize,
83 );
84
85 /// Subtracts `a` from `b` and stores the result on `b`.
86 fn vec_znx_big_sub_ab_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnxBig<Data, B>, a_col: usize);
87
88 /// Subtracts `b` from `a` and stores the result on `b`.
89 fn vec_znx_big_sub_ba_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnxBig<Data, B>, a_col: usize);
90
91 /// Subtracts `b` from `a` and stores the result on `c`.
92 fn vec_znx_big_sub_small_a(
93 &self,
94 res: &mut VecZnxBig<DataMut, B>,
95 res_col: usize,
96 a: &VecZnx<Data>,
97 a_col: usize,
98 b: &VecZnxBig<Data, B>,
99 b_col: usize,
100 );
101
102 /// Subtracts `a` from `res` and stores the result on `res`.
103 fn vec_znx_big_sub_small_a_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
104
105 /// Subtracts `b` from `a` and stores the result on `c`.
106 fn vec_znx_big_sub_small_b(
107 &self,
108 res: &mut VecZnxBig<DataMut, B>,
109 res_col: usize,
110 a: &VecZnxBig<Data, B>,
111 a_col: usize,
112 b: &VecZnx<Data>,
113 b_col: usize,
114 );
115
116 /// Subtracts `res` from `a` and stores the result on `res`.
117 fn vec_znx_big_sub_small_b_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
118
119 /// Normalizes `a` and stores the result on `b`.
120 ///
121 /// # Arguments
122 ///
123 /// * `log_base2k`: normalization basis.
124 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
125 fn vec_znx_big_normalize(
126 &self,
127 log_base2k: usize,
128 res: &mut VecZnx<DataMut>,
129 res_col: usize,
130 a: &VecZnxBig<Data, B>,
131 a_col: usize,
132 scratch: &mut ScratchBorr,
133 );
134
135 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
136 fn vec_znx_big_automorphism(
137 &self,
138 k: i64,
139 res: &mut VecZnxBig<DataMut, B>,
140 res_col: usize,
141 a: &VecZnxBig<Data, B>,
142 a_col: usize,
143 );
144
145 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
146 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<DataMut, B>, a_col: usize);
147}
148
149pub trait VecZnxBigScratch {
150 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
151 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
152}
153
154impl VecZnxBigAlloc<FFT64> for Module<FFT64> {
155 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBigOwned<FFT64> {
156 VecZnxBig::new(self, cols, size)
157 }
158
159 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<FFT64> {
160 VecZnxBig::new_from_bytes(self, cols, size, bytes)
161 }
162
163 // fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
164 // VecZnxBig::from_bytes_borrow(self, 1, cols, size, tmp_bytes)
165 // }
166
167 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
168 bytes_of_vec_znx_big(self, cols, size)
169 }
170}
171
172impl<DataMut, Data> VecZnxBigOps<DataMut, Data, FFT64> for Module<FFT64>
173where
174 DataMut: AsMut<[u8]> + AsRef<[u8]>,
175 Data: AsRef<[u8]>,
176{
177 fn vec_znx_big_add(
178 &self,
179 res: &mut VecZnxBig<DataMut, FFT64>,
180 res_col: usize,
181 a: &VecZnxBig<Data, FFT64>,
182 a_col: usize,
183 b: &VecZnxBig<Data, FFT64>,
184 b_col: usize,
185 ) {
186 #[cfg(debug_assertions)]
187 {
188 assert_eq!(a.n(), self.n());
189 assert_eq!(b.n(), self.n());
190 assert_eq!(res.n(), self.n());
191 assert_ne!(a.as_ptr(), b.as_ptr());
192 }
193 unsafe {
194 vec_znx::vec_znx_add(
195 self.ptr,
196 res.at_mut_ptr(res_col, 0),
197 res.size() as u64,
198 res.sl() as u64,
199 a.at_ptr(a_col, 0),
200 a.size() as u64,
201 a.sl() as u64,
202 b.at_ptr(b_col, 0),
203 b.size() as u64,
204 b.sl() as u64,
205 )
206 }
207 }
208
209 fn vec_znx_big_add_inplace(
210 &self,
211 res: &mut VecZnxBig<DataMut, FFT64>,
212 res_col: usize,
213 a: &VecZnxBig<Data, FFT64>,
214 a_col: usize,
215 ) {
216 #[cfg(debug_assertions)]
217 {
218 assert_eq!(a.n(), self.n());
219 assert_eq!(res.n(), self.n());
220 }
221 unsafe {
222 vec_znx::vec_znx_add(
223 self.ptr,
224 res.at_mut_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 res.at_ptr(res_col, 0),
231 res.size() as u64,
232 res.sl() as u64,
233 )
234 }
235 }
236
237 fn vec_znx_big_sub(
238 &self,
239 res: &mut VecZnxBig<DataMut, FFT64>,
240 res_col: usize,
241 a: &VecZnxBig<Data, FFT64>,
242 a_col: usize,
243 b: &VecZnxBig<Data, FFT64>,
244 b_col: usize,
245 ) {
246 #[cfg(debug_assertions)]
247 {
248 assert_eq!(a.n(), self.n());
249 assert_eq!(b.n(), self.n());
250 assert_eq!(res.n(), self.n());
251 assert_ne!(a.as_ptr(), b.as_ptr());
252 }
253 unsafe {
254 vec_znx::vec_znx_sub(
255 self.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 fn vec_znx_big_sub_ab_inplace(
270 &self,
271 res: &mut VecZnxBig<DataMut, FFT64>,
272 res_col: usize,
273 a: &VecZnxBig<Data, FFT64>,
274 a_col: usize,
275 ) {
276 #[cfg(debug_assertions)]
277 {
278 assert_eq!(a.n(), self.n());
279 assert_eq!(res.n(), self.n());
280 }
281 unsafe {
282 vec_znx::vec_znx_sub(
283 self.ptr,
284 res.at_mut_ptr(res_col, 0),
285 res.size() as u64,
286 res.sl() as u64,
287 res.at_ptr(res_col, 0),
288 res.size() as u64,
289 res.sl() as u64,
290 a.at_ptr(a_col, 0),
291 a.size() as u64,
292 a.sl() as u64,
293 )
294 }
295 }
296
297 fn vec_znx_big_sub_ba_inplace(
298 &self,
299 res: &mut VecZnxBig<DataMut, FFT64>,
300 res_col: usize,
301 a: &VecZnxBig<Data, FFT64>,
302 a_col: usize,
303 ) {
304 #[cfg(debug_assertions)]
305 {
306 assert_eq!(a.n(), self.n());
307 assert_eq!(res.n(), self.n());
308 }
309 unsafe {
310 vec_znx::vec_znx_sub(
311 self.ptr,
312 res.at_mut_ptr(res_col, 0),
313 res.size() as u64,
314 res.sl() as u64,
315 a.at_ptr(a_col, 0),
316 a.size() as u64,
317 a.sl() as u64,
318 res.at_ptr(res_col, 0),
319 res.size() as u64,
320 res.sl() as u64,
321 )
322 }
323 }
324
325 fn vec_znx_big_sub_small_b(
326 &self,
327 res: &mut VecZnxBig<DataMut, FFT64>,
328 res_col: usize,
329 a: &VecZnxBig<Data, FFT64>,
330 a_col: usize,
331 b: &VecZnx<Data>,
332 b_col: usize,
333 ) {
334 #[cfg(debug_assertions)]
335 {
336 assert_eq!(a.n(), self.n());
337 assert_eq!(b.n(), self.n());
338 assert_eq!(res.n(), self.n());
339 assert_ne!(a.as_ptr(), b.as_ptr());
340 }
341 unsafe {
342 vec_znx::vec_znx_sub(
343 self.ptr,
344 res.at_mut_ptr(res_col, 0),
345 res.size() as u64,
346 res.sl() as u64,
347 a.at_ptr(a_col, 0),
348 a.size() as u64,
349 a.sl() as u64,
350 b.at_ptr(b_col, 0),
351 b.size() as u64,
352 b.sl() as u64,
353 )
354 }
355 }
356
357 fn vec_znx_big_sub_small_b_inplace(
358 &self,
359 res: &mut VecZnxBig<DataMut, FFT64>,
360 res_col: usize,
361 a: &VecZnx<Data>,
362 a_col: usize,
363 ) {
364 #[cfg(debug_assertions)]
365 {
366 assert_eq!(a.n(), self.n());
367 assert_eq!(res.n(), self.n());
368 }
369 unsafe {
370 vec_znx::vec_znx_sub(
371 self.ptr,
372 res.at_mut_ptr(res_col, 0),
373 res.size() as u64,
374 res.sl() as u64,
375 a.at_ptr(a_col, 0),
376 a.size() as u64,
377 a.sl() as u64,
378 res.at_ptr(res_col, 0),
379 res.size() as u64,
380 res.sl() as u64,
381 )
382 }
383 }
384
385 fn vec_znx_big_sub_small_a(
386 &self,
387 res: &mut VecZnxBig<DataMut, FFT64>,
388 res_col: usize,
389 a: &VecZnx<Data>,
390 a_col: usize,
391 b: &VecZnxBig<Data, FFT64>,
392 b_col: usize,
393 ) {
394 #[cfg(debug_assertions)]
395 {
396 assert_eq!(a.n(), self.n());
397 assert_eq!(b.n(), self.n());
398 assert_eq!(res.n(), self.n());
399 assert_ne!(a.as_ptr(), b.as_ptr());
400 }
401 unsafe {
402 vec_znx::vec_znx_sub(
403 self.ptr,
404 res.at_mut_ptr(res_col, 0),
405 res.size() as u64,
406 res.sl() as u64,
407 a.at_ptr(a_col, 0),
408 a.size() as u64,
409 a.sl() as u64,
410 b.at_ptr(b_col, 0),
411 b.size() as u64,
412 b.sl() as u64,
413 )
414 }
415 }
416
417 fn vec_znx_big_sub_small_a_inplace(
418 &self,
419 res: &mut VecZnxBig<DataMut, FFT64>,
420 res_col: usize,
421 a: &VecZnx<Data>,
422 a_col: usize,
423 ) {
424 #[cfg(debug_assertions)]
425 {
426 assert_eq!(a.n(), self.n());
427 assert_eq!(res.n(), self.n());
428 }
429 unsafe {
430 vec_znx::vec_znx_sub(
431 self.ptr,
432 res.at_mut_ptr(res_col, 0),
433 res.size() as u64,
434 res.sl() as u64,
435 res.at_ptr(res_col, 0),
436 res.size() as u64,
437 res.sl() as u64,
438 a.at_ptr(a_col, 0),
439 a.size() as u64,
440 a.sl() as u64,
441 )
442 }
443 }
444
445 fn vec_znx_big_add_small(
446 &self,
447 res: &mut VecZnxBig<DataMut, FFT64>,
448 res_col: usize,
449 a: &VecZnxBig<Data, FFT64>,
450 a_col: usize,
451 b: &VecZnx<Data>,
452 b_col: usize,
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(&self, res: &mut VecZnxBig<DataMut, FFT64>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
478 #[cfg(debug_assertions)]
479 {
480 assert_eq!(a.n(), self.n());
481 assert_eq!(res.n(), self.n());
482 }
483 unsafe {
484 vec_znx::vec_znx_add(
485 self.ptr,
486 res.at_mut_ptr(res_col, 0),
487 res.size() as u64,
488 res.sl() as u64,
489 res.at_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 fn vec_znx_big_normalize(
500 &self,
501 log_base2k: usize,
502 res: &mut VecZnx<DataMut>,
503 res_col: usize,
504 a: &VecZnxBig<Data, FFT64>,
505 a_col: usize,
506 scratch: &mut ScratchBorr,
507 ) {
508 #[cfg(debug_assertions)]
509 {
510 assert_eq!(a.n(), self.n());
511 assert_eq!(res.n(), self.n());
512 //(Jay)Note: This is calling VezZnxOps::vec_znx_normalize_tmp_bytes and not VecZnxBigOps::vec_znx_big_normalize_tmp_bytes.
513 // In the FFT backend the tmp sizes are same but will be different in the NTT backend
514 // assert!(tmp_bytes.len() >= <Self as VecZnxOps<DataMut, Data>>::vec_znx_normalize_tmp_bytes(&self));
515 // assert_alignement(tmp_bytes.as_ptr());
516 }
517
518 let (tmp_bytes, _) = scratch.tmp_scalar_slice(<Self as VecZnxBigScratch>::vec_znx_big_normalize_tmp_bytes(
519 &self,
520 ));
521 unsafe {
522 vec_znx::vec_znx_normalize_base2k(
523 self.ptr,
524 log_base2k as u64,
525 res.at_mut_ptr(res_col, 0),
526 res.size() as u64,
527 res.sl() as u64,
528 a.at_ptr(a_col, 0),
529 a.size() as u64,
530 a.sl() as u64,
531 tmp_bytes.as_mut_ptr(),
532 );
533 }
534 }
535
536 fn vec_znx_big_automorphism(
537 &self,
538 k: i64,
539 res: &mut VecZnxBig<DataMut, FFT64>,
540 res_col: usize,
541 a: &VecZnxBig<Data, FFT64>,
542 a_col: usize,
543 ) {
544 #[cfg(debug_assertions)]
545 {
546 assert_eq!(a.n(), self.n());
547 assert_eq!(res.n(), self.n());
548 }
549 unsafe {
550 vec_znx::vec_znx_automorphism(
551 self.ptr,
552 k,
553 res.at_mut_ptr(res_col, 0),
554 res.size() as u64,
555 res.sl() as u64,
556 a.at_ptr(a_col, 0),
557 a.size() as u64,
558 a.sl() as u64,
559 )
560 }
561 }
562
563 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<DataMut, FFT64>, a_col: usize) {
564 #[cfg(debug_assertions)]
565 {
566 assert_eq!(a.n(), self.n());
567 }
568 unsafe {
569 vec_znx::vec_znx_automorphism(
570 self.ptr,
571 k,
572 a.at_mut_ptr(a_col, 0),
573 a.size() as u64,
574 a.sl() as u64,
575 a.at_ptr(a_col, 0),
576 a.size() as u64,
577 a.sl() as u64,
578 )
579 }
580 }
581}
582
583impl<B: Backend> VecZnxBigScratch for Module<B> {
584 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
585 <Self as VecZnxScratch>::vec_znx_normalize_tmp_bytes(self)
586 }
587}