poulpy/base2k/src/vec_znx_ops.rs
base2k/src/vec_znx_ops.rsRaw
1use crate::ffi::vec_znx;
2use crate::znx_base::{ZnxInfos, switch_degree};
3use crate::{Backend, Module, VecZnx, VecZnxOwned, ZnxView, ZnxViewMut, assert_alignement};
4
5pub trait VecZnxAlloc {
6 /// Allocates a new [VecZnx].
7 ///
8 /// # Arguments
9 ///
10 /// * `cols`: the number of polynomials.
11 /// * `size`: the number small polynomials per column.
12 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnxOwned;
13
14 /// Instantiates a new [VecZnx] from a slice of bytes.
15 /// The returned [VecZnx] takes ownership of the slice of bytes.
16 ///
17 /// # Arguments
18 ///
19 /// * `cols`: the number of polynomials.
20 /// * `size`: the number small polynomials per column.
21 ///
22 /// # Panic
23 /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
24 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxOwned;
25
26 // /// Instantiates a new [VecZnx] from a slice of bytes.
27 // /// The returned [VecZnx] does take ownership of the slice of bytes.
28 // ///
29 // /// # Arguments
30 // ///
31 // /// * `cols`: the number of polynomials.
32 // /// * `size`: the number small polynomials per column.
33 // ///
34 // /// # Panic
35 // /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
36 // fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx;
37 // (Jay)TODO
38
39 /// Returns the number of bytes necessary to allocate
40 /// a new [VecZnx] through [VecZnxOps::new_vec_znx_from_bytes]
41 /// or [VecZnxOps::new_vec_znx_from_bytes_borrow].
42 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize;
43}
44
45pub trait VecZnxOps<DataMut, Data> {
46 /// Returns the minimum number of bytes necessary for normalization.
47 fn vec_znx_normalize_tmp_bytes(&self) -> usize;
48
49 /// Normalizes the selected column of `a` and stores the result into the selected column of `res`.
50 fn vec_znx_normalize(
51 &self,
52 log_base2k: usize,
53 res: &mut VecZnx<DataMut>,
54 res_col: usize,
55 a: &VecZnx<Data>,
56 a_col: usize,
57 tmp_bytes: &mut [u8],
58 );
59
60 /// Normalizes the selected column of `a`.
61 fn vec_znx_normalize_inplace(&self, log_base2k: usize, a: &mut VecZnx<DataMut>, a_col: usize, tmp_bytes: &mut [u8]);
62
63 /// Adds the selected column of `a` to the selected column of `b` and writes the result on the selected column of `res`.
64 fn vec_znx_add(
65 &self,
66 res: &mut VecZnx<DataMut>,
67 res_col: usize,
68 a: &VecZnx<Data>,
69 a_col: usize,
70 b: &VecZnx<Data>,
71 b_col: usize,
72 );
73
74 /// Adds the selected column of `a` to the selected column of `b` and writes the result on the selected column of `res`.
75 fn vec_znx_add_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
76
77 /// Subtracts the selected column of `b` from the selected column of `a` and writes the result on the selected column of `res`.
78 fn vec_znx_sub(
79 &self,
80 res: &mut VecZnx<DataMut>,
81 res_col: usize,
82 a: &VecZnx<Data>,
83 a_col: usize,
84 b: &VecZnx<Data>,
85 b_col: usize,
86 );
87
88 /// Subtracts the selected column of `a` from the selected column of `res` inplace.
89 ///
90 /// res[res_col] -= a[a_col]
91 fn vec_znx_sub_ab_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
92
93 /// Subtracts the selected column of `res` from the selected column of `a` and inplace mutates `res`
94 ///
95 /// res[res_col] = a[a_col] - res[res_col]
96 fn vec_znx_sub_ba_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
97
98 // Negates the selected column of `a` and stores the result in `res_col` of `res`.
99 fn vec_znx_negate(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
100
101 /// Negates the selected column of `a`.
102 fn vec_znx_negate_inplace(&self, a: &mut VecZnx<DataMut>, a_col: usize);
103
104 /// Multiplies the selected column of `a` by X^k and stores the result in `res_col` of `res`.
105 fn vec_znx_rotate(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
106
107 /// Multiplies the selected column of `a` by X^k.
108 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize);
109
110 /// Applies the automorphism X^i -> X^ik on the selected column of `a` and stores the result in `res_col` column of `res`.
111 fn vec_znx_automorphism(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
112
113 /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
114 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize);
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(
123 &self,
124 res: &mut Vec<VecZnx<DataMut>>,
125 res_col: usize,
126 a: &VecZnx<Data>,
127 a_col: usize,
128 buf: &mut VecZnx<DataMut>,
129 );
130
131 /// Merges the subrings of the selected column of `a` into the selected column of `res`.
132 ///
133 /// # Panics
134 ///
135 /// This method requires that all [VecZnx] of a have the same ring degree
136 /// and that a.n() * a.len() <= b.n()
137 fn vec_znx_merge(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &Vec<VecZnx<Data>>, a_col: usize);
138}
139
140impl<B: Backend> VecZnxAlloc for Module<B> {
141 //(Jay)TODO: One must define the Scalar generic param here.
142 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnxOwned {
143 VecZnxOwned::new::<i64>(self.n(), cols, size)
144 }
145
146 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize {
147 VecZnxOwned::bytes_of::<i64>(self.n(), cols, size)
148 }
149
150 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxOwned {
151 VecZnxOwned::new_from_bytes::<i64>(self.n(), cols, size, bytes)
152 }
153}
154
155impl<B: Backend, DataMut, Data> VecZnxOps<DataMut, Data> for Module<B>
156where
157 Data: AsRef<[u8]>,
158 DataMut: AsRef<[u8]> + AsMut<[u8]>,
159{
160 fn vec_znx_normalize_tmp_bytes(&self) -> usize {
161 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
162 }
163
164 fn vec_znx_normalize(
165 &self,
166 log_base2k: usize,
167 res: &mut VecZnx<DataMut>,
168 res_col: usize,
169 a: &VecZnx<Data>,
170 a_col: usize,
171 tmp_bytes: &mut [u8],
172 ) {
173 #[cfg(debug_assertions)]
174 {
175 assert_eq!(a.n(), self.n());
176 assert_eq!(res.n(), self.n());
177 assert!(tmp_bytes.len() >= <Self as VecZnxOps<DataMut, Data>>::vec_znx_normalize_tmp_bytes(&self));
178 assert_alignement(tmp_bytes.as_ptr());
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(&self, log_base2k: usize, a: &mut VecZnx<DataMut>, a_col: usize, tmp_bytes: &mut [u8]) {
196 unsafe {
197 let a_ptr: *const VecZnx<_> = a;
198 Self::vec_znx_normalize(self, log_base2k, a, a_col, &*a_ptr, a_col, tmp_bytes);
199 }
200 }
201
202 fn vec_znx_add(
203 &self,
204 res: &mut VecZnx<DataMut>,
205 res_col: usize,
206 a: &VecZnx<Data>,
207 a_col: usize,
208 b: &VecZnx<Data>,
209 b_col: usize,
210 ) {
211 #[cfg(debug_assertions)]
212 {
213 assert_eq!(a.n(), self.n());
214 assert_eq!(b.n(), self.n());
215 assert_eq!(res.n(), self.n());
216 assert_ne!(a.as_ptr(), b.as_ptr());
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 b.at_ptr(b_col, 0),
228 b.size() as u64,
229 b.sl() as u64,
230 )
231 }
232 }
233
234 fn vec_znx_add_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
235 #[cfg(debug_assertions)]
236 {
237 assert_eq!(a.n(), self.n());
238 assert_eq!(res.n(), self.n());
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 res.at_ptr(res_col, 0),
250 res.size() as u64,
251 res.sl() as u64,
252 )
253 }
254 }
255
256 fn vec_znx_sub(
257 &self,
258 res: &mut VecZnx<DataMut>,
259 res_col: usize,
260 a: &VecZnx<Data>,
261 a_col: usize,
262 b: &VecZnx<Data>,
263 b_col: usize,
264 ) {
265 #[cfg(debug_assertions)]
266 {
267 assert_eq!(a.n(), self.n());
268 assert_eq!(b.n(), self.n());
269 assert_eq!(res.n(), self.n());
270 assert_ne!(a.as_ptr(), b.as_ptr());
271 }
272 unsafe {
273 vec_znx::vec_znx_sub(
274 self.ptr,
275 res.at_mut_ptr(res_col, 0),
276 res.size() as u64,
277 res.sl() as u64,
278 a.at_ptr(a_col, 0),
279 a.size() as u64,
280 a.sl() as u64,
281 b.at_ptr(b_col, 0),
282 b.size() as u64,
283 b.sl() as u64,
284 )
285 }
286 }
287
288 fn vec_znx_sub_ab_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
289 #[cfg(debug_assertions)]
290 {
291 assert_eq!(a.n(), self.n());
292 assert_eq!(res.n(), self.n());
293 }
294 unsafe {
295 vec_znx::vec_znx_sub(
296 self.ptr,
297 res.at_mut_ptr(res_col, 0),
298 res.size() as u64,
299 res.sl() as u64,
300 res.at_ptr(res_col, 0),
301 res.size() as u64,
302 res.sl() as u64,
303 a.at_ptr(a_col, 0),
304 a.size() as u64,
305 a.sl() as u64,
306 )
307 }
308 }
309
310 fn vec_znx_sub_ba_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
311 #[cfg(debug_assertions)]
312 {
313 assert_eq!(a.n(), self.n());
314 assert_eq!(res.n(), self.n());
315 }
316 unsafe {
317 vec_znx::vec_znx_sub(
318 self.ptr,
319 res.at_mut_ptr(res_col, 0),
320 res.size() as u64,
321 res.sl() as u64,
322 a.at_ptr(a_col, 0),
323 a.size() as u64,
324 a.sl() as u64,
325 res.at_ptr(res_col, 0),
326 res.size() as u64,
327 res.sl() as u64,
328 )
329 }
330 }
331
332 fn vec_znx_negate(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
333 #[cfg(debug_assertions)]
334 {
335 assert_eq!(a.n(), self.n());
336 assert_eq!(res.n(), self.n());
337 }
338 unsafe {
339 vec_znx::vec_znx_negate(
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 )
348 }
349 }
350
351 fn vec_znx_negate_inplace(&self, a: &mut VecZnx<DataMut>, a_col: usize) {
352 unsafe {
353 let a_ref: *const VecZnx<_> = a;
354 Self::vec_znx_negate(self, a, a_col, a_ref.as_ref().unwrap(), a_col);
355 }
356 }
357
358 fn vec_znx_rotate(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
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_rotate(
366 self.ptr,
367 k,
368 res.at_mut_ptr(res_col, 0),
369 res.size() as u64,
370 res.sl() as u64,
371 a.at_ptr(a_col, 0),
372 a.size() as u64,
373 a.sl() as u64,
374 )
375 }
376 }
377
378 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize) {
379 #[cfg(debug_assertions)]
380 {
381 assert_eq!(a.n(), self.n());
382 }
383 unsafe {
384 vec_znx::vec_znx_rotate(
385 self.ptr,
386 k,
387 a.at_mut_ptr(a_col, 0),
388 a.size() as u64,
389 a.sl() as u64,
390 a.at_ptr(a_col, 0),
391 a.size() as u64,
392 a.sl() as u64,
393 )
394 }
395 }
396
397 fn vec_znx_automorphism(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
398 #[cfg(debug_assertions)]
399 {
400 assert_eq!(a.n(), self.n());
401 assert_eq!(res.n(), self.n());
402 }
403 unsafe {
404 vec_znx::vec_znx_automorphism(
405 self.ptr,
406 k,
407 res.at_mut_ptr(res_col, 0),
408 res.size() as u64,
409 res.sl() as u64,
410 a.at_ptr(a_col, 0),
411 a.size() as u64,
412 a.sl() as u64,
413 )
414 }
415 }
416
417 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize) {
418 #[cfg(debug_assertions)]
419 {
420 assert_eq!(a.n(), self.n());
421 }
422 unsafe {
423 vec_znx::vec_znx_automorphism(
424 self.ptr,
425 k,
426 a.at_mut_ptr(a_col, 0),
427 a.size() as u64,
428 a.sl() as u64,
429 a.at_ptr(a_col, 0),
430 a.size() as u64,
431 a.sl() as u64,
432 )
433 }
434 }
435
436 fn vec_znx_split(
437 &self,
438 res: &mut Vec<VecZnx<DataMut>>,
439 res_col: usize,
440 a: &VecZnx<Data>,
441 a_col: usize,
442 buf: &mut VecZnx<DataMut>,
443 ) {
444 let (n_in, n_out) = (a.n(), res[0].n());
445
446 debug_assert!(
447 n_out < n_in,
448 "invalid a: output ring degree should be smaller"
449 );
450 res[1..].iter().for_each(|bi| {
451 debug_assert_eq!(
452 bi.n(),
453 n_out,
454 "invalid input a: all VecZnx must have the same degree"
455 )
456 });
457
458 res.iter_mut().enumerate().for_each(|(i, bi)| {
459 if i == 0 {
460 switch_degree(bi, res_col, a, a_col);
461 self.vec_znx_rotate(-1, buf, 0, a, a_col);
462 } else {
463 switch_degree(bi, res_col, buf, a_col);
464 <Self as VecZnxOps<DataMut, Data>>::vec_znx_rotate_inplace(self, -1, buf, a_col);
465 }
466 })
467 }
468
469 fn vec_znx_merge(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &Vec<VecZnx<Data>>, a_col: usize) {
470 let (n_in, n_out) = (res.n(), a[0].n());
471
472 debug_assert!(
473 n_out < n_in,
474 "invalid a: output ring degree should be smaller"
475 );
476 a[1..].iter().for_each(|ai| {
477 debug_assert_eq!(
478 ai.n(),
479 n_out,
480 "invalid input a: all VecZnx must have the same degree"
481 )
482 });
483
484 a.iter().enumerate().for_each(|(_, ai)| {
485 switch_degree(res, res_col, ai, a_col);
486 <Self as VecZnxOps<DataMut, Data>>::vec_znx_rotate_inplace(self, -1, res, res_col);
487 });
488
489 <Self as VecZnxOps<DataMut, Data>>::vec_znx_rotate_inplace(self, a.len() as i64, res, res_col);
490 }
491}