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