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 fn vec_znx_sub_ab_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
90
91 // /// Subtracts the selected column of `a` from the selected column of `res` and negates the selected column of `res`.
92 // fn vec_znx_sub_ba_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
93
94 // Negates the selected column of `a` and stores the result in `res_col` of `res`.
95 fn vec_znx_negate(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
96
97 /// Negates the selected column of `a`.
98 fn vec_znx_negate_inplace(&self, a: &mut VecZnx<DataMut>, a_col: usize);
99
100 /// Multiplies the selected column of `a` by X^k and stores the result in `res_col` of `res`.
101 fn vec_znx_rotate(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
102
103 /// Multiplies the selected column of `a` by X^k.
104 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize);
105
106 /// Applies the automorphism X^i -> X^ik on the selected column of `a` and stores the result in `res_col` column of `res`.
107 fn vec_znx_automorphism(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
108
109 /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
110 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize);
111
112 /// Splits the selected columns of `b` into subrings and copies them them into the selected column of `res`.
113 ///
114 /// # Panics
115 ///
116 /// This method requires that all [VecZnx] of b have the same ring degree
117 /// and that b.n() * b.len() <= a.n()
118 fn vec_znx_split(
119 &self,
120 res: &mut Vec<VecZnx<DataMut>>,
121 res_col: usize,
122 a: &VecZnx<Data>,
123 a_col: usize,
124 buf: &mut VecZnx<DataMut>,
125 );
126
127 /// Merges the subrings of the selected column of `a` into the selected column of `res`.
128 ///
129 /// # Panics
130 ///
131 /// This method requires that all [VecZnx] of a have the same ring degree
132 /// and that a.n() * a.len() <= b.n()
133 fn vec_znx_merge(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &Vec<VecZnx<Data>>, a_col: usize);
134}
135
136impl<B: Backend> VecZnxAlloc for Module<B> {
137 //(Jay)TODO: One must define the Scalar generic param here.
138 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnxOwned {
139 VecZnxOwned::new(self.n(), cols, size)
140 }
141
142 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize {
143 VecZnxOwned::bytes_of(self.n(), cols, size)
144 }
145
146 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxOwned {
147 VecZnxOwned::new_from_bytes(self.n(), cols, size, bytes)
148 }
149}
150
151impl<B: Backend, DataMut, Data> VecZnxOps<DataMut, Data> for Module<B>
152where
153 Data: AsRef<[u8]>,
154 DataMut: AsRef<[u8]> + AsMut<[u8]>,
155{
156 fn vec_znx_normalize_tmp_bytes(&self) -> usize {
157 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
158 }
159
160 fn vec_znx_normalize(
161 &self,
162 log_base2k: usize,
163 res: &mut VecZnx<DataMut>,
164 res_col: usize,
165 a: &VecZnx<Data>,
166 a_col: usize,
167 tmp_bytes: &mut [u8],
168 ) {
169 #[cfg(debug_assertions)]
170 {
171 assert_eq!(a.n(), self.n());
172 assert_eq!(res.n(), self.n());
173 assert!(tmp_bytes.len() >= Self::vec_znx_normalize_tmp_bytes(&self));
174 assert_alignement(tmp_bytes.as_ptr());
175 }
176 unsafe {
177 vec_znx::vec_znx_normalize_base2k(
178 self.ptr,
179 log_base2k as u64,
180 res.at_mut_ptr(res_col, 0),
181 res.size() as u64,
182 res.sl() as u64,
183 a.at_ptr(a_col, 0),
184 a.size() as u64,
185 a.sl() as u64,
186 tmp_bytes.as_mut_ptr(),
187 );
188 }
189 }
190
191 fn vec_znx_normalize_inplace(&self, log_base2k: usize, a: &mut VecZnx<DataMut>, a_col: usize, tmp_bytes: &mut [u8]) {
192 unsafe {
193 let a_ptr: *mut VecZnx = a as *mut VecZnx;
194 Self::vec_znx_normalize(
195 self,
196 log_base2k,
197 &mut *a_ptr,
198 a_col,
199 &*a_ptr,
200 a_col,
201 tmp_bytes,
202 );
203 }
204 }
205
206 fn vec_znx_add(
207 &self,
208 res: &mut VecZnx<DataMut>,
209 res_col: usize,
210 a: &VecZnx<Data>,
211 a_col: usize,
212 b: &VecZnx<Data>,
213 b_col: usize,
214 ) {
215 #[cfg(debug_assertions)]
216 {
217 assert_eq!(a.n(), self.n());
218 assert_eq!(b.n(), self.n());
219 assert_eq!(res.n(), self.n());
220 assert_ne!(a.as_ptr(), b.as_ptr());
221 }
222 unsafe {
223 vec_znx::vec_znx_add(
224 self.ptr,
225 res.at_mut_ptr(res_col, 0),
226 res.size() as u64,
227 res.sl() as u64,
228 a.at_ptr(a_col, 0),
229 a.size() as u64,
230 a.sl() as u64,
231 b.at_ptr(b_col, 0),
232 b.size() as u64,
233 b.sl() as u64,
234 )
235 }
236 }
237
238 fn vec_znx_add_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
239 unsafe {
240 Self::vec_znx_add(&self, res, res_col, a, a_col, res, res_col);
241 }
242 }
243
244 fn vec_znx_sub(
245 &self,
246 res: &mut VecZnx<DataMut>,
247 res_col: usize,
248 a: &VecZnx<Data>,
249 a_col: usize,
250 b: &VecZnx<Data>,
251 b_col: usize,
252 ) {
253 #[cfg(debug_assertions)]
254 {
255 assert_eq!(a.n(), self.n());
256 assert_eq!(b.n(), self.n());
257 assert_eq!(res.n(), self.n());
258 assert_ne!(a.as_ptr(), b.as_ptr());
259 }
260 unsafe {
261 vec_znx::vec_znx_sub(
262 self.ptr,
263 res.at_mut_ptr(res_col, 0),
264 res.size() as u64,
265 res.sl() as u64,
266 a.at_ptr(a_col, 0),
267 a.size() as u64,
268 a.sl() as u64,
269 b.at_ptr(b_col, 0),
270 b.size() as u64,
271 b.sl() as u64,
272 )
273 }
274 }
275
276 fn vec_znx_sub_ab_inplace(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
277 unsafe {
278 let res_ptr: *mut VecZnx = res as *mut VecZnx;
279 Self::vec_znx_sub(self, res, res_col, a, a_col, res, res_col);
280 }
281 }
282
283 // fn vec_znx_sub_ba_inplace(&self, res: &mut VecZnx, res_col: usize, a: &VecZnx, a_col: usize) {
284 // unsafe {
285 // let res_ptr: *mut VecZnx = res as *mut VecZnx;
286 // Self::vec_znx_sub(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
287 // }
288 // }
289
290 fn vec_znx_negate(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
291 #[cfg(debug_assertions)]
292 {
293 assert_eq!(a.n(), self.n());
294 assert_eq!(res.n(), self.n());
295 }
296 unsafe {
297 vec_znx::vec_znx_negate(
298 self.ptr,
299 res.at_mut_ptr(res_col, 0),
300 res.size() as u64,
301 res.sl() as u64,
302 a.at_ptr(a_col, 0),
303 a.size() as u64,
304 a.sl() as u64,
305 )
306 }
307 }
308
309 fn vec_znx_negate_inplace(&self, a: &mut VecZnx<DataMut>, a_col: usize) {
310 unsafe {
311 Self::vec_znx_negate(self, a, a_col, a, a_col);
312 }
313 }
314
315 fn vec_znx_rotate(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
316 #[cfg(debug_assertions)]
317 {
318 assert_eq!(a.n(), self.n());
319 assert_eq!(res.n(), self.n());
320 }
321 unsafe {
322 vec_znx::vec_znx_rotate(
323 self.ptr,
324 k,
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_rotate_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize) {
336 unsafe {
337 Self::vec_znx_rotate(self, k, a, a_col, a, a_col);
338 }
339 }
340
341 fn vec_znx_automorphism(&self, k: i64, res: &mut VecZnx<DataMut>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
342 #[cfg(debug_assertions)]
343 {
344 assert_eq!(a.n(), self.n());
345 assert_eq!(res.n(), self.n());
346 }
347 unsafe {
348 vec_znx::vec_znx_automorphism(
349 self.ptr,
350 k,
351 res.at_mut_ptr(res_col, 0),
352 res.size() as u64,
353 res.sl() as u64,
354 a.at_ptr(a_col, 0),
355 a.size() as u64,
356 a.sl() as u64,
357 )
358 }
359 }
360
361 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx<DataMut>, a_col: usize) {
362 unsafe {
363 Self::vec_znx_automorphism(self, k, a, a_col, a, a_col);
364 }
365 }
366
367 fn vec_znx_split(
368 &self,
369 res: &mut Vec<VecZnx<DataMut>>,
370 res_col: usize,
371 a: &VecZnx<Data>,
372 a_col: usize,
373 buf: &mut VecZnx<DataMut>,
374 ) {
375 let (n_in, n_out) = (a.n(), res[0].n());
376
377 debug_assert!(
378 n_out < n_in,
379 "invalid a: output ring degree should be smaller"
380 );
381 res[1..].iter().for_each(|bi| {
382 debug_assert_eq!(
383 bi.n(),
384 n_out,
385 "invalid input a: all VecZnx must have the same degree"
386 )
387 });
388
389 res.iter_mut().enumerate().for_each(|(i, bi)| {
390 if i == 0 {
391 switch_degree(bi, res_col, a, a_col);
392 self.vec_znx_rotate(-1, buf, 0, a, a_col);
393 } else {
394 switch_degree(bi, res_col, buf, a_col);
395 self.vec_znx_rotate_inplace(-1, buf, a_col);
396 }
397 })
398 }
399
400 fn vec_znx_merge(&self, res: &mut VecZnx<DataMut>, res_col: usize, a: &Vec<VecZnx<Data>>, a_col: usize) {
401 let (n_in, n_out) = (res.n(), a[0].n());
402
403 debug_assert!(
404 n_out < n_in,
405 "invalid a: output ring degree should be smaller"
406 );
407 a[1..].iter().for_each(|ai| {
408 debug_assert_eq!(
409 ai.n(),
410 n_out,
411 "invalid input a: all VecZnx must have the same degree"
412 )
413 });
414
415 a.iter().enumerate().for_each(|(_, ai)| {
416 switch_degree(res, res_col, ai, a_col);
417 self.vec_znx_rotate_inplace(-1, res, res_col);
418 });
419
420 self.vec_znx_rotate_inplace(a.len() as i64, res, res_col);
421 }
422}