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, 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 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(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig, tmp_bytes: &mut [u8]);
145
146 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize;
147
148 fn vec_znx_big_range_normalize_base2k(
149 &self,
150 log_base2k: usize,
151 res: &mut VecZnx,
152 a: &VecZnxBig,
153 a_range_begin: usize,
154 a_range_xend: usize,
155 a_range_step: usize,
156 tmp_bytes: &mut [u8],
157 );
158
159 fn vec_znx_big_automorphism(&self, gal_el: i64, b: &mut VecZnxBig, a: &VecZnxBig);
160
161 fn vec_znx_big_automorphism_inplace(&self, gal_el: i64, a: &mut VecZnxBig);
162}
163
164impl VecZnxBigOps for Module {
165 fn new_vec_znx_big(&self, cols: usize) -> VecZnxBig {
166 let mut data: Vec<u8> = alloc_aligned::<u8>(self.bytes_of_vec_znx_big(cols));
167 let ptr: *mut u8 = data.as_mut_ptr();
168 VecZnxBig {
169 data: data,
170 ptr: ptr,
171 n: self.n(),
172 cols: cols,
173 backend: self.backend(),
174 }
175 }
176
177 fn new_vec_znx_big_from_bytes(&self, cols: usize, bytes: &mut [u8]) -> VecZnxBig {
178 VecZnxBig::from_bytes(self, cols, bytes)
179 }
180
181 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, tmp_bytes: &mut [u8]) -> VecZnxBig {
182 VecZnxBig::from_bytes_borrow(self, cols, tmp_bytes)
183 }
184
185 fn bytes_of_vec_znx_big(&self, cols: usize) -> usize {
186 unsafe { vec_znx_big::bytes_of_vec_znx_big(self.ptr, cols as u64) as usize }
187 }
188
189 fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VecZnxBig, a: &VecZnx) {
190 unsafe {
191 vec_znx_big::vec_znx_big_sub_small_a(
192 self.ptr,
193 b.ptr as *mut vec_znx_big_t,
194 b.cols() as u64,
195 a.as_ptr(),
196 a.cols() as u64,
197 a.n() as u64,
198 b.ptr as *mut vec_znx_big_t,
199 b.cols() as u64,
200 )
201 }
202 }
203
204 fn vec_znx_big_sub_small_a(&self, c: &mut VecZnxBig, a: &VecZnx, b: &VecZnxBig) {
205 unsafe {
206 vec_znx_big::vec_znx_big_sub_small_a(
207 self.ptr,
208 c.ptr as *mut vec_znx_big_t,
209 c.cols() as u64,
210 a.as_ptr(),
211 a.cols() as u64,
212 a.n() as u64,
213 b.ptr as *mut vec_znx_big_t,
214 b.cols() as u64,
215 )
216 }
217 }
218
219 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig, a: &VecZnx, b: &VecZnxBig) {
220 unsafe {
221 vec_znx_big::vec_znx_big_add_small(
222 self.ptr,
223 c.ptr as *mut vec_znx_big_t,
224 c.cols() as u64,
225 b.ptr as *mut vec_znx_big_t,
226 b.cols() as u64,
227 a.as_ptr(),
228 a.cols() as u64,
229 a.n() as u64,
230 )
231 }
232 }
233
234 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig, a: &VecZnx) {
235 unsafe {
236 vec_znx_big::vec_znx_big_add_small(
237 self.ptr,
238 b.ptr as *mut vec_znx_big_t,
239 b.cols() as u64,
240 b.ptr as *mut vec_znx_big_t,
241 b.cols() as u64,
242 a.as_ptr(),
243 a.cols() as u64,
244 a.n() as u64,
245 )
246 }
247 }
248
249 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
250 unsafe { vec_znx_big::vec_znx_big_normalize_base2k_tmp_bytes(self.ptr) as usize }
251 }
252
253 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig, tmp_bytes: &mut [u8]) {
254 debug_assert!(
255 tmp_bytes.len() >= <Module as VecZnxBigOps>::vec_znx_big_normalize_tmp_bytes(self),
256 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",
257 tmp_bytes.len(),
258 <Module as VecZnxBigOps>::vec_znx_big_normalize_tmp_bytes(self)
259 );
260 #[cfg(debug_assertions)]
261 {
262 assert_alignement(tmp_bytes.as_ptr())
263 }
264 unsafe {
265 vec_znx_big::vec_znx_big_normalize_base2k(
266 self.ptr,
267 log_base2k as u64,
268 b.as_mut_ptr(),
269 b.cols() as u64,
270 b.n() as u64,
271 a.ptr as *mut vec_znx_big_t,
272 a.cols() as u64,
273 tmp_bytes.as_mut_ptr(),
274 )
275 }
276 }
277
278 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize {
279 unsafe { vec_znx_big::vec_znx_big_range_normalize_base2k_tmp_bytes(self.ptr) as usize }
280 }
281
282 fn vec_znx_big_range_normalize_base2k(
283 &self,
284 log_base2k: usize,
285 res: &mut VecZnx,
286 a: &VecZnxBig,
287 a_range_begin: usize,
288 a_range_xend: usize,
289 a_range_step: usize,
290 tmp_bytes: &mut [u8],
291 ) {
292 debug_assert!(
293 tmp_bytes.len() >= <Module as VecZnxBigOps>::vec_znx_big_range_normalize_base2k_tmp_bytes(self),
294 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_range_normalize_base2k_tmp_bytes()={}",
295 tmp_bytes.len(),
296 <Module as VecZnxBigOps>::vec_znx_big_range_normalize_base2k_tmp_bytes(self)
297 );
298 #[cfg(debug_assertions)]
299 {
300 assert_alignement(tmp_bytes.as_ptr())
301 }
302 unsafe {
303 vec_znx_big::vec_znx_big_range_normalize_base2k(
304 self.ptr,
305 log_base2k as u64,
306 res.as_mut_ptr(),
307 res.cols() as u64,
308 res.n() as u64,
309 a.ptr as *mut vec_znx_big_t,
310 a_range_begin as u64,
311 a_range_xend as u64,
312 a_range_step as u64,
313 tmp_bytes.as_mut_ptr(),
314 );
315 }
316 }
317
318 fn vec_znx_big_automorphism(&self, gal_el: i64, b: &mut VecZnxBig, a: &VecZnxBig) {
319 unsafe {
320 vec_znx_big::vec_znx_big_automorphism(
321 self.ptr,
322 gal_el,
323 b.ptr as *mut vec_znx_big_t,
324 b.cols() as u64,
325 a.ptr as *mut vec_znx_big_t,
326 a.cols() as u64,
327 );
328 }
329 }
330
331 fn vec_znx_big_automorphism_inplace(&self, gal_el: i64, a: &mut VecZnxBig) {
332 unsafe {
333 vec_znx_big::vec_znx_big_automorphism(
334 self.ptr,
335 gal_el,
336 a.ptr as *mut vec_znx_big_t,
337 a.cols() as u64,
338 a.ptr as *mut vec_znx_big_t,
339 a.cols() as u64,
340 );
341 }
342 }
343}