poulpy/base2k/src/vec_znx_big_ops.rs
base2k/src/vec_znx_big_ops.rsRaw
1use std::cmp::min;
2
3use crate::ffi::vec_znx;
4use crate::internals::{apply_binary_op, apply_unary_op, ffi_binary_op_factory_type_1, ffi_ternary_op_factory};
5use crate::{Backend, FFT64, Module, VecZnx, VecZnxBig, VecZnxOps, ZnxBase, ZnxBasics, ZnxInfos, ZnxLayout, assert_alignement};
6
7pub trait VecZnxBigOps<B: Backend> {
8 /// Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
9 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBig<B>;
10
11 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
12 ///
13 /// Behavior: takes ownership of the backing array.
14 ///
15 /// # Arguments
16 ///
17 /// * `cols`: the number of polynomials..
18 /// * `size`: the number of polynomials per column.
19 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
20 ///
21 /// # Panics
22 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
23 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxBig<B>;
24
25 /// Returns a new [VecZnxBig] with the provided bytes array as backing array.
26 ///
27 /// Behavior: the backing array is only borrowed.
28 ///
29 /// # Arguments
30 ///
31 /// * `cols`: the number of polynomials..
32 /// * `size`: the number of polynomials per column.
33 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_big].
34 ///
35 /// # Panics
36 /// If `bytes.len()` < [Module::bytes_of_vec_znx_big].
37 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<B>;
38
39 /// Returns the minimum number of bytes necessary to allocate
40 /// a new [VecZnxBig] through [VecZnxBig::from_bytes].
41 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize;
42
43 /// Adds `a` to `b` and stores the result on `c`.
44 fn vec_znx_big_add(&self, c: &mut VecZnxBig<B>, a: &VecZnxBig<B>, b: &VecZnxBig<B>);
45
46 /// Adds `a` to `b` and stores the result on `b`.
47 fn vec_znx_big_add_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
48
49 /// Adds `a` to `b` and stores the result on `c`.
50 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
51
52 /// Adds `a` to `b` and stores the result on `b`.
53 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
54
55 /// Subtracts `a` to `b` and stores the result on `c`.
56 fn vec_znx_big_sub(&self, c: &mut VecZnxBig<B>, a: &VecZnxBig<B>, b: &VecZnxBig<B>);
57
58 /// Subtracts `a` to `b` and stores the result on `b`.
59 fn vec_znx_big_sub_ab_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
60
61 /// Subtracts `b` to `a` and stores the result on `b`.
62 fn vec_znx_big_sub_ba_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
63
64 /// Subtracts `b` to `a` and stores the result on `c`.
65 fn vec_znx_big_sub_small_ab(&self, c: &mut VecZnxBig<B>, a: &VecZnx, b: &VecZnxBig<B>);
66
67 /// Subtracts `a` to `b` and stores the result on `b`.
68 fn vec_znx_big_sub_small_ab_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
69
70 /// Subtracts `b` to `a` and stores the result on `c`.
71 fn vec_znx_big_sub_small_ba(&self, c: &mut VecZnxBig<B>, a: &VecZnxBig<B>, b: &VecZnx);
72
73 /// Subtracts `b` to `a` and stores the result on `b`.
74 fn vec_znx_big_sub_small_ba_inplace(&self, b: &mut VecZnxBig<B>, a: &VecZnx);
75
76 /// Returns the minimum number of bytes to apply [VecZnxBigOps::vec_znx_big_normalize].
77 fn vec_znx_big_normalize_tmp_bytes(&self, cols: usize) -> usize;
78
79 /// Normalizes `a` and stores the result on `b`.
80 ///
81 /// # Arguments
82 ///
83 /// * `log_base2k`: normalization basis.
84 /// * `tmp_bytes`: scratch space of size at least [VecZnxBigOps::vec_znx_big_normalize].
85 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<B>, tmp_bytes: &mut [u8]);
86
87 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
88 fn vec_znx_big_automorphism(&self, k: i64, b: &mut VecZnxBig<B>, a: &VecZnxBig<B>);
89
90 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
91 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<B>);
92}
93
94impl VecZnxBigOps<FFT64> for Module<FFT64> {
95 fn new_vec_znx_big(&self, cols: usize, size: usize) -> VecZnxBig<FFT64> {
96 VecZnxBig::new(self, cols, size)
97 }
98
99 fn new_vec_znx_big_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnxBig<FFT64> {
100 VecZnxBig::from_bytes(self, cols, size, bytes)
101 }
102
103 fn new_vec_znx_big_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnxBig<FFT64> {
104 VecZnxBig::from_bytes_borrow(self, cols, size, tmp_bytes)
105 }
106
107 fn bytes_of_vec_znx_big(&self, cols: usize, size: usize) -> usize {
108 VecZnxBig::bytes_of(self, cols, size)
109 }
110
111 fn vec_znx_big_add(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>, b: &VecZnxBig<FFT64>) {
112 let op = ffi_ternary_op_factory(
113 self.ptr,
114 c.size(),
115 c.sl(),
116 a.size(),
117 a.sl(),
118 b.size(),
119 b.sl(),
120 vec_znx::vec_znx_add,
121 );
122 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnxBig<FFT64>, VecZnxBig<FFT64>, false>(self, c, a, b, op);
123 }
124
125 fn vec_znx_big_add_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
126 unsafe {
127 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
128 Self::vec_znx_big_add(self, &mut *b_ptr, a, &*b_ptr);
129 }
130 }
131
132 fn vec_znx_big_sub(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>, b: &VecZnxBig<FFT64>) {
133 let op = ffi_ternary_op_factory(
134 self.ptr,
135 c.size(),
136 c.sl(),
137 a.size(),
138 a.sl(),
139 b.size(),
140 b.sl(),
141 vec_znx::vec_znx_sub,
142 );
143 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnxBig<FFT64>, VecZnxBig<FFT64>, true>(self, c, a, b, op);
144 }
145
146 fn vec_znx_big_sub_ab_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
147 unsafe {
148 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
149 Self::vec_znx_big_sub(self, &mut *b_ptr, a, &*b_ptr);
150 }
151 }
152
153 fn vec_znx_big_sub_ba_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
154 unsafe {
155 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
156 Self::vec_znx_big_sub(self, &mut *b_ptr, &*b_ptr, a);
157 }
158 }
159
160 fn vec_znx_big_sub_small_ba(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>, b: &VecZnx) {
161 let op = ffi_ternary_op_factory(
162 self.ptr,
163 c.size(),
164 c.sl(),
165 a.size(),
166 a.sl(),
167 b.size(),
168 b.sl(),
169 vec_znx::vec_znx_sub,
170 );
171 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnxBig<FFT64>, VecZnx, true>(self, c, a, b, op);
172 }
173
174 fn vec_znx_big_sub_small_ba_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
175 unsafe {
176 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
177 Self::vec_znx_big_sub_small_ba(self, &mut *b_ptr, &*b_ptr, a);
178 }
179 }
180
181 fn vec_znx_big_sub_small_ab(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
182 let op = ffi_ternary_op_factory(
183 self.ptr,
184 c.size(),
185 c.sl(),
186 a.size(),
187 a.sl(),
188 b.size(),
189 b.sl(),
190 vec_znx::vec_znx_sub,
191 );
192 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnx, VecZnxBig<FFT64>, true>(self, c, a, b, op);
193 }
194
195 fn vec_znx_big_sub_small_ab_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
196 unsafe {
197 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
198 Self::vec_znx_big_sub_small_ab(self, &mut *b_ptr, a, &*b_ptr);
199 }
200 }
201
202 fn vec_znx_big_add_small(&self, c: &mut VecZnxBig<FFT64>, a: &VecZnx, b: &VecZnxBig<FFT64>) {
203 let op = ffi_ternary_op_factory(
204 self.ptr,
205 c.size(),
206 c.sl(),
207 a.size(),
208 a.sl(),
209 b.size(),
210 b.sl(),
211 vec_znx::vec_znx_add,
212 );
213 apply_binary_op::<FFT64, VecZnxBig<FFT64>, VecZnx, VecZnxBig<FFT64>, false>(self, c, a, b, op);
214 }
215
216 fn vec_znx_big_add_small_inplace(&self, b: &mut VecZnxBig<FFT64>, a: &VecZnx) {
217 unsafe {
218 let b_ptr: *mut VecZnxBig<FFT64> = b as *mut VecZnxBig<FFT64>;
219 Self::vec_znx_big_add_small(self, &mut *b_ptr, a, &*b_ptr);
220 }
221 }
222
223 fn vec_znx_big_normalize_tmp_bytes(&self, cols: usize) -> usize {
224 Self::vec_znx_normalize_tmp_bytes(self, cols)
225 }
226
227 fn vec_znx_big_normalize(&self, log_base2k: usize, b: &mut VecZnx, a: &VecZnxBig<FFT64>, tmp_bytes: &mut [u8]) {
228 #[cfg(debug_assertions)]
229 {
230 assert!(tmp_bytes.len() >= Self::vec_znx_big_normalize_tmp_bytes(&self, a.cols()));
231 assert_alignement(tmp_bytes.as_ptr());
232 }
233
234 let a_size: usize = a.size();
235 let b_size: usize = b.size();
236 let a_sl: usize = a.sl();
237 let b_sl: usize = b.sl();
238 let a_cols: usize = a.cols();
239 let b_cols: usize = b.cols();
240 let min_cols: usize = min(a_cols, b_cols);
241 (0..min_cols).for_each(|i| unsafe {
242 vec_znx::vec_znx_normalize_base2k(
243 self.ptr,
244 log_base2k as u64,
245 b.at_mut_ptr(i, 0),
246 b_size as u64,
247 b_sl as u64,
248 a.at_ptr(i, 0),
249 a_size as u64,
250 a_sl as u64,
251 tmp_bytes.as_mut_ptr(),
252 );
253 });
254
255 (min_cols..b_cols).for_each(|i| (0..b_size).for_each(|j| b.zero_at(i, j)));
256 }
257
258 fn vec_znx_big_automorphism(&self, k: i64, b: &mut VecZnxBig<FFT64>, a: &VecZnxBig<FFT64>) {
259 let op = ffi_binary_op_factory_type_1(
260 self.ptr,
261 k,
262 b.size(),
263 b.sl(),
264 a.size(),
265 a.sl(),
266 vec_znx::vec_znx_automorphism,
267 );
268 apply_unary_op::<FFT64, VecZnxBig<FFT64>>(self, b, a, op);
269 }
270
271 fn vec_znx_big_automorphism_inplace(&self, k: i64, a: &mut VecZnxBig<FFT64>) {
272 unsafe {
273 let a_ptr: *mut VecZnxBig<FFT64> = a as *mut VecZnxBig<FFT64>;
274 Self::vec_znx_big_automorphism(self, k, &mut *a_ptr, &*a_ptr);
275 }
276 }
277}