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