poulpy/base2k/src/vec_znx_big_ops.rs
base2k/src/vec_znx_big_ops.rsRaw
1use crate::ffi::vec_znx_big::{self, vec_znx_big_t};
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_big::vec_znx_big_add(
175 self.ptr,
176 res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big_t,
177 res.size() as u64,
178 a.at_ptr(a_col * res.size(), 0) as *const vec_znx_big_t,
179 a.size() as u64,
180 b.at_ptr(b_col * res.size(), 0) as *const vec_znx_big_t,
181 b.size() as u64,
182 )
183 }
184 }
185
186 fn vec_znx_big_add_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
187 unsafe {
188 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
189 Self::vec_znx_big_add(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
190 }
191 }
192
193 fn vec_znx_big_sub(
194 &self,
195 res: &mut VecZnxBig<FFT64>,
196 res_col: usize,
197 a: &VecZnxBig<FFT64>,
198 a_col: usize,
199 b: &VecZnxBig<FFT64>,
200 b_col: usize,
201 ) {
202 #[cfg(debug_assertions)]
203 {
204 assert_eq!(a.n(), self.n());
205 assert_eq!(b.n(), self.n());
206 assert_eq!(res.n(), self.n());
207 assert_ne!(a.as_ptr(), b.as_ptr());
208 }
209 unsafe {
210 vec_znx_big::vec_znx_big_sub(
211 self.ptr,
212 res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big_t,
213 res.size() as u64,
214 a.at_ptr(a_col * res.size(), 0) as *const vec_znx_big_t,
215 a.size() as u64,
216 b.at_ptr(b_col * res.size(), 0) as *const vec_znx_big_t,
217 b.size() as u64,
218 )
219 }
220 }
221
222 fn vec_znx_big_sub_ab_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
223 unsafe {
224 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
225 Self::vec_znx_big_sub(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
226 }
227 }
228
229 fn vec_znx_big_sub_ba_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
230 unsafe {
231 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
232 Self::vec_znx_big_sub(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
233 }
234 }
235
236 fn vec_znx_big_sub_small_b(
237 &self,
238 res: &mut VecZnxBig<FFT64>,
239 res_col: usize,
240 a: &VecZnxBig<FFT64>,
241 a_col: usize,
242 b: &VecZnx,
243 b_col: usize,
244 ) {
245 #[cfg(debug_assertions)]
246 {
247 assert_eq!(a.n(), self.n());
248 assert_eq!(b.n(), self.n());
249 assert_eq!(res.n(), self.n());
250 assert_ne!(a.as_ptr(), b.as_ptr());
251 }
252 unsafe {
253 vec_znx_big::vec_znx_big_sub_small_b(
254 self.ptr,
255 res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big_t,
256 res.size() as u64,
257 a.at_ptr(a_col * a.size(), 0) as *const vec_znx_big_t,
258 a.size() as u64,
259 b.at_ptr(b_col, 0),
260 b.size() as u64,
261 b.sl() as u64,
262 )
263 }
264 }
265
266 fn vec_znx_big_sub_small_b_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
267 unsafe {
268 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
269 Self::vec_znx_big_sub_small_b(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
270 }
271 }
272
273 fn vec_znx_big_sub_small_a(
274 &self,
275 res: &mut VecZnxBig<FFT64>,
276 res_col: usize,
277 a: &VecZnx,
278 a_col: usize,
279 b: &VecZnxBig<FFT64>,
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_big::vec_znx_big_sub_small_a(
291 self.ptr,
292 res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big_t,
293 res.size() as u64,
294 a.at_ptr(a_col, 0),
295 a.size() as u64,
296 a.sl() as u64,
297 b.at_ptr(b_col * b.size(), 0) as *const vec_znx_big_t,
298 b.size() as u64,
299 )
300 }
301 }
302
303 fn vec_znx_big_sub_small_a_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
304 unsafe {
305 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
306 Self::vec_znx_big_sub_small_a(self, &mut *res_ptr, res_col, a, a_col, &*res_ptr, res_col);
307 }
308 }
309
310 fn vec_znx_big_add_small(
311 &self,
312 res: &mut VecZnxBig<FFT64>,
313 res_col: usize,
314 a: &VecZnxBig<FFT64>,
315 a_col: usize,
316 b: &VecZnx,
317 b_col: usize,
318 ) {
319 #[cfg(debug_assertions)]
320 {
321 assert_eq!(a.n(), self.n());
322 assert_eq!(b.n(), self.n());
323 assert_eq!(res.n(), self.n());
324 assert_ne!(a.as_ptr(), b.as_ptr());
325 }
326 unsafe {
327 vec_znx_big::vec_znx_big_add_small(
328 self.ptr,
329 res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big_t,
330 res.size() as u64,
331 a.at_ptr(a_col * a.size(), 0) as *const vec_znx_big_t,
332 a.size() as u64,
333 b.at_ptr(b_col, 0),
334 b.size() as u64,
335 b.sl() as u64,
336 )
337 }
338 }
339
340 fn vec_znx_big_add_small_inplace(&self, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnx, a_col: usize) {
341 unsafe {
342 let res_ptr: *mut VecZnxBig<FFT64> = res as *mut VecZnxBig<FFT64>;
343 Self::vec_znx_big_add_small(self, &mut *res_ptr, res_col, &*res_ptr, res_col, a, a_col);
344 }
345 }
346
347 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
348 Self::vec_znx_normalize_tmp_bytes(self)
349 }
350
351 fn vec_znx_big_normalize(
352 &self,
353 log_base2k: usize,
354 res: &mut VecZnx,
355 res_col: usize,
356 a: &VecZnxBig<FFT64>,
357 a_col: usize,
358 tmp_bytes: &mut [u8],
359 ) {
360 #[cfg(debug_assertions)]
361 {
362 assert_eq!(a.n(), self.n());
363 assert_eq!(res.n(), self.n());
364 assert!(tmp_bytes.len() >= Self::vec_znx_normalize_tmp_bytes(&self));
365 assert_alignement(tmp_bytes.as_ptr());
366 }
367 unsafe {
368 vec_znx_big::vec_znx_big_normalize_base2k(
369 self.ptr,
370 log_base2k as u64,
371 res.at_mut_ptr(res_col, 0),
372 res.size() as u64,
373 res.sl() as u64,
374 a.at_ptr(a_col * a.size(), 0) as *const vec_znx_big_t,
375 a.size() as u64,
376 tmp_bytes.as_mut_ptr(),
377 );
378 }
379 }
380
381 fn vec_znx_big_automorphism(&self, k: i64, res: &mut VecZnxBig<FFT64>, res_col: usize, a: &VecZnxBig<FFT64>, a_col: usize) {
382 #[cfg(debug_assertions)]
383 {
384 assert_eq!(a.n(), self.n());
385 assert_eq!(res.n(), self.n());
386 }
387 unsafe {
388 vec_znx_big::vec_znx_big_automorphism(
389 self.ptr,
390 k,
391 res.at_mut_ptr(res_col * res.size(), 0) as *mut vec_znx_big_t,
392 res.size() as u64,
393 a.at_ptr(a_col * a.size(), 0) as *const vec_znx_big_t,
394 a.size() as u64,
395 )
396 }
397 }
398
399 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<FFT64>, a_col: usize) {
400 unsafe {
401 let a_ptr: *mut VecZnxBig<FFT64> = a as *mut VecZnxBig<FFT64>;
402 Self::vec_znx_big_automorphism(self, k, &mut *a_ptr, a_col, &*a_ptr, a_col);
403 }
404 }
405}