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