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