poulpy/base2k/src/vec_znx_ops.rs
base2k/src/vec_znx_ops.rsRaw
1use crate::ffi::vec_znx;
2use crate::{
3 Backend, Module, Scratch, VecZnx, VecZnxOwned, VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero,
4};
5use itertools::izip;
6use std::cmp::min;
7
8pub trait VecZnxAlloc {
9 /// Allocates a new [VecZnx].
10 ///
11 /// # Arguments
12 ///
13 /// * `cols`: the number of polynomials.
14 /// * `size`: the number small polynomials per column.
15 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnxOwned;
16
17 /// Instantiates a new [VecZnx] from a slice of bytes.
18 /// The returned [VecZnx] takes ownership of the slice of bytes.
19 ///
20 /// # Arguments
21 ///
22 /// * `cols`: the number of polynomials.
23 /// * `size`: the number small polynomials per column.
24 ///
25 /// # Panic
26 /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
27 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxOwned;
28
29 /// Returns the number of bytes necessary to allocate
30 /// a new [VecZnx] through [VecZnxOps::new_vec_znx_from_bytes]
31 /// or [VecZnxOps::new_vec_znx_from_bytes_borrow].
32 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize;
33}
34
35pub trait VecZnxOps {
36 /// Normalizes the selected column of `a` and stores the result into the selected column of `res`.
37 fn vec_znx_normalize<R, A>(&self, log_base2k: usize, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
38 where
39 R: VecZnxToMut,
40 A: VecZnxToRef;
41
42 /// Normalizes the selected column of `a`.
43 fn vec_znx_normalize_inplace<A>(&self, log_base2k: usize, a: &mut A, a_col: usize, scratch: &mut Scratch)
44 where
45 A: VecZnxToMut;
46
47 /// Adds the selected column of `a` to the selected column of `b` and writes the result on the selected column of `res`.
48 fn vec_znx_add<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
49 where
50 R: VecZnxToMut,
51 A: VecZnxToRef,
52 B: VecZnxToRef;
53
54 /// Adds the selected column of `a` to the selected column of `b` and writes the result on the selected column of `res`.
55 fn vec_znx_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
56 where
57 R: VecZnxToMut,
58 A: VecZnxToRef;
59
60 /// Subtracts the selected column of `b` from the selected column of `a` and writes the result on the selected column of `res`.
61 fn vec_znx_sub<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
62 where
63 R: VecZnxToMut,
64 A: VecZnxToRef,
65 B: VecZnxToRef;
66
67 /// Subtracts the selected column of `a` from the selected column of `res` inplace.
68 ///
69 /// res[res_col] -= a[a_col]
70 fn vec_znx_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
71 where
72 R: VecZnxToMut,
73 A: VecZnxToRef;
74
75 /// Subtracts the selected column of `res` from the selected column of `a` and inplace mutates `res`
76 ///
77 /// res[res_col] = a[a_col] - res[res_col]
78 fn vec_znx_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
79 where
80 R: VecZnxToMut,
81 A: VecZnxToRef;
82
83 // Negates the selected column of `a` and stores the result in `res_col` of `res`.
84 fn vec_znx_negate<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
85 where
86 R: VecZnxToMut,
87 A: VecZnxToRef;
88
89 /// Negates the selected column of `a`.
90 fn vec_znx_negate_inplace<A>(&self, a: &mut A, a_col: usize)
91 where
92 A: VecZnxToMut;
93
94 /// Multiplies the selected column of `a` by X^k and stores the result in `res_col` of `res`.
95 fn vec_znx_rotate<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
96 where
97 R: VecZnxToMut,
98 A: VecZnxToRef;
99
100 /// Multiplies the selected column of `a` by X^k.
101 fn vec_znx_rotate_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
102 where
103 A: VecZnxToMut;
104
105 /// Applies the automorphism X^i -> X^ik on the selected column of `a` and stores the result in `res_col` column of `res`.
106 fn vec_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
107 where
108 R: VecZnxToMut,
109 A: VecZnxToRef;
110
111 /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
112 fn vec_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
113 where
114 A: VecZnxToMut;
115
116 /// Splits the selected columns of `b` into subrings and copies them them into the selected column of `res`.
117 ///
118 /// # Panics
119 ///
120 /// This method requires that all [VecZnx] of b have the same ring degree
121 /// and that b.n() * b.len() <= a.n()
122 fn vec_znx_split<R, A>(&self, res: &mut Vec<R>, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
123 where
124 R: VecZnxToMut,
125 A: VecZnxToRef;
126
127 /// Merges the subrings of the selected column of `a` into the selected column of `res`.
128 ///
129 /// # Panics
130 ///
131 /// This method requires that all [VecZnx] of a have the same ring degree
132 /// and that a.n() * a.len() <= b.n()
133 fn vec_znx_merge<R, A>(&self, res: &mut R, res_col: usize, a: Vec<A>, a_col: usize)
134 where
135 R: VecZnxToMut,
136 A: VecZnxToRef;
137
138 fn switch_degree<R, A>(&self, r: &mut R, col_b: usize, a: &A, col_a: usize)
139 where
140 R: VecZnxToMut,
141 A: VecZnxToRef;
142}
143
144pub trait VecZnxScratch {
145 /// Returns the minimum number of bytes necessary for normalization.
146 fn vec_znx_normalize_tmp_bytes(&self) -> usize;
147}
148
149impl<B: Backend> VecZnxAlloc for Module<B> {
150 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnxOwned {
151 VecZnxOwned::new::<i64>(self.n(), cols, size)
152 }
153
154 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize {
155 VecZnxOwned::bytes_of::<i64>(self.n(), cols, size)
156 }
157
158 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxOwned {
159 VecZnxOwned::new_from_bytes::<i64>(self.n(), cols, size, bytes)
160 }
161}
162
163impl<BACKEND: Backend> VecZnxOps for Module<BACKEND> {
164 fn vec_znx_normalize<R, A>(&self, log_base2k: usize, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
165 where
166 R: VecZnxToMut,
167 A: VecZnxToRef,
168 {
169 let a: VecZnx<&[u8]> = a.to_ref();
170 let mut res: VecZnx<&mut [u8]> = res.to_mut();
171
172 #[cfg(debug_assertions)]
173 {
174 assert_eq!(a.n(), self.n());
175 assert_eq!(res.n(), self.n());
176 }
177
178 let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_normalize_tmp_bytes());
179
180 unsafe {
181 vec_znx::vec_znx_normalize_base2k(
182 self.ptr,
183 log_base2k as u64,
184 res.at_mut_ptr(res_col, 0),
185 res.size() as u64,
186 res.sl() as u64,
187 a.at_ptr(a_col, 0),
188 a.size() as u64,
189 a.sl() as u64,
190 tmp_bytes.as_mut_ptr(),
191 );
192 }
193 }
194
195 fn vec_znx_normalize_inplace<A>(&self, log_base2k: usize, a: &mut A, a_col: usize, scratch: &mut Scratch)
196 where
197 A: VecZnxToMut,
198 {
199 let mut a: VecZnx<&mut [u8]> = a.to_mut();
200
201 #[cfg(debug_assertions)]
202 {
203 assert_eq!(a.n(), self.n());
204 }
205
206 let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_normalize_tmp_bytes());
207
208 unsafe {
209 vec_znx::vec_znx_normalize_base2k(
210 self.ptr,
211 log_base2k as u64,
212 a.at_mut_ptr(a_col, 0),
213 a.size() as u64,
214 a.sl() as u64,
215 a.at_ptr(a_col, 0),
216 a.size() as u64,
217 a.sl() as u64,
218 tmp_bytes.as_mut_ptr(),
219 );
220 }
221 }
222
223 fn vec_znx_add<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
224 where
225 R: VecZnxToMut,
226 A: VecZnxToRef,
227 B: VecZnxToRef,
228 {
229 let a: VecZnx<&[u8]> = a.to_ref();
230 let b: VecZnx<&[u8]> = b.to_ref();
231 let mut res: VecZnx<&mut [u8]> = res.to_mut();
232
233 #[cfg(debug_assertions)]
234 {
235 assert_eq!(a.n(), self.n());
236 assert_eq!(b.n(), self.n());
237 assert_eq!(res.n(), self.n());
238 assert_ne!(a.as_ptr(), b.as_ptr());
239 }
240 unsafe {
241 vec_znx::vec_znx_add(
242 self.ptr,
243 res.at_mut_ptr(res_col, 0),
244 res.size() as u64,
245 res.sl() as u64,
246 a.at_ptr(a_col, 0),
247 a.size() as u64,
248 a.sl() as u64,
249 b.at_ptr(b_col, 0),
250 b.size() as u64,
251 b.sl() as u64,
252 )
253 }
254 }
255
256 fn vec_znx_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
257 where
258 R: VecZnxToMut,
259 A: VecZnxToRef,
260 {
261 let a: VecZnx<&[u8]> = a.to_ref();
262 let mut res: VecZnx<&mut [u8]> = res.to_mut();
263
264 #[cfg(debug_assertions)]
265 {
266 assert_eq!(a.n(), self.n());
267 assert_eq!(res.n(), self.n());
268 }
269 unsafe {
270 vec_znx::vec_znx_add(
271 self.ptr,
272 res.at_mut_ptr(res_col, 0),
273 res.size() as u64,
274 res.sl() as u64,
275 a.at_ptr(a_col, 0),
276 a.size() as u64,
277 a.sl() as u64,
278 res.at_ptr(res_col, 0),
279 res.size() as u64,
280 res.sl() as u64,
281 )
282 }
283 }
284
285 fn vec_znx_sub<R, A, B>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
286 where
287 R: VecZnxToMut,
288 A: VecZnxToRef,
289 B: VecZnxToRef,
290 {
291 let a: VecZnx<&[u8]> = a.to_ref();
292 let b: VecZnx<&[u8]> = b.to_ref();
293 let mut res: VecZnx<&mut [u8]> = res.to_mut();
294
295 #[cfg(debug_assertions)]
296 {
297 assert_eq!(a.n(), self.n());
298 assert_eq!(b.n(), self.n());
299 assert_eq!(res.n(), self.n());
300 assert_ne!(a.as_ptr(), b.as_ptr());
301 }
302 unsafe {
303 vec_znx::vec_znx_sub(
304 self.ptr,
305 res.at_mut_ptr(res_col, 0),
306 res.size() as u64,
307 res.sl() as u64,
308 a.at_ptr(a_col, 0),
309 a.size() as u64,
310 a.sl() as u64,
311 b.at_ptr(b_col, 0),
312 b.size() as u64,
313 b.sl() as u64,
314 )
315 }
316 }
317
318 fn vec_znx_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
319 where
320 R: VecZnxToMut,
321 A: VecZnxToRef,
322 {
323 let a: VecZnx<&[u8]> = a.to_ref();
324 let mut res: VecZnx<&mut [u8]> = res.to_mut();
325 #[cfg(debug_assertions)]
326 {
327 assert_eq!(a.n(), self.n());
328 assert_eq!(res.n(), self.n());
329 }
330 unsafe {
331 vec_znx::vec_znx_sub(
332 self.ptr,
333 res.at_mut_ptr(res_col, 0),
334 res.size() as u64,
335 res.sl() as u64,
336 res.at_ptr(res_col, 0),
337 res.size() as u64,
338 res.sl() as u64,
339 a.at_ptr(a_col, 0),
340 a.size() as u64,
341 a.sl() as u64,
342 )
343 }
344 }
345
346 fn vec_znx_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
347 where
348 R: VecZnxToMut,
349 A: VecZnxToRef,
350 {
351 let a: VecZnx<&[u8]> = a.to_ref();
352 let mut res: VecZnx<&mut [u8]> = res.to_mut();
353 #[cfg(debug_assertions)]
354 {
355 assert_eq!(a.n(), self.n());
356 assert_eq!(res.n(), self.n());
357 }
358 unsafe {
359 vec_znx::vec_znx_sub(
360 self.ptr,
361 res.at_mut_ptr(res_col, 0),
362 res.size() as u64,
363 res.sl() as u64,
364 a.at_ptr(a_col, 0),
365 a.size() as u64,
366 a.sl() as u64,
367 res.at_ptr(res_col, 0),
368 res.size() as u64,
369 res.sl() as u64,
370 )
371 }
372 }
373
374 fn vec_znx_negate<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
375 where
376 R: VecZnxToMut,
377 A: VecZnxToRef,
378 {
379 let a: VecZnx<&[u8]> = a.to_ref();
380 let mut res: VecZnx<&mut [u8]> = res.to_mut();
381 #[cfg(debug_assertions)]
382 {
383 assert_eq!(a.n(), self.n());
384 assert_eq!(res.n(), self.n());
385 }
386 unsafe {
387 vec_znx::vec_znx_negate(
388 self.ptr,
389 res.at_mut_ptr(res_col, 0),
390 res.size() as u64,
391 res.sl() as u64,
392 a.at_ptr(a_col, 0),
393 a.size() as u64,
394 a.sl() as u64,
395 )
396 }
397 }
398
399 fn vec_znx_negate_inplace<A>(&self, a: &mut A, a_col: usize)
400 where
401 A: VecZnxToMut,
402 {
403 let mut a: VecZnx<&mut [u8]> = a.to_mut();
404 #[cfg(debug_assertions)]
405 {
406 assert_eq!(a.n(), self.n());
407 }
408 unsafe {
409 vec_znx::vec_znx_negate(
410 self.ptr,
411 a.at_mut_ptr(a_col, 0),
412 a.size() as u64,
413 a.sl() as u64,
414 a.at_ptr(a_col, 0),
415 a.size() as u64,
416 a.sl() as u64,
417 )
418 }
419 }
420
421 fn vec_znx_rotate<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
422 where
423 R: VecZnxToMut,
424 A: VecZnxToRef,
425 {
426 let a: VecZnx<&[u8]> = a.to_ref();
427 let mut res: VecZnx<&mut [u8]> = res.to_mut();
428 #[cfg(debug_assertions)]
429 {
430 assert_eq!(a.n(), self.n());
431 assert_eq!(res.n(), self.n());
432 }
433 unsafe {
434 vec_znx::vec_znx_rotate(
435 self.ptr,
436 k,
437 res.at_mut_ptr(res_col, 0),
438 res.size() as u64,
439 res.sl() as u64,
440 a.at_ptr(a_col, 0),
441 a.size() as u64,
442 a.sl() as u64,
443 )
444 }
445 }
446
447 fn vec_znx_rotate_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
448 where
449 A: VecZnxToMut,
450 {
451 let mut a: VecZnx<&mut [u8]> = a.to_mut();
452 #[cfg(debug_assertions)]
453 {
454 assert_eq!(a.n(), self.n());
455 }
456 unsafe {
457 vec_znx::vec_znx_rotate(
458 self.ptr,
459 k,
460 a.at_mut_ptr(a_col, 0),
461 a.size() as u64,
462 a.sl() as u64,
463 a.at_ptr(a_col, 0),
464 a.size() as u64,
465 a.sl() as u64,
466 )
467 }
468 }
469
470 fn vec_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
471 where
472 R: VecZnxToMut,
473 A: VecZnxToRef,
474 {
475 let a: VecZnx<&[u8]> = a.to_ref();
476 let mut res: VecZnx<&mut [u8]> = res.to_mut();
477 #[cfg(debug_assertions)]
478 {
479 assert_eq!(a.n(), self.n());
480 assert_eq!(res.n(), self.n());
481 }
482 unsafe {
483 vec_znx::vec_znx_automorphism(
484 self.ptr,
485 k,
486 res.at_mut_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_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
497 where
498 A: VecZnxToMut,
499 {
500 let mut a: VecZnx<&mut [u8]> = a.to_mut();
501 #[cfg(debug_assertions)]
502 {
503 assert_eq!(a.n(), self.n());
504 }
505 unsafe {
506 vec_znx::vec_znx_automorphism(
507 self.ptr,
508 k,
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 fn vec_znx_split<R, A>(&self, res: &mut Vec<R>, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
520 where
521 R: VecZnxToMut,
522 A: VecZnxToRef,
523 {
524 let a: VecZnx<&[u8]> = a.to_ref();
525
526 let (n_in, n_out) = (a.n(), res[0].to_mut().n());
527
528 let (mut buf, _) = scratch.tmp_vec_znx(self, 1, a.size());
529
530 debug_assert!(
531 n_out < n_in,
532 "invalid a: output ring degree should be smaller"
533 );
534 res[1..].iter_mut().for_each(|bi| {
535 debug_assert_eq!(
536 bi.to_mut().n(),
537 n_out,
538 "invalid input a: all VecZnx must have the same degree"
539 )
540 });
541
542 res.iter_mut().enumerate().for_each(|(i, bi)| {
543 if i == 0 {
544 self.switch_degree(bi, res_col, &a, a_col);
545 self.vec_znx_rotate(-1, &mut buf, 0, &a, a_col);
546 } else {
547 self.switch_degree(bi, res_col, &mut buf, a_col);
548 self.vec_znx_rotate_inplace(-1, &mut buf, a_col);
549 }
550 })
551 }
552
553 fn vec_znx_merge<R, A>(&self, res: &mut R, res_col: usize, a: Vec<A>, a_col: usize)
554 where
555 R: VecZnxToMut,
556 A: VecZnxToRef,
557 {
558 let mut res: VecZnx<&mut [u8]> = res.to_mut();
559
560 let (n_in, n_out) = (res.n(), a[0].to_ref().n());
561
562 debug_assert!(
563 n_out < n_in,
564 "invalid a: output ring degree should be smaller"
565 );
566 a[1..].iter().for_each(|ai| {
567 debug_assert_eq!(
568 ai.to_ref().n(),
569 n_out,
570 "invalid input a: all VecZnx must have the same degree"
571 )
572 });
573
574 a.iter().enumerate().for_each(|(_, ai)| {
575 self.switch_degree(&mut res, res_col, ai, a_col);
576 self.vec_znx_rotate_inplace(-1, &mut res, res_col);
577 });
578
579 self.vec_znx_rotate_inplace(a.len() as i64, &mut res, res_col);
580 }
581
582 fn switch_degree<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
583 where
584 R: VecZnxToMut,
585 A: VecZnxToRef,
586 {
587 let a: VecZnx<&[u8]> = a.to_ref();
588 let mut res: VecZnx<&mut [u8]> = res.to_mut();
589
590 let (n_in, n_out) = (a.n(), res.n());
591 let (gap_in, gap_out): (usize, usize);
592
593 if n_in > n_out {
594 (gap_in, gap_out) = (n_in / n_out, 1)
595 } else {
596 (gap_in, gap_out) = (1, n_out / n_in);
597 res.zero();
598 }
599
600 let size: usize = min(a.size(), res.size());
601
602 (0..size).for_each(|i| {
603 izip!(
604 a.at(a_col, i).iter().step_by(gap_in),
605 res.at_mut(res_col, i).iter_mut().step_by(gap_out)
606 )
607 .for_each(|(x_in, x_out)| *x_out = *x_in);
608 });
609 }
610}
611
612impl<B: Backend> VecZnxScratch for Module<B> {
613 fn vec_znx_normalize_tmp_bytes(&self) -> usize {
614 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
615 }
616}