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(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 polynomials per column.
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 polynomials per column.
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 /// Subtracts `a` to `b` and stores the result on `b`.
153 fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
154
155 /// Subtracts `b` to `a` and stores the result on `c`.
156 fn vec_znx_big_sub_small_a(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
157
158 /// Adds `a` to `b` and stores the result on `c`.
159 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
160
161 /// Adds `a` to `b` and stores the result on `b`.
162 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
163
164 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
165 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
166
167 /// Normalizes `a` and stores the result on `b`.
168 ///
169 /// # Arguments
170 ///
171 /// * `log_base2k`: normalization basis.
172 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
173 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<B>, tmp_bytes: &mut [u8]);
174
175 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_range_normalize_base2k].
176 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize;
177
178 /// Normalize `a`, taking into account column interleaving and stores the result on `b`.
179 ///
180 /// # Arguments
181 ///
182 /// * `log_base2k`: normalization basis.
183 /// * `a_range_begin`: column to start.
184 /// * `a_range_end`: column to end.
185 /// * `a_range_step`: column step size.
186 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_range_normalize_base2k_tmp_bytes].
187 fn vec_znx_big_range_normalize_base2k(
188 &self,
189 log_base2k: usize,
190 b: &mut VecZnx,
191 a: &VecZnxBig<B>,
192 a_range_begin: usize,
193 a_range_xend: usize,
194 a_range_step: usize,
195 tmp_bytes: &mut [u8],
196 );
197
198 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
199 fn vec_znx_big_automorphism(&self, k: i64, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
200
201 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
202 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<B>);
203}
204
205impl VecZnxBigOps<FFT64> for Module<FFT64> {
206 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBig<FFT64> {
207 VecZnxBig::new(self, cols, size)
208 }
209
210 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxBig<FFT64> {
211 VecZnxBig::from_bytes(self, cols, size, bytes)
212 }
213
214 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
215 VecZnxBig::from_bytes_borrow(self, cols, size, tmp_bytes)
216 }
217
218 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
219 VecZnxBig::bytes_of(self, cols, size)
220 }
221
222 fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
223 unsafe {
224 vec_znx_big::vec_znx_big_sub_small_a(
225 self.ptr,
226 b.ptr as *mut vec_znx_big_t,
227 b.poly_count() as u64,
228 a.as_ptr(),
229 a.poly_count() as u64,
230 a.n() as u64,
231 b.ptr as *mut vec_znx_big_t,
232 b.poly_count() as u64,
233 )
234 }
235 }
236
237 fn vec_znx_big_sub_small_a(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
238 unsafe {
239 vec_znx_big::vec_znx_big_sub_small_a(
240 self.ptr,
241 c.ptr as *mut vec_znx_big_t,
242 c.poly_count() as u64,
243 a.as_ptr(),
244 a.poly_count() as u64,
245 a.n() as u64,
246 b.ptr as *mut vec_znx_big_t,
247 b.poly_count() as u64,
248 )
249 }
250 }
251
252 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
253 unsafe {
254 vec_znx_big::vec_znx_big_add_small(
255 self.ptr,
256 c.ptr as *mut vec_znx_big_t,
257 c.poly_count() as u64,
258 b.ptr as *mut vec_znx_big_t,
259 b.poly_count() as u64,
260 a.as_ptr(),
261 a.poly_count() as u64,
262 a.n() as u64,
263 )
264 }
265 }
266
267 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
268 unsafe {
269 vec_znx_big::vec_znx_big_add_small(
270 self.ptr,
271 b.ptr as *mut vec_znx_big_t,
272 b.poly_count() as u64,
273 b.ptr as *mut vec_znx_big_t,
274 b.poly_count() as u64,
275 a.as_ptr(),
276 a.poly_count() as u64,
277 a.n() as u64,
278 )
279 }
280 }
281
282 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
283 unsafe { vec_znx_big::vec_znx_big_normalize_base2k_tmp_bytes(self.ptr) as usize }
284 }
285
286 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<FFT64>, tmp_bytes: &mut [u8]) {
287 debug_assert!(
288 tmp_bytes.len() >= Self::vec_znx_big_normalize_tmp_bytes(self),
289 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",
290 tmp_bytes.len(),
291 Self::vec_znx_big_normalize_tmp_bytes(self)
292 );
293 #[cfg(debug_assertions)]
294 {
295 assert_alignement(tmp_bytes.as_ptr())
296 }
297 unsafe {
298 vec_znx_big::vec_znx_big_normalize_base2k(
299 self.ptr,
300 log_base2k as u64,
301 b.as_mut_ptr(),
302 b.size() as u64,
303 b.n() as u64,
304 a.ptr as *mut vec_znx_big_t,
305 a.size() as u64,
306 tmp_bytes.as_mut_ptr(),
307 )
308 }
309 }
310
311 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize {
312 unsafe { vec_znx_big::vec_znx_big_range_normalize_base2k_tmp_bytes(self.ptr) as usize }
313 }
314
315 fn vec_znx_big_range_normalize_base2k(
316 &self,
317 log_base2k: usize,
318 res: &mut VecZnx,
319 a: &VecZnxBig<FFT64>,
320 a_range_begin: usize,
321 a_range_xend: usize,
322 a_range_step: usize,
323 tmp_bytes: &mut [u8],
324 ) {
325 debug_assert!(
326 tmp_bytes.len() >= Self::vec_znx_big_range_normalize_base2k_tmp_bytes(self),
327 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_range_normalize_base2k_tmp_bytes()={}",
328 tmp_bytes.len(),
329 Self::vec_znx_big_range_normalize_base2k_tmp_bytes(self)
330 );
331 #[cfg(debug_assertions)]
332 {
333 assert_alignement(tmp_bytes.as_ptr())
334 }
335 unsafe {
336 vec_znx_big::vec_znx_big_range_normalize_base2k(
337 self.ptr,
338 log_base2k as u64,
339 res.as_mut_ptr(),
340 res.size() as u64,
341 res.n() as u64,
342 a.ptr as *mut vec_znx_big_t,
343 a_range_begin as u64,
344 a_range_xend as u64,
345 a_range_step as u64,
346 tmp_bytes.as_mut_ptr(),
347 );
348 }
349 }
350
351 fn vec_znx_big_automorphism(&self, gal_el: i64, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
352 unsafe {
353 vec_znx_big::vec_znx_big_automorphism(
354 self.ptr,
355 gal_el,
356 b.ptr as *mut vec_znx_big_t,
357 b.poly_count() as u64,
358 a.ptr as *mut vec_znx_big_t,
359 a.poly_count() as u64,
360 );
361 }
362 }
363
364 fn vec_znx_big_automorphism_inplace(&self, gal_el: i64, a: &mut VecZnxBig<FFT64>) {
365 unsafe {
366 vec_znx_big::vec_znx_big_automorphism(
367 self.ptr,
368 gal_el,
369 a.ptr as *mut vec_znx_big_t,
370 a.poly_count() as u64,
371 a.ptr as *mut vec_znx_big_t,
372 a.poly_count() as u64,
373 );
374 }
375 }
376}