poulpy/base2k/src/vec_znx_big.rs
base2k/src/vec_znx_big.rsRaw
1use crate::ffi::vec_znx_big::{self, vec_znx_big_t};
2use crate::{Backend, FFT64, Infos, Module, VecZnx, VecZnxDft, alloc_aligned, assert_alignement};
3use std::marker::PhantomData;
4
5pub struct VecZnxBig<B: Backend> {
6 pub data: Vec<u8>,
7 pub ptr: *mut u8,
8 pub n: usize,
9 pub cols: usize,
10 pub limbs: usize,
11 pub _marker: PhantomData<B>,
12}
13
14impl VecZnxBig<FFT64> {
15
16 pub fn new(module: &Module<FFT64>, cols: usize, limbs: usize) -> Self {
17 #[cfg(debug_assertions)]
18 {
19 assert!(cols > 0);
20 assert!(limbs > 0);
21 }
22 let mut data: Vec<u8> = alloc_aligned::<u8>(module.bytes_of_vec_znx_big(cols, limbs));
23 let ptr: *mut u8 = data.as_mut_ptr();
24 Self {
25 data: data,
26 ptr: ptr,
27 n: module.n(),
28 cols: cols,
29 limbs: limbs,
30 _marker: PhantomData,
31 }
32 }
33
34 /// Returns a new [VecZnxBig] with the provided data as backing array.
35 /// User must ensure that data is properly alligned and that
36 /// the size of data is at least equal to [Module::bytes_of_vec_znx_big].
37 pub fn from_bytes(module: &Module<FFT64>, cols: usize, limbs: usize, bytes: &mut [u8]) -> Self {
38 #[cfg(debug_assertions)]
39 {
40 assert!(cols > 0);
41 assert!(limbs > 0);
42 assert_eq!(bytes.len(), module.bytes_of_vec_znx_big(cols, limbs));
43 assert_alignement(bytes.as_ptr())
44 };
45 unsafe {
46 Self {
47 data: Vec::from_raw_parts(bytes.as_mut_ptr(), bytes.len(), bytes.len()),
48 ptr: bytes.as_mut_ptr(),
49 n: module.n(),
50 cols: cols,
51 limbs: limbs,
52 _marker: PhantomData,
53 }
54 }
55 }
56
57 pub fn from_bytes_borrow(module: &Module<FFT64>, cols: usize, limbs: usize, bytes: &mut [u8]) -> Self {
58 #[cfg(debug_assertions)]
59 {
60 assert!(cols > 0);
61 assert!(limbs > 0);
62 assert_eq!(bytes.len(), module.bytes_of_vec_znx_big(cols, limbs));
63 assert_alignement(bytes.as_ptr());
64 }
65 Self {
66 data: Vec::new(),
67 ptr: bytes.as_mut_ptr(),
68 n: module.n(),
69 cols: cols,
70 limbs: limbs,
71 _marker: PhantomData,
72 }
73 }
74
75 pub fn as_vec_znx_dft(&mut self) -> VecZnxDft<FFT64> {
76 VecZnxDft::<FFT64> {
77 data: Vec::new(),
78 ptr: self.ptr,
79 n: self.n,
80 cols: self.cols,
81 limbs: self.limbs,
82 _marker: self._marker,
83 }
84 }
85
86 /// Returns a non-mutable pointer to the backedn slice of the receiver.
87 pub fn as_ptr(&self) -> *const i64 {
88 self.ptr as *const i64
89 }
90
91 /// Returns a mutable pointer to the backedn slice of the receiver.
92 pub fn as_mut_ptr(&mut self) -> *mut i64 {
93 self.ptr as *mut i64
94 }
95
96 pub fn raw(&self) -> &[i64] {
97 unsafe { &std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
98 }
99
100 pub fn raw_mut(&mut self) -> &mut [i64] {
101 let ptr: *mut i64 = self.ptr as *mut i64;
102 let size: usize = self.n() * self.poly_count();
103 unsafe { std::slice::from_raw_parts_mut(ptr, size) }
104 }
105
106 pub fn at_ptr(&self, i: usize, j: usize) -> *const i64 {
107 #[cfg(debug_assertions)]
108 {
109 assert!(i < self.cols());
110 assert!(j < self.limbs());
111 }
112 let offset: usize = self.n * (j * self.cols() + i);
113 self.as_ptr().wrapping_add(offset)
114 }
115
116 /// Returns a non-mutable reference to the i-th limb.
117 /// The returned array is of size [Self::n()] * [Self::cols()].
118 pub fn at_limb(&self, i: usize) -> &[i64] {
119 unsafe { std::slice::from_raw_parts(self.at_ptr(0, i), self.n * self.cols()) }
120 }
121
122 /// Returns a non-mutable reference to the (i, j)-th poly.
123 /// The returned array is of size [Self::n()].
124 pub fn at_poly(&self, i: usize, j: usize) -> &[i64] {
125 unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n) }
126 }
127
128 /// Returns a mutable pointer starting a the (i, j)-th small poly.
129 pub fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut i64 {
130 #[cfg(debug_assertions)]
131 {
132 assert!(i < self.cols());
133 assert!(j < self.limbs());
134 }
135
136 let offset: usize = self.n * (j * self.cols() + i);
137 self.as_mut_ptr().wrapping_add(offset)
138 }
139
140 /// Returns a mutable reference to the i-th limb.
141 /// The returned array is of size [Self::n()] * [Self::cols()].
142 pub fn at_limb_mut(&mut self, i: usize) -> &mut [i64] {
143 unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(0, i), self.n * self.cols()) }
144 }
145
146 /// Returns a mutable reference to the (i, j)-th poly.
147 /// The returned array is of size [Self::n()].
148 pub fn at_poly_mut(&mut self, i: usize, j: usize) -> &mut [i64] {
149 unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n) }
150 }
151
152 pub fn print(&self, n: usize) {
153 (0..self.limbs()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]));
154 }
155}
156
157impl<B: Backend> Infos for VecZnxBig<B> {
158 fn log_n(&self) -> usize {
159 (usize::BITS - (self.n - 1).leading_zeros()) as _
160 }
161
162 fn n(&self) -> usize {
163 self.n
164 }
165
166 fn cols(&self) -> usize {
167 self.cols
168 }
169
170 fn rows(&self) -> usize {
171 1
172 }
173
174 fn limbs(&self) -> usize {
175 self.limbs
176 }
177
178 fn poly_count(&self) -> usize {
179 self.cols * self.limbs
180 }
181}
182
183pub trait VecZnxBigOps<B: Backend> {
184 /// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
185 fn new_vec_znx_big(&self, cols: usize, limbs: usize) -> VecZnxBig<B>;
186
187 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
188 ///
189 /// Behavior: takes ownership of the backing array.
190 ///
191 /// # Arguments
192 ///
193 /// * `cols`: the number of polynomials..
194 /// * `limbs`: the number of limbs (a.k.a small polynomials) per polynomial.
195 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
196 ///
197 /// # Panics
198 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
199 fn new_vec_znx_big_from_bytes(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxBig<B>;
200
201 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
202 ///
203 /// Behavior: the backing array is only borrowed.
204 ///
205 /// # Arguments
206 ///
207 /// * `cols`: the number of polynomials..
208 /// * `limbs`: the number of limbs (a.k.a small polynomials) per polynomial.
209 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
210 ///
211 /// # Panics
212 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
213 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<B>;
214
215 /// Returns the minimum number of bytes necessary to allocate
216 /// a new [VecZnxBig] through [VecZnxBig::from_bytes].
217 fn bytes_of_vec_znx_big(&self, cols: usize, limbs: usize) -> usize;
218
219 /// b[VecZnxBig] <- b[VecZnxBig] - a[VecZnx]
220 ///
221 /// # Behavior
222 ///
223 /// [VecZnxBig] (3 cols and 4 limbs)
224 /// [a0, b0, c0] [a1, b1, c1] [a2, b2, c2] [a3, b3, c3]
225 /// -
226 /// [VecZnx] (2 cols and 3 limbs)
227 /// [d0, e0] [d1, e1] [d2, e2]
228 /// =
229 /// [a0-d0, b0-e0, c0] [a1-d1, b1-e1, c1] [a2-d2, b2-e2, c2] [a3, b3, c3]
230 fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
231
232 /// c <- b - a
233 fn vec_znx_big_sub_small_a(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
234
235 /// c <- b + a
236 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
237
238 /// b <- b + a
239 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
240
241 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
242
243 /// b <- normalize(a)
244 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<B>, tmp_bytes: &mut [u8]);
245
246 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize;
247
248 fn vec_znx_big_range_normalize_base2k(
249 &self,
250 log_base2k: usize,
251 res: &mut VecZnx,
252 a: &VecZnxBig<B>,
253 a_range_begin: usize,
254 a_range_xend: usize,
255 a_range_step: usize,
256 tmp_bytes: &mut [u8],
257 );
258
259 fn vec_znx_big_automorphism(&self, gal_el: i64, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
260
261 fn vec_znx_big_automorphism_inplace(&self, gal_el: i64, a: &mut VecZnxBig<B>);
262}
263
264impl VecZnxBigOps<FFT64> for Module<FFT64> {
265 fn new_vec_znx_big(&self, cols: usize, limbs: usize) -> VecZnxBig<FFT64> {
266 VecZnxBig::new(self, cols, limbs)
267 }
268
269 fn new_vec_znx_big_from_bytes(&self, cols: usize, limbs: usize, bytes: &mut [u8]) -> VecZnxBig<FFT64> {
270 VecZnxBig::from_bytes(self, cols, limbs, bytes)
271 }
272
273 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, limbs: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
274 VecZnxBig::from_bytes_borrow(self, cols, limbs, tmp_bytes)
275 }
276
277 fn bytes_of_vec_znx_big(&self, cols: usize, limbs: usize) -> usize {
278 unsafe { vec_znx_big::bytes_of_vec_znx_big(self.ptr, limbs as u64) as usize * cols }
279 }
280
281 /// [VecZnxBig] (3 cols and 4 limbs)
282 /// [a0, b0, c0] [a1, b1, c1] [a2, b2, c2] [a3, b3, c3]
283 /// -
284 /// [VecZnx] (2 cols and 3 limbs)
285 /// [d0, e0] [d1, e1] [d2, e2]
286 /// =
287 /// [a0-d0, b0-e0, c0] [a1-d1, b1-e1, c1] [a2-d2, b2-e2, c2] [a3, b3, c3]
288 fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
289 unsafe {
290 vec_znx_big::vec_znx_big_sub_small_a(
291 self.ptr,
292 b.ptr as *mut vec_znx_big_t,
293 b.poly_count() as u64,
294 a.as_ptr(),
295 a.poly_count() as u64,
296 a.n() as u64,
297 b.ptr as *mut vec_znx_big_t,
298 b.poly_count() as u64,
299 )
300 }
301 }
302
303 fn vec_znx_big_sub_small_a(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
304 unsafe {
305 vec_znx_big::vec_znx_big_sub_small_a(
306 self.ptr,
307 c.ptr as *mut vec_znx_big_t,
308 c.poly_count() as u64,
309 a.as_ptr(),
310 a.poly_count() as u64,
311 a.n() as u64,
312 b.ptr as *mut vec_znx_big_t,
313 b.poly_count() as u64,
314 )
315 }
316 }
317
318 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
319 unsafe {
320 vec_znx_big::vec_znx_big_add_small(
321 self.ptr,
322 c.ptr as *mut vec_znx_big_t,
323 c.poly_count() as u64,
324 b.ptr as *mut vec_znx_big_t,
325 b.poly_count() as u64,
326 a.as_ptr(),
327 a.poly_count() as u64,
328 a.n() as u64,
329 )
330 }
331 }
332
333 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
334 unsafe {
335 vec_znx_big::vec_znx_big_add_small(
336 self.ptr,
337 b.ptr as *mut vec_znx_big_t,
338 b.poly_count() as u64,
339 b.ptr as *mut vec_znx_big_t,
340 b.poly_count() as u64,
341 a.as_ptr(),
342 a.poly_count() as u64,
343 a.n() as u64,
344 )
345 }
346 }
347
348 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
349 unsafe { vec_znx_big::vec_znx_big_normalize_base2k_tmp_bytes(self.ptr) as usize }
350 }
351
352 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<FFT64>, tmp_bytes: &mut [u8]) {
353 debug_assert!(
354 tmp_bytes.len() >= Self::vec_znx_big_normalize_tmp_bytes(self),
355 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",
356 tmp_bytes.len(),
357 Self::vec_znx_big_normalize_tmp_bytes(self)
358 );
359 #[cfg(debug_assertions)]
360 {
361 assert_alignement(tmp_bytes.as_ptr())
362 }
363 unsafe {
364 vec_znx_big::vec_znx_big_normalize_base2k(
365 self.ptr,
366 log_base2k as u64,
367 b.as_mut_ptr(),
368 b.limbs() as u64,
369 b.n() as u64,
370 a.ptr as *mut vec_znx_big_t,
371 a.limbs() as u64,
372 tmp_bytes.as_mut_ptr(),
373 )
374 }
375 }
376
377 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize {
378 unsafe { vec_znx_big::vec_znx_big_range_normalize_base2k_tmp_bytes(self.ptr) as usize }
379 }
380
381 fn vec_znx_big_range_normalize_base2k(
382 &self,
383 log_base2k: usize,
384 res: &mut VecZnx,
385 a: &VecZnxBig<FFT64>,
386 a_range_begin: usize,
387 a_range_xend: usize,
388 a_range_step: usize,
389 tmp_bytes: &mut [u8],
390 ) {
391 debug_assert!(
392 tmp_bytes.len() >= Self::vec_znx_big_range_normalize_base2k_tmp_bytes(self),
393 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_range_normalize_base2k_tmp_bytes()={}",
394 tmp_bytes.len(),
395 Self::vec_znx_big_range_normalize_base2k_tmp_bytes(self)
396 );
397 #[cfg(debug_assertions)]
398 {
399 assert_alignement(tmp_bytes.as_ptr())
400 }
401 unsafe {
402 vec_znx_big::vec_znx_big_range_normalize_base2k(
403 self.ptr,
404 log_base2k as u64,
405 res.as_mut_ptr(),
406 res.limbs() as u64,
407 res.n() as u64,
408 a.ptr as *mut vec_znx_big_t,
409 a_range_begin as u64,
410 a_range_xend as u64,
411 a_range_step as u64,
412 tmp_bytes.as_mut_ptr(),
413 );
414 }
415 }
416
417 fn vec_znx_big_automorphism(&self, gal_el: i64, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
418 unsafe {
419 vec_znx_big::vec_znx_big_automorphism(
420 self.ptr,
421 gal_el,
422 b.ptr as *mut vec_znx_big_t,
423 b.poly_count() as u64,
424 a.ptr as *mut vec_znx_big_t,
425 a.poly_count() as u64,
426 );
427 }
428 }
429
430 fn vec_znx_big_automorphism_inplace(&self, gal_el: i64, a: &mut VecZnxBig<FFT64>) {
431 unsafe {
432 vec_znx_big::vec_znx_big_automorphism(
433 self.ptr,
434 gal_el,
435 a.ptr as *mut vec_znx_big_t,
436 a.poly_count() as u64,
437 a.ptr as *mut vec_znx_big_t,
438 a.poly_count() as u64,
439 );
440 }
441 }
442}