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, ZnxLayout, ZnxSliceSize};
3use crate::{Backend, FFT64, Module, VecZnx, VecZnxBig, VecZnxOps, assert_alignement};
4
5pub trait VecZnxBigOps<B: Backend> {
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) -> VecZnxBig<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>) -> VecZnxBig<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 /// Adds `a` to `b` and stores the result on `c`.
42 fn vec_znx_big_add(
43 &self,
44 res: &mut VecZnxBig<B>,
45 res_col: usize,
46 a: &VecZnxBig<B>,
47 a_col: usize,
48 b: &VecZnxBig<B>,
49 b_col: usize,
50 );
51
52 /// Adds `a` to `b` and stores the result on `b`.
53 fn vec_znx_big_add_inplace(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnxBig<B>, a_col: usize);
54
55 /// Adds `a` to `b` and stores the result on `c`.
56 fn vec_znx_big_add_small(
57 &self,
58 res: &mut VecZnxBig<B>,
59 res_col: usize,
60 a: &VecZnxBig<B>,
61 a_col: usize,
62 b: &VecZnx,
63 b_col: usize,
64 );
65
66 /// Adds `a` to `b` and stores the result on `b`.
67 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnx, a_col: usize);
68
69 /// Subtracts `a` to `b` and stores the result on `c`.
70 fn vec_znx_big_sub(
71 &self,
72 res: &mut VecZnxBig<B>,
73 res_col: usize,
74 a: &VecZnxBig<B>,
75 a_col: usize,
76 b: &VecZnxBig<B>,
77 b_col: usize,
78 );
79
80 /// Subtracts `a` to `b` and stores the result on `b`.
81 fn vec_znx_big_sub_ab_inplace(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnxBig<B>, a_col: usize);
82
83 /// Subtracts `b` to `a` and stores the result on `b`.
84 fn vec_znx_big_sub_ba_inplace(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnxBig<B>, a_col: usize);
85
86 /// Subtracts `b` to `a` and stores the result on `c`.
87 fn vec_znx_big_sub_small_a(
88 &self,
89 res: &mut VecZnxBig<B>,
90 res_col: usize,
91 a: &VecZnx,
92 a_col: usize,
93 b: &VecZnxBig<B>,
94 b_col: usize,
95 );
96
97 /// Subtracts `a` to `b` and stores the result on `b`.
98 fn vec_znx_big_sub_small_a_inplace(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnx, a_col: usize);
99
100 /// Subtracts `b` to `a` and stores the result on `c`.
101 fn vec_znx_big_sub_small_b(
102 &self,
103 res: &mut VecZnxBig<B>,
104 res_col: usize,
105 a: &VecZnxBig<B>,
106 a_col: usize,
107 b: &VecZnx,
108 b_col: usize,
109 );
110
111 /// Subtracts `b` to `a` and stores the result on `b`.
112 fn vec_znx_big_sub_small_b_inplace(&self, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnx, a_col: usize);
113
114 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
115 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
116
117 /// Normalizes `a` and stores the result on `b`.
118 ///
119 /// # Arguments
120 ///
121 /// * `log_base2k`: normalization basis.
122 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
123 fn vec_znx_big_normalize(
124 &self,
125 log_base2k: usize,
126 res: &mut VecZnx,
127 res_col: usize,
128 a: &VecZnxBig<B>,
129 a_col: usize,
130 tmp_bytes: &mut [u8],
131 );
132
133 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
134 fn vec_znx_big_automorphism(&self, k: i64, res: &mut VecZnxBig<B>, res_col: usize, a: &VecZnxBig<B>, a_col: usize);
135
136 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
137 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<B>, a_col: usize);
138}
139
140impl VecZnxBigOps<FFT64> for Module<FFT64> {
141 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBig<FFT64> {
142 VecZnxBig::new(self, 1, cols, size)
143 }
144
145 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBig<FFT64> {
146 VecZnxBig::from_bytes(self, 1, cols, size, bytes)
147 }
148
149 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
150 VecZnxBig::from_bytes_borrow(self, 1, cols, size, tmp_bytes)
151 }
152
153 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
154 VecZnxBig::bytes_of(self, 1, cols, size)
155 }
156
157 fn vec_znx_big_add(
158 &self,
159 res: &mut VecZnxBig<FFT64>,
160 res_col: usize,
161 a: &VecZnxBig<FFT64>,
162 a_col: usize,
163 b: &VecZnxBig<FFT64>,
164 b_col: usize,
165 ) {
166 #[cfg(debug_assertions)]
167 {
168 assert_eq!(a.n(), self.n());
169 assert_eq!(b.n(), self.n());
170 assert_eq!(res.n(), self.n());
171 assert_ne!(a.as_ptr(), b.as_ptr());
172 }
173 unsafe {
174 vec_znx::vec_znx_add(
175 self.ptr,
176 res.at_mut_ptr(res_col, 0),
177 res.size() as u64,
178 res.sl() as u64,
179 a.at_ptr(a_col, 0),
180 a.size() as u64,
181 a.sl() as u64,
182 b.at_ptr(b_col, 0),
183 b.size() as u64,
184 b.sl() as u64,
185 )
186 }
187 }
188
189 fn vec_znx_big_add_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
190 unsafe {
191 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
192 Self::vec_znx_big_add(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
193 }
194 }
195
196 fn vec_znx_big_sub(
197 &self,
198 res: &mut VecZnxBig<FFT64>,
199 res_col: usize,
200 a: &VecZnxBig<FFT64>,
201 a_col: usize,
202 b: &VecZnxBig<FFT64>,
203 b_col: usize,
204 ) {
205 #[cfg(debug_assertions)]
206 {
207 assert_eq!(a.n(), self.n());
208 assert_eq!(b.n(), self.n());
209 assert_eq!(res.n(), self.n());
210 assert_ne!(a.as_ptr(), b.as_ptr());
211 }
212 unsafe {
213 vec_znx::vec_znx_sub(
214 self.ptr,
215 res.at_mut_ptr(res_col, 0),
216 res.size() as u64,
217 res.sl() as u64,
218 a.at_ptr(a_col, 0),
219 a.size() as u64,
220 a.sl() as u64,
221 b.at_ptr(b_col, 0),
222 b.size() as u64,
223 b.sl() as u64,
224 )
225 }
226 }
227
228 fn vec_znx_big_sub_ab_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
229 unsafe {
230 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
231 Self::vec_znx_big_sub(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
232 }
233 }
234
235 fn vec_znx_big_sub_ba_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
236 unsafe {
237 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
238 Self::vec_znx_big_sub(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
239 }
240 }
241
242 fn vec_znx_big_sub_small_b(
243 &self,
244 res: &mut VecZnxBig<FFT64>,
245 res_col: usize,
246 a: &VecZnxBig<FFT64>,
247 a_col: usize,
248 b: &VecZnx,
249 b_col: usize,
250 ) {
251 #[cfg(debug_assertions)]
252 {
253 assert_eq!(a.n(), self.n());
254 assert_eq!(b.n(), self.n());
255 assert_eq!(res.n(), self.n());
256 assert_ne!(a.as_ptr(), b.as_ptr());
257 }
258 unsafe {
259 vec_znx::vec_znx_sub(
260 self.ptr,
261 res.at_mut_ptr(res_col, 0),
262 res.size() as u64,
263 res.sl() as u64,
264 a.at_ptr(a_col, 0),
265 a.size() as u64,
266 a.sl() as u64,
267 b.at_ptr(b_col, 0),
268 b.size() as u64,
269 b.sl() as u64,
270 )
271 }
272 }
273
274 fn vec_znx_big_sub_small_b_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
275 unsafe {
276 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
277 Self::vec_znx_big_sub_small_b(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
278 }
279 }
280
281 fn vec_znx_big_sub_small_a(
282 &self,
283 res: &mut VecZnxBig<FFT64>,
284 res_col: usize,
285 a: &VecZnx,
286 a_col: usize,
287 b: &VecZnxBig<FFT64>,
288 b_col: usize,
289 ) {
290 #[cfg(debug_assertions)]
291 {
292 assert_eq!(a.n(), self.n());
293 assert_eq!(b.n(), self.n());
294 assert_eq!(res.n(), self.n());
295 assert_ne!(a.as_ptr(), b.as_ptr());
296 }
297 unsafe {
298 vec_znx::vec_znx_sub(
299 self.ptr,
300 res.at_mut_ptr(res_col, 0),
301 res.size() as u64,
302 res.sl() as u64,
303 a.at_ptr(a_col, 0),
304 a.size() as u64,
305 a.sl() as u64,
306 b.at_ptr(b_col, 0),
307 b.size() as u64,
308 b.sl() as u64,
309 )
310 }
311 }
312
313 fn vec_znx_big_sub_small_a_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
314 unsafe {
315 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
316 Self::vec_znx_big_sub_small_a(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
317 }
318 }
319
320 fn vec_znx_big_add_small(
321 &self,
322 res: &mut VecZnxBig<FFT64>,
323 res_col: usize,
324 a: &VecZnxBig<FFT64>,
325 a_col: usize,
326 b: &VecZnx,
327 b_col: usize,
328 ) {
329 #[cfg(debug_assertions)]
330 {
331 assert_eq!(a.n(), self.n());
332 assert_eq!(b.n(), self.n());
333 assert_eq!(res.n(), self.n());
334 assert_ne!(a.as_ptr(), b.as_ptr());
335 }
336 unsafe {
337 vec_znx::vec_znx_add(
338 self.ptr,
339 res.at_mut_ptr(res_col, 0),
340 res.size() as u64,
341 res.sl() as u64,
342 a.at_ptr(a_col, 0),
343 a.size() as u64,
344 a.sl() as u64,
345 b.at_ptr(b_col, 0),
346 b.size() as u64,
347 b.sl() as u64,
348 )
349 }
350 }
351
352 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
353 unsafe {
354 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
355 Self::vec_znx_big_add_small(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
356 }
357 }
358
359 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
360 Self::vec_znx_normalize_tmp_bytes(self)
361 }
362
363 fn vec_znx_big_normalize(
364 &self,
365 log_base2k: usize,
366 res: &mut VecZnx,
367 res_col: usize,
368 a: &VecZnxBig<FFT64>,
369 a_col: usize,
370 tmp_bytes: &mut [u8],
371 ) {
372 #[cfg(debug_assertions)]
373 {
374 assert_eq!(a.n(), self.n());
375 assert_eq!(res.n(), self.n());
376 assert!(tmp_bytes.len() >= Self::vec_znx_normalize_tmp_bytes(&self));
377 assert_alignement(tmp_bytes.as_ptr());
378 }
379 unsafe {
380 vec_znx::vec_znx_normalize_base2k(
381 self.ptr,
382 log_base2k as u64,
383 res.at_mut_ptr(res_col, 0),
384 res.size() as u64,
385 res.sl() as u64,
386 a.at_ptr(a_col, 0),
387 a.size() as u64,
388 a.sl() as u64,
389 tmp_bytes.as_mut_ptr(),
390 );
391 }
392 }
393
394 fn vec_znx_big_automorphism(&self, k: i64, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
395 #[cfg(debug_assertions)]
396 {
397 assert_eq!(a.n(), self.n());
398 assert_eq!(res.n(), self.n());
399 }
400 unsafe {
401 vec_znx::vec_znx_automorphism(
402 self.ptr,
403 k,
404 res.at_mut_ptr(res_col, 0),
405 res.size() as u64,
406 res.sl() as u64,
407 a.at_ptr(a_col, 0),
408 a.size() as u64,
409 a.sl() as u64,
410 )
411 }
412 }
413
414 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<FFT64>, a_col: usize) {
415 unsafe {
416 let a_ptr: *mut VecZnxBig<FFT64> = a as *mut VecZnxBig<FFT64>;
417 Self::vec_znx_big_automorphism(self, k, &mut *a_ptr, a_col, &*a_ptr, a_col);
418 }
419 }
420}