poulpy/base2k/src/vec_znx_ops.rs
base2k/src/vec_znx_ops.rsRaw
1use crate::ffi::vec_znx;
2use crate::znx_base::{ZnxAlloc, ZnxInfos, ZnxLayout, ZnxSliceSize, switch_degree};
3use crate::{Backend, Module, VEC_ZNX_ROWS, VecZnx, assert_alignement};
4pub trait VecZnxOps {
5 /// Allocates a new [VecZnx].
6 ///
7 /// # Arguments
8 ///
9 /// * `cols`: the number of polynomials.
10 /// * `size`: the number small polynomials per column.
11 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnx;
12
13 /// Instantiates a new [VecZnx] from a slice of bytes.
14 /// The returned [VecZnx] takes ownership of the slice of bytes.
15 ///
16 /// # Arguments
17 ///
18 /// * `cols`: the number of polynomials.
19 /// * `size`: the number small polynomials per column.
20 ///
21 /// # Panic
22 /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
23 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnx;
24
25 /// Instantiates a new [VecZnx] from a slice of bytes.
26 /// The returned [VecZnx] does take ownership of the slice of bytes.
27 ///
28 /// # Arguments
29 ///
30 /// * `cols`: the number of polynomials.
31 /// * `size`: the number small polynomials per column.
32 ///
33 /// # Panic
34 /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
35 fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx;
36
37 /// Returns the number of bytes necessary to allocate
38 /// a new [VecZnx] through [VecZnxOps::new_vec_znx_from_bytes]
39 /// or [VecZnxOps::new_vec_znx_from_bytes_borrow].
40 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize;
41
42 /// Returns the minimum number of bytes necessary for normalization.
43 fn vec_znx_normalize_tmp_bytes(&self) -> usize;
44
45 /// Normalizes the selected column of `a` and stores the result into the selected column of `res`.
46 fn vec_znx_normalize(
47 &self,
48 log_base2k: usize,
49 res: &mut VecZnx,
50 res_col: usize,
51 a: &VecZnx,
52 a_col: usize,
53 tmp_bytes: &mut [u8],
54 );
55
56 /// Normalizes the selected column of `a`.
57 fn vec_znx_normalize_inplace(&self, log_base2k: usize, a: &mut VecZnx, a_col: usize, tmp_bytes: &mut [u8]);
58
59 /// Adds the selected column of `a` to the selected column of `b` and write the result on the selected column of `c`.
60 fn vec_znx_add(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize, b: &VecZnx, b_col: usize);
61
62 /// Adds the selected column of `a` to the selected column of `b` and write the result on the selected column of `res`.
63 fn vec_znx_add_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize);
64
65 /// Subtracts the selected column of `b` to the selected column of `a` and write the result on the selected column of `res`.
66 fn vec_znx_sub(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize, b: &VecZnx, b_col: usize);
67
68 /// Subtracts the selected column of `a` to the selected column of `res`.
69 fn vec_znx_sub_ab_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize);
70
71 /// Subtracts the selected column of `a` to the selected column of `res` and negates the selected column of `res`.
72 fn vec_znx_sub_ba_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize);
73
74 // Negates the selected column of `a` and stores the result on the selected column of `res`.
75 fn vec_znx_negate(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize);
76
77 /// Negates the selected column of `a`.
78 fn vec_znx_negate_inplace(&self, a: &mut VecZnx, a_col: usize);
79
80 /// Multiplies the selected column of `a` by X^k and stores the result on the selected column of `res`.
81 fn vec_znx_rotate(&self, k: i64, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize);
82
83 /// Multiplies the selected column of `a` by X^k.
84 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx, a_col: usize);
85
86 /// Applies the automorphism X^i -> X^ik on the selected column of `a` and stores the result on the selected column of `res`.
87 fn vec_znx_automorphism(&self, k: i64, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize);
88
89 /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
90 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx, a_col: usize);
91
92 /// Splits the selected columns of `b` into subrings and copies them them into the selected column of `res`.
93 ///
94 /// # Panics
95 ///
96 /// This method requires that all [VecZnx] of b have the same ring degree
97 /// and that b.n() * b.len() <= a.n()
98 fn vec_znx_split(&self, res: &mut Vec<VecZnx>, res_col: usize, a: &VecZnx, a_col: usize, buf: &mut VecZnx);
99
100 /// Merges the subrings of the selected column of `a` into the selected column of `res`.
101 ///
102 /// # Panics
103 ///
104 /// This method requires that all [VecZnx] of a have the same ring degree
105 /// and that a.n() * a.len() <= b.n()
106 fn vec_znx_merge(&self, res: &mut VecZnx, res_col: usize, a: &Vec<VecZnx>, a_col: usize);
107}
108
109impl<B: Backend> VecZnxOps for Module<B> {
110 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnx {
111 VecZnx::new(self, VEC_ZNX_ROWS, cols, size)
112 }
113
114 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize {
115 VecZnx::bytes_of(self, VEC_ZNX_ROWS, cols, size)
116 }
117
118 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnx {
119 VecZnx::from_bytes(self, VEC_ZNX_ROWS, cols, size, bytes)
120 }
121
122 fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx {
123 VecZnx::from_bytes_borrow(self, VEC_ZNX_ROWS, cols, size, tmp_bytes)
124 }
125
126 fn vec_znx_normalize_tmp_bytes(&self) -> usize {
127 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
128 }
129
130 fn vec_znx_normalize(
131 &self,
132 log_base2k: usize,
133 res: &mut VecZnx,
134 res_col: usize,
135 a: &VecZnx,
136 a_col: usize,
137 tmp_bytes: &mut [u8],
138 ) {
139 #[cfg(debug_assertions)]
140 {
141 assert_eq!(a.n(), self.n());
142 assert_eq!(res.n(), self.n());
143 assert!(tmp_bytes.len() >= Self::vec_znx_normalize_tmp_bytes(&self));
144 assert_alignement(tmp_bytes.as_ptr());
145 }
146 unsafe {
147 vec_znx::vec_znx_normalize_base2k(
148 self.ptr,
149 log_base2k as u64,
150 res.at_mut_ptr(res_col, 0),
151 res.size() as u64,
152 res.sl() as u64,
153 a.at_ptr(a_col, 0),
154 a.size() as u64,
155 a.sl() as u64,
156 tmp_bytes.as_mut_ptr(),
157 );
158 }
159 }
160
161 fn vec_znx_normalize_inplace(&self, log_base2k: usize, a: &mut VecZnx, a_col: usize, tmp_bytes: &mut [u8]) {
162 unsafe {
163 let a_ptr: *mut VecZnx = a as *mut VecZnx;
164 Self::vec_znx_normalize(
165 self,
166 log_base2k,
167 &mut *a_ptr,
168 a_col,
169 &*a_ptr,
170 a_col,
171 tmp_bytes,
172 );
173 }
174 }
175
176 fn vec_znx_add(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize, b: &VecZnx, b_col: usize) {
177 #[cfg(debug_assertions)]
178 {
179 assert_eq!(a.n(), self.n());
180 assert_eq!(b.n(), self.n());
181 assert_eq!(res.n(), self.n());
182 assert_ne!(a.as_ptr(), b.as_ptr());
183 }
184 unsafe {
185 vec_znx::vec_znx_add(
186 self.ptr,
187 res.at_mut_ptr(res_col, 0),
188 res.size() as u64,
189 res.sl() as u64,
190 a.at_ptr(a_col, 0),
191 a.size() as u64,
192 a.sl() as u64,
193 b.at_ptr(b_col, 0),
194 b.size() as u64,
195 b.sl() as u64,
196 )
197 }
198 }
199
200 fn vec_znx_add_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
201 unsafe {
202 let res_ptr: *mut VecZnx = res as *mut VecZnx;
203 Self::vec_znx_add(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
204 }
205 }
206
207 fn vec_znx_sub(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize, b: &VecZnx, b_col: usize) {
208 #[cfg(debug_assertions)]
209 {
210 assert_eq!(a.n(), self.n());
211 assert_eq!(b.n(), self.n());
212 assert_eq!(res.n(), self.n());
213 assert_ne!(a.as_ptr(), b.as_ptr());
214 }
215 unsafe {
216 vec_znx::vec_znx_sub(
217 self.ptr,
218 res.at_mut_ptr(res_col, 0),
219 res.size() as u64,
220 res.sl() as u64,
221 a.at_ptr(a_col, 0),
222 a.size() as u64,
223 a.sl() as u64,
224 b.at_ptr(b_col, 0),
225 b.size() as u64,
226 b.sl() as u64,
227 )
228 }
229 }
230
231 fn vec_znx_sub_ab_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
232 unsafe {
233 let res_ptr: *mut VecZnx = res as *mut VecZnx;
234 Self::vec_znx_sub(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
235 }
236 }
237
238 fn vec_znx_sub_ba_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
239 unsafe {
240 let res_ptr: *mut VecZnx = res as *mut VecZnx;
241 Self::vec_znx_sub(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
242 }
243 }
244
245 fn vec_znx_negate(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
246 #[cfg(debug_assertions)]
247 {
248 assert_eq!(a.n(), self.n());
249 assert_eq!(res.n(), self.n());
250 }
251 unsafe {
252 vec_znx::vec_znx_negate(
253 self.ptr,
254 res.at_mut_ptr(res_col, 0),
255 res.size() as u64,
256 res.sl() as u64,
257 a.at_ptr(a_col, 0),
258 a.size() as u64,
259 a.sl() as u64,
260 )
261 }
262 }
263
264 fn vec_znx_negate_inplace(&self, a: &mut VecZnx, a_col: usize) {
265 unsafe {
266 let a_ptr: *mut VecZnx = a as *mut VecZnx;
267 Self::vec_znx_negate(self, &mut *a_ptr, a_col, &*a_ptr, a_col);
268 }
269 }
270
271 fn vec_znx_rotate(&self, k: i64, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
272 #[cfg(debug_assertions)]
273 {
274 assert_eq!(a.n(), self.n());
275 assert_eq!(res.n(), self.n());
276 }
277 unsafe {
278 vec_znx::vec_znx_rotate(
279 self.ptr,
280 k,
281 res.at_mut_ptr(res_col, 0),
282 res.size() as u64,
283 res.sl() as u64,
284 a.at_ptr(a_col, 0),
285 a.size() as u64,
286 a.sl() as u64,
287 )
288 }
289 }
290
291 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx, a_col: usize) {
292 unsafe {
293 let a_ptr: *mut VecZnx = a as *mut VecZnx;
294 Self::vec_znx_rotate(self, k, &mut *a_ptr, a_col, &*a_ptr, a_col);
295 }
296 }
297
298 fn vec_znx_automorphism(&self, k: i64, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
299 #[cfg(debug_assertions)]
300 {
301 assert_eq!(a.n(), self.n());
302 assert_eq!(res.n(), self.n());
303 }
304 unsafe {
305 vec_znx::vec_znx_automorphism(
306 self.ptr,
307 k,
308 res.at_mut_ptr(res_col, 0),
309 res.size() as u64,
310 res.sl() as u64,
311 a.at_ptr(a_col, 0),
312 a.size() as u64,
313 a.sl() as u64,
314 )
315 }
316 }
317
318 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx, a_col: usize) {
319 unsafe {
320 let a_ptr: *mut VecZnx = a as *mut VecZnx;
321 Self::vec_znx_automorphism(self, k, &mut *a_ptr, a_col, &*a_ptr, a_col);
322 }
323 }
324
325 fn vec_znx_split(&self, res: &mut Vec<VecZnx>, res_col: usize, a: &VecZnx, a_col: usize, buf: &mut VecZnx) {
326 let (n_in, n_out) = (a.n(), res[0].n());
327
328 debug_assert!(
329 n_out < n_in,
330 "invalid a: output ring degree should be smaller"
331 );
332 res[1..].iter().for_each(|bi| {
333 debug_assert_eq!(
334 bi.n(),
335 n_out,
336 "invalid input a: all VecZnx must have the same degree"
337 )
338 });
339
340 res.iter_mut().enumerate().for_each(|(i, bi)| {
341 if i == 0 {
342 switch_degree(bi, res_col, a, a_col);
343 self.vec_znx_rotate(-1, buf, 0, a, a_col);
344 } else {
345 switch_degree(bi, res_col, buf, a_col);
346 self.vec_znx_rotate_inplace(-1, buf, a_col);
347 }
348 })
349 }
350
351 fn vec_znx_merge(&self, res: &mut VecZnx, res_col: usize, a: &Vec<VecZnx>, a_col: usize) {
352 let (n_in, n_out) = (res.n(), a[0].n());
353
354 debug_assert!(
355 n_out < n_in,
356 "invalid a: output ring degree should be smaller"
357 );
358 a[1..].iter().for_each(|ai| {
359 debug_assert_eq!(
360 ai.n(),
361 n_out,
362 "invalid input a: all VecZnx must have the same degree"
363 )
364 });
365
366 a.iter().enumerate().for_each(|(_, ai)| {
367 switch_degree(res, res_col, ai, a_col);
368 self.vec_znx_rotate_inplace(-1, res, res_col);
369 });
370
371 self.vec_znx_rotate_inplace(a.len() as i64, res, res_col);
372 }
373}