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