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