poulpy/base2k/src/vec_znx_big_ops.rs
base2k/src/vec_znx_big_ops.rsRaw
1use crate::ffi::vec_znx_big::vec_znx_big_t;
2use crate::ffi::{vec_znx, vec_znx_big};
3use crate::internals::{apply_binary_op, ffi_ternary_op_factory};
4use crate::{Backend, FFT64, Module, VecZnx, VecZnxBig, ZnxBase, ZnxInfos, ZnxLayout, assert_alignement};
5
6pub trait VecZnxBigOps<B: Backend> {
7 /// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
8 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBig<B>;
9
10 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
11 ///
12 /// Behavior: takes ownership of the backing array.
13 ///
14 /// # Arguments
15 ///
16 /// * `cols`: the number of polynomials..
17 /// * `size`: the number of polynomials per column.
18 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
19 ///
20 /// # Panics
21 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
22 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxBig<B>;
23
24 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
25 ///
26 /// Behavior: the backing array is only borrowed.
27 ///
28 /// # Arguments
29 ///
30 /// * `cols`: the number of polynomials..
31 /// * `size`: the number of polynomials per column.
32 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
33 ///
34 /// # Panics
35 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
36 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<B>;
37
38 /// Returns the minimum number of bytes necessary to allocate
39 /// a new [VecZnxBig] through [VecZnxBig::from_bytes].
40 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize;
41
42 /// Adds `a` to `b` and stores the result on `c`.
43 fn vec_znx_big_add(&self, c: &mut VecZnxBig<B>, a: &VecZnxBig<B>, b: &VecZnxBig<B>);
44
45 /// Adds `a` to `b` and stores the result on `b`.
46 fn vec_znx_big_add_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
47
48 /// Adds `a` to `b` and stores the result on `c`.
49 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
50
51 /// Adds `a` to `b` and stores the result on `b`.
52 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
53
54 /// Subtracts `a` to `b` and stores the result on `c`.
55 fn vec_znx_big_sub(&self, c: &mut VecZnxBig<B>, a: &VecZnxBig<B>, b: &VecZnxBig<B>);
56
57 /// Subtracts `a` to `b` and stores the result on `b`.
58 fn vec_znx_big_sub_ab_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
59
60 /// Subtracts `b` to `a` and stores the result on `b`.
61 fn vec_znx_big_sub_ba_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
62
63 /// Subtracts `b` to `a` and stores the result on `c`.
64 fn vec_znx_big_sub_small_ab(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
65
66 /// Subtracts `a` to `b` and stores the result on `b`.
67 fn vec_znx_big_sub_small_ab_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
68
69 /// Subtracts `b` to `a` and stores the result on `c`.
70 fn vec_znx_big_sub_small_ba(&self, c: &mut VecZnxBig<B>, a: &VecZnxBig<B>, b: &VecZnx);
71
72 /// Subtracts `b` to `a` and stores the result on `b`.
73 fn vec_znx_big_sub_small_ba_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
74
75 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
76 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize;
77
78 /// Normalizes `a` and stores the result on `b`.
79 ///
80 /// # Arguments
81 ///
82 /// * `log_base2k`: normalization basis.
83 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
84 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<B>, tmp_bytes: &mut [u8]);
85
86 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_range_normalize_base2k].
87 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize;
88
89 /// Normalize `a`, taking into account column interleaving and stores the result on `b`.
90 ///
91 /// # Arguments
92 ///
93 /// * `log_base2k`: normalization basis.
94 /// * `a_range_begin`: column to start.
95 /// * `a_range_end`: column to end.
96 /// * `a_range_step`: column step size.
97 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_range_normalize_base2k_tmp_bytes].
98 fn vec_znx_big_range_normalize_base2k(
99 &self,
100 log_base2k: usize,
101 b: &mut VecZnx,
102 a: &VecZnxBig<B>,
103 a_range_begin: usize,
104 a_range_xend: usize,
105 a_range_step: usize,
106 tmp_bytes: &mut [u8],
107 );
108
109 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
110 fn vec_znx_big_automorphism(&self, k: i64, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
111
112 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
113 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<B>);
114}
115
116impl VecZnxBigOps<FFT64> for Module<FFT64> {
117 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBig<FFT64> {
118 VecZnxBig::new(self, cols, size)
119 }
120
121 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxBig<FFT64> {
122 VecZnxBig::from_bytes(self, cols, size, bytes)
123 }
124
125 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
126 VecZnxBig::from_bytes_borrow(self, cols, size, tmp_bytes)
127 }
128
129 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
130 VecZnxBig::bytes_of(self, cols, size)
131 }
132
133 fn vec_znx_big_add(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>, b: &VecZnxBig<FFT64>) {
134 let op = ffi_ternary_op_factory(
135 self.ptr,
136 c.size(),
137 c.sl(),
138 a.size(),
139 a.sl(),
140 b.size(),
141 b.sl(),
142 vec_znx::vec_znx_add,
143 );
144 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnxBig<FFT64>, VecZnxBig<FFT64>, false>(self, c, a, b, op);
145 }
146
147 fn vec_znx_big_add_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
148 unsafe {
149 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
150 Self::vec_znx_big_add(self, &mut *b_ptr, a, &*b_ptr);
151 }
152 }
153
154 fn vec_znx_big_sub(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>, b: &VecZnxBig<FFT64>) {
155 let op = ffi_ternary_op_factory(
156 self.ptr,
157 c.size(),
158 c.sl(),
159 a.size(),
160 a.sl(),
161 b.size(),
162 b.sl(),
163 vec_znx::vec_znx_sub,
164 );
165 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnxBig<FFT64>, VecZnxBig<FFT64>, true>(self, c, a, b, op);
166 }
167
168 fn vec_znx_big_sub_ab_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
169 unsafe {
170 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
171 Self::vec_znx_big_sub(self, &mut *b_ptr, a, &*b_ptr);
172 }
173 }
174
175 fn vec_znx_big_sub_ba_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
176 unsafe {
177 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
178 Self::vec_znx_big_sub(self, &mut *b_ptr, &*b_ptr, a);
179 }
180 }
181
182 fn vec_znx_big_sub_small_ba(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>, b: &VecZnx) {
183 let op = ffi_ternary_op_factory(
184 self.ptr,
185 c.size(),
186 c.sl(),
187 a.size(),
188 a.sl(),
189 b.size(),
190 b.sl(),
191 vec_znx::vec_znx_sub,
192 );
193 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnxBig<FFT64>, VecZnx, true>(self, c, a, b, op);
194 }
195
196 fn vec_znx_big_sub_small_ba_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
197 unsafe {
198 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
199 Self::vec_znx_big_sub_small_ba(self, &mut *b_ptr, &*b_ptr, a);
200 }
201 }
202
203 fn vec_znx_big_sub_small_ab(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
204 let op = ffi_ternary_op_factory(
205 self.ptr,
206 c.size(),
207 c.sl(),
208 a.size(),
209 a.sl(),
210 b.size(),
211 b.sl(),
212 vec_znx::vec_znx_sub,
213 );
214 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnx, VecZnxBig<FFT64>, true>(self, c, a, b, op);
215 }
216
217 fn vec_znx_big_sub_small_ab_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
218 unsafe {
219 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
220 Self::vec_znx_big_sub_small_ab(self, &mut *b_ptr, a, &*b_ptr);
221 }
222 }
223
224 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
225 let op = ffi_ternary_op_factory(
226 self.ptr,
227 c.size(),
228 c.sl(),
229 a.size(),
230 a.sl(),
231 b.size(),
232 b.sl(),
233 vec_znx::vec_znx_add,
234 );
235 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnx, VecZnxBig<FFT64>, false>(self, c, a, b, op);
236 }
237
238 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
239 unsafe {
240 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
241 Self::vec_znx_big_add_small(self, &mut *b_ptr, a, &*b_ptr);
242 }
243 }
244
245 fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
246 unsafe { vec_znx_big::vec_znx_big_normalize_base2k_tmp_bytes(self.ptr) as usize }
247 }
248
249 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<FFT64>, tmp_bytes: &mut [u8]) {
250 debug_assert!(
251 tmp_bytes.len() >= Self::vec_znx_big_normalize_tmp_bytes(self),
252 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",
253 tmp_bytes.len(),
254 Self::vec_znx_big_normalize_tmp_bytes(self)
255 );
256 #[cfg(debug_assertions)]
257 {
258 assert_alignement(tmp_bytes.as_ptr())
259 }
260 unsafe {
261 vec_znx_big::vec_znx_big_normalize_base2k(
262 self.ptr,
263 log_base2k as u64,
264 b.as_mut_ptr(),
265 b.size() as u64,
266 b.n() as u64,
267 a.ptr as *mut vec_znx_big_t,
268 a.size() as u64,
269 tmp_bytes.as_mut_ptr(),
270 )
271 }
272 }
273
274 fn vec_znx_big_range_normalize_base2k_tmp_bytes(&self) -> usize {
275 unsafe { vec_znx_big::vec_znx_big_range_normalize_base2k_tmp_bytes(self.ptr) as usize }
276 }
277
278 fn vec_znx_big_range_normalize_base2k(
279 &self,
280 log_base2k: usize,
281 res: &mut VecZnx,
282 a: &VecZnxBig<FFT64>,
283 a_range_begin: usize,
284 a_range_xend: usize,
285 a_range_step: usize,
286 tmp_bytes: &mut [u8],
287 ) {
288 debug_assert!(
289 tmp_bytes.len() >= Self::vec_znx_big_range_normalize_base2k_tmp_bytes(self),
290 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_range_normalize_base2k_tmp_bytes()={}",
291 tmp_bytes.len(),
292 Self::vec_znx_big_range_normalize_base2k_tmp_bytes(self)
293 );
294 #[cfg(debug_assertions)]
295 {
296 assert_alignement(tmp_bytes.as_ptr())
297 }
298 unsafe {
299 vec_znx_big::vec_znx_big_range_normalize_base2k(
300 self.ptr,
301 log_base2k as u64,
302 res.as_mut_ptr(),
303 res.size() as u64,
304 res.n() as u64,
305 a.ptr as *mut vec_znx_big_t,
306 a_range_begin as u64,
307 a_range_xend as u64,
308 a_range_step as u64,
309 tmp_bytes.as_mut_ptr(),
310 );
311 }
312 }
313
314 fn vec_znx_big_automorphism(&self, gal_el: i64, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
315 unsafe {
316 vec_znx_big::vec_znx_big_automorphism(
317 self.ptr,
318 gal_el,
319 b.ptr as *mut vec_znx_big_t,
320 b.poly_count() as u64,
321 a.ptr as *mut vec_znx_big_t,
322 a.poly_count() as u64,
323 );
324 }
325 }
326
327 fn vec_znx_big_automorphism_inplace(&self, gal_el: i64, a: &mut VecZnxBig<FFT64>) {
328 unsafe {
329 vec_znx_big::vec_znx_big_automorphism(
330 self.ptr,
331 gal_el,
332 a.ptr as *mut vec_znx_big_t,
333 a.poly_count() as u64,
334 a.ptr as *mut vec_znx_big_t,
335 a.poly_count() as u64,
336 );
337 }
338 }
339}