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