poulpy/base2k/src/vec_znx_big_ops.rs
base2k/src/vec_znx_big_ops.rsRaw
1use crate::ffi::vec_znx;
2use crate::znx_base::{ZnxAlloc, ZnxInfos, ZnxView, ZnxViewMut};
3use crate::{Backend, DataView, FFT64, Module, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxOps, assert_alignement};
4
5pub trait VecZnxBigAlloc<B> {
6 /// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
7 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBigOwned<B>;
8
9 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
10 ///
11 /// Behavior: takes ownership of the backing array.
12 ///
13 /// # Arguments
14 ///
15 /// * `cols`: the number of polynomials..
16 /// * `size`: the number of polynomials per column.
17 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
18 ///
19 /// # Panics
20 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
21 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<B>;
22
23 // /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
24 // ///
25 // /// Behavior: the backing array is only borrowed.
26 // ///
27 // /// # Arguments
28 // ///
29 // /// * `cols`: the number of polynomials..
30 // /// * `size`: the number of polynomials per column.
31 // /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
32 // ///
33 // /// # Panics
34 // /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
35 // fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<B>;
36
37 /// Returns the minimum number of bytes necessary to allocate
38 /// a new [VecZnxBig] through [VecZnxBig::from_bytes].
39 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize;
40}
41
42pub trait VecZnxBigOps<DataMut, Data, B> {
43 /// Adds `a` to `b` and stores the result on `c`.
44 fn vec_znx_big_add(
45 &self,
46 res: &mut VecZnxBig<DataMut, B>,
47 res_col: usize,
48 a: &VecZnxBig<Data, B>,
49 a_col: usize,
50 b: &VecZnxBig<Data, B>,
51 b_col: usize,
52 );
53
54 /// Adds `a` to `b` and stores the result on `b`.
55 fn vec_znx_big_add_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnxBig<Data, B>, a_col: usize);
56
57 /// Adds `a` to `b` and stores the result on `c`.
58 fn vec_znx_big_add_small(
59 &self,
60 res: &mut VecZnxBig<DataMut, B>,
61 res_col: usize,
62 a: &VecZnxBig<Data, B>,
63 a_col: usize,
64 b: &VecZnx<Data>,
65 b_col: usize,
66 );
67
68 /// Adds `a` to `b` and stores the result on `b`.
69 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
70
71 /// Subtracts `a` to `b` and stores the result on `c`.
72 fn vec_znx_big_sub(
73 &self,
74 res: &mut VecZnxBig<DataMut, B>,
75 res_col: usize,
76 a: &VecZnxBig<Data, B>,
77 a_col: usize,
78 b: &VecZnxBig<Data, B>,
79 b_col: usize,
80 );
81
82 /// Subtracts `a` to `b` and stores the result on `b`.
83 fn vec_znx_big_sub_ab_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnxBig<Data, B>, a_col: usize);
84
85 /// Subtracts `b` to `a` and stores the result on `b`.
86 fn vec_znx_big_sub_ba_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnxBig<Data, B>, a_col: usize);
87
88 /// Subtracts `b` to `a` and stores the result on `c`.
89 fn vec_znx_big_sub_small_a(
90 &self,
91 res: &mut VecZnxBig<DataMut, B>,
92 res_col: usize,
93 a: &VecZnx<Data>,
94 a_col: usize,
95 b: &VecZnxBig<Data, B>,
96 b_col: usize,
97 );
98
99 /// Subtracts `a` to `b` and stores the result on `b`.
100 fn vec_znx_big_sub_small_a_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
101
102 /// Subtracts `b` to `a` and stores the result on `c`.
103 fn vec_znx_big_sub_small_b(
104 &self,
105 res: &mut VecZnxBig<DataMut, B>,
106 res_col: usize,
107 a: &VecZnxBig<Data, B>,
108 a_col: usize,
109 b: &VecZnx<Data>,
110 b_col: usize,
111 );
112
113 /// Subtracts `b` to `a` and stores the result on `b`.
114 fn vec_znx_big_sub_small_b_inplace(&self, res: &mut VecZnxBig<DataMut, B>, res_col: usize, a: &VecZnx<Data>, a_col: usize);
115
116 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
117 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
118
119 /// Normalizes `a` and stores the result on `b`.
120 ///
121 /// # Arguments
122 ///
123 /// * `log_base2k`: normalization basis.
124 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
125 fn vec_znx_big_normalize(
126 &self,
127 log_base2k: usize,
128 res: &mut VecZnx<DataMut>,
129 res_col: usize,
130 a: &VecZnxBig<Data, B>,
131 a_col: usize,
132 tmp_bytes: &mut [u8],
133 );
134
135 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
136 fn vec_znx_big_automorphism(
137 &self,
138 k: i64,
139 res: &mut VecZnxBig<DataMut, B>,
140 res_col: usize,
141 a: &VecZnxBig<Data, B>,
142 a_col: usize,
143 );
144
145 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
146 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<DataMut, B>, a_col: usize);
147}
148
149impl VecZnxBigAlloc<FFT64> for Module<FFT64> {
150 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBigOwned<FFT64> {
151 VecZnxBig::new(self, cols, size)
152 }
153
154 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<FFT64> {
155 VecZnxBig::new_from_bytes(self, cols, size, bytes)
156 }
157
158 // fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
159 // VecZnxBig::from_bytes_borrow(self, 1, cols, size, tmp_bytes)
160 // }
161
162 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
163 VecZnxBig::bytes_of(self, cols, size)
164 }
165}
166
167impl<DataMut, Data> VecZnxBigOps<DataMut, Data, FFT64> for Module<FFT64>
168where
169 DataMut: AsMut<[u8]> + AsRef<[u8]>,
170 Data: AsRef<[u8]>,
171{
172 fn vec_znx_big_add(
173 &self,
174 res: &mut VecZnxBig<DataMut, FFT64>,
175 res_col: usize,
176 a: &VecZnxBig<Data, FFT64>,
177 a_col: usize,
178 b: &VecZnxBig<Data, FFT64>,
179 b_col: usize,
180 ) {
181 #[cfg(debug_assertions)]
182 {
183 assert_eq!(a.n(), self.n());
184 assert_eq!(b.n(), self.n());
185 assert_eq!(res.n(), self.n());
186 assert_ne!(a.as_ptr(), b.as_ptr());
187 }
188 unsafe {
189 vec_znx::vec_znx_add(
190 self.ptr,
191 res.at_mut_ptr(res_col, 0),
192 res.size() as u64,
193 res.sl() as u64,
194 a.at_ptr(a_col, 0),
195 a.size() as u64,
196 a.sl() as u64,
197 b.at_ptr(b_col, 0),
198 b.size() as u64,
199 b.sl() as u64,
200 )
201 }
202 }
203
204 fn vec_znx_big_add_inplace(
205 &self,
206 res: &mut VecZnxBig<DataMut, FFT64>,
207 res_col: usize,
208 a: &VecZnxBig<Data, FFT64>,
209 a_col: usize,
210 ) {
211 unsafe {
212 Self::vec_znx_big_add(self, res, res_col, a, a_col, res, res_col);
213 }
214 }
215
216 fn vec_znx_big_sub(
217 &self,
218 res: &mut VecZnxBig<DataMut, FFT64>,
219 res_col: usize,
220 a: &VecZnxBig<Data, FFT64>,
221 a_col: usize,
222 b: &VecZnxBig<Data, FFT64>,
223 b_col: usize,
224 ) {
225 #[cfg(debug_assertions)]
226 {
227 assert_eq!(a.n(), self.n());
228 assert_eq!(b.n(), self.n());
229 assert_eq!(res.n(), self.n());
230 assert_ne!(a.as_ptr(), b.as_ptr());
231 }
232 unsafe {
233 vec_znx::vec_znx_sub(
234 self.ptr,
235 res.at_mut_ptr(res_col, 0),
236 res.size() as u64,
237 res.sl() as u64,
238 a.at_ptr(a_col, 0),
239 a.size() as u64,
240 a.sl() as u64,
241 b.at_ptr(b_col, 0),
242 b.size() as u64,
243 b.sl() as u64,
244 )
245 }
246 }
247
248 //(Jay)TODO: check whether definitions sub_ab, sub_ba make sense to you
249 fn vec_znx_big_sub_ab_inplace(
250 &self,
251 res: &mut VecZnxBig<DataMut, FFT64>,
252 res_col: usize,
253 a: &VecZnxBig<Data, FFT64>,
254 a_col: usize,
255 ) {
256 unsafe {
257 Self::vec_znx_big_sub(self, res, res_col, a, a_col, res, res_col);
258 }
259 }
260
261 fn vec_znx_big_sub_ba_inplace(
262 &self,
263 res: &mut VecZnxBig<DataMut, FFT64>,
264 res_col: usize,
265 a: &VecZnxBig<Data, FFT64>,
266 a_col: usize,
267 ) {
268 unsafe {
269 Self::vec_znx_big_sub(self, res, res_col, res, res_col, a, a_col);
270 }
271 }
272
273 fn vec_znx_big_sub_small_b(
274 &self,
275 res: &mut VecZnxBig<DataMut, FFT64>,
276 res_col: usize,
277 a: &VecZnxBig<Data, FFT64>,
278 a_col: usize,
279 b: &VecZnx<Data>,
280 b_col: usize,
281 ) {
282 #[cfg(debug_assertions)]
283 {
284 assert_eq!(a.n(), self.n());
285 assert_eq!(b.n(), self.n());
286 assert_eq!(res.n(), self.n());
287 assert_ne!(a.as_ptr(), b.as_ptr());
288 }
289 unsafe {
290 vec_znx::vec_znx_sub(
291 self.ptr,
292 res.at_mut_ptr(res_col, 0),
293 res.size() as u64,
294 res.sl() as u64,
295 a.at_ptr(a_col, 0),
296 a.size() as u64,
297 a.sl() as u64,
298 b.at_ptr(b_col, 0),
299 b.size() as u64,
300 b.sl() as u64,
301 )
302 }
303 }
304
305 fn vec_znx_big_sub_small_b_inplace(
306 &self,
307 res: &mut VecZnxBig<DataMut, FFT64>,
308 res_col: usize,
309 a: &VecZnx<Data>,
310 a_col: usize,
311 ) {
312 unsafe {
313 Self::vec_znx_big_sub_small_b(self, res, res_col, res, res_col, a, a_col);
314 }
315 }
316
317 fn vec_znx_big_sub_small_a(
318 &self,
319 res: &mut VecZnxBig<DataMut, FFT64>,
320 res_col: usize,
321 a: &VecZnx<Data>,
322 a_col: usize,
323 b: &VecZnxBig<Data, FFT64>,
324 b_col: usize,
325 ) {
326 #[cfg(debug_assertions)]
327 {
328 assert_eq!(a.n(), self.n());
329 assert_eq!(b.n(), self.n());
330 assert_eq!(res.n(), self.n());
331 assert_ne!(a.as_ptr(), b.as_ptr());
332 }
333 unsafe {
334 vec_znx::vec_znx_sub(
335 self.ptr,
336 res.at_mut_ptr(res_col, 0),
337 res.size() as u64,
338 res.sl() as u64,
339 a.at_ptr(a_col, 0),
340 a.size() as u64,
341 a.sl() as u64,
342 b.at_ptr(b_col, 0),
343 b.size() as u64,
344 b.sl() as u64,
345 )
346 }
347 }
348
349 fn vec_znx_big_sub_small_a_inplace(
350 &self,
351 res: &mut VecZnxBig<DataMut, FFT64>,
352 res_col: usize,
353 a: &VecZnx<Data>,
354 a_col: usize,
355 ) {
356 unsafe {
357 Self::vec_znx_big_sub_small_a(self, res, res_col, a, a_col, res, res_col);
358 }
359 }
360
361 fn vec_znx_big_add_small(
362 &self,
363 res: &mut VecZnxBig<DataMut, FFT64>,
364 res_col: usize,
365 a: &VecZnxBig<Data, FFT64>,
366 a_col: usize,
367 b: &VecZnx<Data>,
368 b_col: usize,
369 ) {
370 #[cfg(debug_assertions)]
371 {
372 assert_eq!(a.n(), self.n());
373 assert_eq!(b.n(), self.n());
374 assert_eq!(res.n(), self.n());
375 assert_ne!(a.as_ptr(), b.as_ptr());
376 }
377 unsafe {
378 vec_znx::vec_znx_add(
379 self.ptr,
380 res.at_mut_ptr(res_col, 0),
381 res.size() as u64,
382 res.sl() as u64,
383 a.at_ptr(a_col, 0),
384 a.size() as u64,
385 a.sl() as u64,
386 b.at_ptr(b_col, 0),
387 b.size() as u64,
388 b.sl() as u64,
389 )
390 }
391 }
392
393 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<DataMut, FFT64>, res_col: usize, a: &VecZnx<Data>, a_col: usize) {
394 Self::vec_znx_big_add_small(self, res, res_col, res, res_col, a, a_col);
395 }
396
397 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
398 Self::vec_znx_normalize_tmp_bytes(self)
399 }
400
401 fn vec_znx_big_normalize(
402 &self,
403 log_base2k: usize,
404 res: &mut VecZnx<DataMut>,
405 res_col: usize,
406 a: &VecZnxBig<Data, FFT64>,
407 a_col: usize,
408 tmp_bytes: &mut [u8],
409 ) {
410 #[cfg(debug_assertions)]
411 {
412 assert_eq!(a.n(), self.n());
413 assert_eq!(res.n(), self.n());
414 assert!(tmp_bytes.len() >= Self::vec_znx_normalize_tmp_bytes(&self));
415 assert_alignement(tmp_bytes.as_ptr());
416 }
417 unsafe {
418 vec_znx::vec_znx_normalize_base2k(
419 self.ptr,
420 log_base2k as u64,
421 res.at_mut_ptr(res_col, 0),
422 res.size() as u64,
423 res.sl() as u64,
424 a.at_ptr(a_col, 0),
425 a.size() as u64,
426 a.sl() as u64,
427 tmp_bytes.as_mut_ptr(),
428 );
429 }
430 }
431
432 fn vec_znx_big_automorphism(
433 &self,
434 k: i64,
435 res: &mut VecZnxBig<DataMut, FFT64>,
436 res_col: usize,
437 a: &VecZnxBig<Data, FFT64>,
438 a_col: usize,
439 ) {
440 #[cfg(debug_assertions)]
441 {
442 assert_eq!(a.n(), self.n());
443 assert_eq!(res.n(), self.n());
444 }
445 unsafe {
446 vec_znx::vec_znx_automorphism(
447 self.ptr,
448 k,
449 res.at_mut_ptr(res_col, 0),
450 res.size() as u64,
451 res.sl() as u64,
452 a.at_ptr(a_col, 0),
453 a.size() as u64,
454 a.sl() as u64,
455 )
456 }
457 }
458
459 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<DataMut, FFT64>, a_col: usize) {
460 unsafe {
461 Self::vec_znx_big_automorphism(self, k, a, a_col, a, a_col);
462 }
463 }
464}