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