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