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