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