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