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, FFT64, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxScratch, bytes_of_vec_znx_big};
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 /// Normalizes `a` and stores the result on `b`.
117 ///
118 /// # Arguments
119 ///
120 /// * `log_base2k`: normalization basis.
121 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
122 fn vec_znx_big_normalize(
123 &self,
124 log_base2k: usize,
125 res: &mut VecZnx<DataMut>,
126 res_col: usize,
127 a: &VecZnxBig<Data, B>,
128 a_col: usize,
129 scratch: &mut Scratch,
130 );
131
132 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
133 fn vec_znx_big_automorphism(
134 &self,
135 k: i64,
136 res: &mut VecZnxBig<DataMut, B>,
137 res_col: usize,
138 a: &VecZnxBig<Data, B>,
139 a_col: usize,
140 );
141
142 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
143 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<DataMut, B>, a_col: usize);
144}
145
146pub trait VecZnxBigScratch {
147 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
148 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
149}
150
151impl VecZnxBigAlloc<FFT64> for Module<FFT64> {
152 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBigOwned<FFT64> {
153 VecZnxBig::new(self, cols, size)
154 }
155
156 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<FFT64> {
157 VecZnxBig::new_from_bytes(self, cols, size, bytes)
158 }
159
160 // fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
161 // VecZnxBig::from_bytes_borrow(self, 1, cols, size, tmp_bytes)
162 // }
163
164 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
165 bytes_of_vec_znx_big(self, cols, size)
166 }
167}
168
169impl<DataMut, Data> VecZnxBigOps<DataMut, Data, FFT64> for Module<FFT64>
170where
171 DataMut: AsMut<[u8]> + AsRef<[u8]>,
172 Data: AsRef<[u8]>,
173{
174 fn vec_znx_big_add(
175 &self,
176 res: &mut VecZnxBig<DataMut, FFT64>,
177 res_col: usize,
178 a: &VecZnxBig<Data, FFT64>,
179 a_col: usize,
180 b: &VecZnxBig<Data, FFT64>,
181 b_col: usize,
182 ) {
183 #[cfg(debug_assertions)]
184 {
185 assert_eq!(a.n(), self.n());
186 assert_eq!(b.n(), self.n());
187 assert_eq!(res.n(), self.n());
188 assert_ne!(a.as_ptr(), b.as_ptr());
189 }
190 unsafe {
191 vec_znx::vec_znx_add(
192 self.ptr,
193 res.at_mut_ptr(res_col, 0),
194 res.size() as u64,
195 res.sl() as u64,
196 a.at_ptr(a_col, 0),
197 a.size() as u64,
198 a.sl() as u64,
199 b.at_ptr(b_col, 0),
200 b.size() as u64,
201 b.sl() as u64,
202 )
203 }
204 }
205
206 fn vec_znx_big_add_inplace(
207 &self,
208 res: &mut VecZnxBig<DataMut, FFT64>,
209 res_col: usize,
210 a: &VecZnxBig<Data, FFT64>,
211 a_col: usize,
212 ) {
213 #[cfg(debug_assertions)]
214 {
215 assert_eq!(a.n(), self.n());
216 assert_eq!(res.n(), self.n());
217 }
218 unsafe {
219 vec_znx::vec_znx_add(
220 self.ptr,
221 res.at_mut_ptr(res_col, 0),
222 res.size() as u64,
223 res.sl() as u64,
224 a.at_ptr(a_col, 0),
225 a.size() as u64,
226 a.sl() as u64,
227 res.at_ptr(res_col, 0),
228 res.size() as u64,
229 res.sl() as u64,
230 )
231 }
232 }
233
234 fn vec_znx_big_sub(
235 &self,
236 res: &mut VecZnxBig<DataMut, FFT64>,
237 res_col: usize,
238 a: &VecZnxBig<Data, FFT64>,
239 a_col: usize,
240 b: &VecZnxBig<Data, FFT64>,
241 b_col: usize,
242 ) {
243 #[cfg(debug_assertions)]
244 {
245 assert_eq!(a.n(), self.n());
246 assert_eq!(b.n(), self.n());
247 assert_eq!(res.n(), self.n());
248 assert_ne!(a.as_ptr(), b.as_ptr());
249 }
250 unsafe {
251 vec_znx::vec_znx_sub(
252 self.ptr,
253 res.at_mut_ptr(res_col, 0),
254 res.size() as u64,
255 res.sl() as u64,
256 a.at_ptr(a_col, 0),
257 a.size() as u64,
258 a.sl() as u64,
259 b.at_ptr(b_col, 0),
260 b.size() as u64,
261 b.sl() as u64,
262 )
263 }
264 }
265
266 fn vec_znx_big_sub_ab_inplace(
267 &self,
268 res: &mut VecZnxBig<DataMut, FFT64>,
269 res_col: usize,
270 a: &VecZnxBig<Data, FFT64>,
271 a_col: usize,
272 ) {
273 #[cfg(debug_assertions)]
274 {
275 assert_eq!(a.n(), self.n());
276 assert_eq!(res.n(), self.n());
277 }
278 unsafe {
279 vec_znx::vec_znx_sub(
280 self.ptr,
281 res.at_mut_ptr(res_col, 0),
282 res.size() as u64,
283 res.sl() as u64,
284 res.at_ptr(res_col, 0),
285 res.size() as u64,
286 res.sl() as u64,
287 a.at_ptr(a_col, 0),
288 a.size() as u64,
289 a.sl() as u64,
290 )
291 }
292 }
293
294 fn vec_znx_big_sub_ba_inplace(
295 &self,
296 res: &mut VecZnxBig<DataMut, FFT64>,
297 res_col: usize,
298 a: &VecZnxBig<Data, FFT64>,
299 a_col: usize,
300 ) {
301 #[cfg(debug_assertions)]
302 {
303 assert_eq!(a.n(), self.n());
304 assert_eq!(res.n(), self.n());
305 }
306 unsafe {
307 vec_znx::vec_znx_sub(
308 self.ptr,
309 res.at_mut_ptr(res_col, 0),
310 res.size() as u64,
311 res.sl() as u64,
312 a.at_ptr(a_col, 0),
313 a.size() as u64,
314 a.sl() as u64,
315 res.at_ptr(res_col, 0),
316 res.size() as u64,
317 res.sl() as u64,
318 )
319 }
320 }
321
322 fn vec_znx_big_sub_small_b(
323 &self,
324 res: &mut VecZnxBig<DataMut, FFT64>,
325 res_col: usize,
326 a: &VecZnxBig<Data, FFT64>,
327 a_col: usize,
328 b: &VecZnx<Data>,
329 b_col: usize,
330 ) {
331 #[cfg(debug_assertions)]
332 {
333 assert_eq!(a.n(), self.n());
334 assert_eq!(b.n(), self.n());
335 assert_eq!(res.n(), self.n());
336 assert_ne!(a.as_ptr(), b.as_ptr());
337 }
338 unsafe {
339 vec_znx::vec_znx_sub(
340 self.ptr,
341 res.at_mut_ptr(res_col, 0),
342 res.size() as u64,
343 res.sl() as u64,
344 a.at_ptr(a_col, 0),
345 a.size() as u64,
346 a.sl() as u64,
347 b.at_ptr(b_col, 0),
348 b.size() as u64,
349 b.sl() as u64,
350 )
351 }
352 }
353
354 fn vec_znx_big_sub_small_b_inplace(
355 &self,
356 res: &mut VecZnxBig<DataMut, FFT64>,
357 res_col: usize,
358 a: &VecZnx<Data>,
359 a_col: usize,
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(
383 &self,
384 res: &mut VecZnxBig<DataMut, FFT64>,
385 res_col: usize,
386 a: &VecZnx<Data>,
387 a_col: usize,
388 b: &VecZnxBig<Data, FFT64>,
389 b_col: usize,
390 ) {
391 #[cfg(debug_assertions)]
392 {
393 assert_eq!(a.n(), self.n());
394 assert_eq!(b.n(), self.n());
395 assert_eq!(res.n(), self.n());
396 assert_ne!(a.as_ptr(), b.as_ptr());
397 }
398 unsafe {
399 vec_znx::vec_znx_sub(
400 self.ptr,
401 res.at_mut_ptr(res_col, 0),
402 res.size() as u64,
403 res.sl() as u64,
404 a.at_ptr(a_col, 0),
405 a.size() as u64,
406 a.sl() as u64,
407 b.at_ptr(b_col, 0),
408 b.size() as u64,
409 b.sl() as u64,
410 )
411 }
412 }
413
414 fn vec_znx_big_sub_small_a_inplace(
415 &self,
416 res: &mut VecZnxBig<DataMut, FFT64>,
417 res_col: usize,
418 a: &VecZnx<Data>,
419 a_col: usize,
420 ) {
421 #[cfg(debug_assertions)]
422 {
423 assert_eq!(a.n(), self.n());
424 assert_eq!(res.n(), self.n());
425 }
426 unsafe {
427 vec_znx::vec_znx_sub(
428 self.ptr,
429 res.at_mut_ptr(res_col, 0),
430 res.size() as u64,
431 res.sl() as u64,
432 res.at_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 )
439 }
440 }
441
442 fn vec_znx_big_add_small(
443 &self,
444 res: &mut VecZnxBig<DataMut, FFT64>,
445 res_col: usize,
446 a: &VecZnxBig<Data, FFT64>,
447 a_col: usize,
448 b: &VecZnx<Data>,
449 b_col: usize,
450 ) {
451 #[cfg(debug_assertions)]
452 {
453 assert_eq!(a.n(), self.n());
454 assert_eq!(b.n(), self.n());
455 assert_eq!(res.n(), self.n());
456 assert_ne!(a.as_ptr(), b.as_ptr());
457 }
458 unsafe {
459 vec_znx::vec_znx_add(
460 self.ptr,
461 res.at_mut_ptr(res_col, 0),
462 res.size() as u64,
463 res.sl() as u64,
464 a.at_ptr(a_col, 0),
465 a.size() as u64,
466 a.sl() as u64,
467 b.at_ptr(b_col, 0),
468 b.size() as u64,
469 b.sl() as u64,
470 )
471 }
472 }
473
474 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<DataMut, FFT64>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
475 #[cfg(debug_assertions)]
476 {
477 assert_eq!(a.n(), self.n());
478 assert_eq!(res.n(), self.n());
479 }
480 unsafe {
481 vec_znx::vec_znx_add(
482 self.ptr,
483 res.at_mut_ptr(res_col, 0),
484 res.size() as u64,
485 res.sl() as u64,
486 res.at_ptr(res_col, 0),
487 res.size() as u64,
488 res.sl() as u64,
489 a.at_ptr(a_col, 0),
490 a.size() as u64,
491 a.sl() as u64,
492 )
493 }
494 }
495
496 fn vec_znx_big_normalize(
497 &self,
498 log_base2k: usize,
499 res: &mut VecZnx<DataMut>,
500 res_col: usize,
501 a: &VecZnxBig<Data, FFT64>,
502 a_col: usize,
503 scratch: &mut Scratch,
504 ) {
505 #[cfg(debug_assertions)]
506 {
507 assert_eq!(a.n(), self.n());
508 assert_eq!(res.n(), self.n());
509 //(Jay)Note: This is calling VezZnxOps::vec_znx_normalize_tmp_bytes and not VecZnxBigOps::vec_znx_big_normalize_tmp_bytes.
510 // In the FFT backend the tmp sizes are same but will be different in the NTT backend
511 // assert!(tmp_bytes.len() >= <Self as VecZnxOps<DataMut, Data>>::vec_znx_normalize_tmp_bytes(&self));
512 // assert_alignement(tmp_bytes.as_ptr());
513 }
514
515 let (tmp_bytes, _) = scratch.tmp_scalar_slice(<Self as VecZnxBigScratch>::vec_znx_big_normalize_tmp_bytes(
516 &self,
517 ));
518 unsafe {
519 vec_znx::vec_znx_normalize_base2k(
520 self.ptr,
521 log_base2k as u64,
522 res.at_mut_ptr(res_col, 0),
523 res.size() as u64,
524 res.sl() as u64,
525 a.at_ptr(a_col, 0),
526 a.size() as u64,
527 a.sl() as u64,
528 tmp_bytes.as_mut_ptr(),
529 );
530 }
531 }
532
533 fn vec_znx_big_automorphism(
534 &self,
535 k: i64,
536 res: &mut VecZnxBig<DataMut, FFT64>,
537 res_col: usize,
538 a: &VecZnxBig<Data, FFT64>,
539 a_col: usize,
540 ) {
541 #[cfg(debug_assertions)]
542 {
543 assert_eq!(a.n(), self.n());
544 assert_eq!(res.n(), self.n());
545 }
546 unsafe {
547 vec_znx::vec_znx_automorphism(
548 self.ptr,
549 k,
550 res.at_mut_ptr(res_col, 0),
551 res.size() as u64,
552 res.sl() as u64,
553 a.at_ptr(a_col, 0),
554 a.size() as u64,
555 a.sl() as u64,
556 )
557 }
558 }
559
560 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<DataMut, FFT64>, a_col: usize) {
561 #[cfg(debug_assertions)]
562 {
563 assert_eq!(a.n(), self.n());
564 }
565 unsafe {
566 vec_znx::vec_znx_automorphism(
567 self.ptr,
568 k,
569 a.at_mut_ptr(a_col, 0),
570 a.size() as u64,
571 a.sl() as u64,
572 a.at_ptr(a_col, 0),
573 a.size() as u64,
574 a.sl() as u64,
575 )
576 }
577 }
578}
579
580impl<B: Backend> VecZnxBigScratch for Module<B> {
581 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
582 <Self as VecZnxScratch>::vec_znx_normalize_tmp_bytes(self)
583 }
584}