poulpy/backend/src/vec_znx_dft_ops.rs
backend/src/vec_znx_dft_ops.rsRaw
1use crate::ffi::{vec_znx_big, vec_znx_dft};
2use crate::vec_znx_dft::bytes_of_vec_znx_dft;
3use crate::znx_base::ZnxInfos;
4use crate::{
5 Backend, Scratch, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef,
6 ZnxSliceSize,
7};
8use crate::{FFT64, Module, ZnxView, ZnxViewMut, ZnxZero};
9use std::cmp::min;
10
11pub trait VecZnxDftAlloc<B: Backend> {
12 /// Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
13 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B>;
14
15 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
16 ///
17 /// Behavior: takes ownership of the backing array.
18 ///
19 /// # Arguments
20 ///
21 /// * `cols`: the number of cols of the [VecZnxDft].
22 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
23 ///
24 /// # Panics
25 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
26 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B>;
27
28 /// Returns a new [VecZnxDft] with the provided bytes array as backing array.
29 ///
30 /// # Arguments
31 ///
32 /// * `cols`: the number of cols of the [VecZnxDft].
33 /// * `bytes`: a byte array of size at least [Module::bytes_of_vec_znx_dft].
34 ///
35 /// # Panics
36 /// If `bytes.len()` < [Module::bytes_of_vec_znx_dft].
37 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
38}
39
40pub trait VecZnxDftOps<B: Backend> {
41 /// Returns the minimum number of bytes necessary to allocate
42 /// a new [VecZnxDft] through [VecZnxDft::from_bytes].
43 fn vec_znx_idft_tmp_bytes(&self) -> usize;
44
45 fn vec_znx_dft_add<R, A, D>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
46 where
47 R: VecZnxDftToMut<B>,
48 A: VecZnxDftToRef<B>,
49 D: VecZnxDftToRef<B>;
50
51 fn vec_znx_dft_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
52 where
53 R: VecZnxDftToMut<B>,
54 A: VecZnxDftToRef<B>;
55
56 fn vec_znx_dft_sub<R, A, D>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
57 where
58 R: VecZnxDftToMut<B>,
59 A: VecZnxDftToRef<B>,
60 D: VecZnxDftToRef<B>;
61
62 fn vec_znx_dft_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
63 where
64 R: VecZnxDftToMut<B>,
65 A: VecZnxDftToRef<B>;
66
67 fn vec_znx_dft_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
68 where
69 R: VecZnxDftToMut<B>,
70 A: VecZnxDftToRef<B>;
71
72 fn vec_znx_dft_copy<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
73 where
74 R: VecZnxDftToMut<B>,
75 A: VecZnxDftToRef<B>;
76
77 /// b <- IDFT(a), uses a as scratch space.
78 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
79 where
80 R: VecZnxBigToMut<B>,
81 A: VecZnxDftToMut<B>;
82
83 /// Consumes a to return IDFT(a) in big coeff space.
84 fn vec_znx_idft_consume<D>(&self, a: VecZnxDft<D, B>) -> VecZnxBig<D, FFT64>
85 where
86 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>;
87
88 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
89 where
90 R: VecZnxBigToMut<B>,
91 A: VecZnxDftToRef<B>;
92
93 fn vec_znx_dft<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
94 where
95 R: VecZnxDftToMut<FFT64>,
96 A: VecZnxToRef;
97}
98
99impl<B: Backend> VecZnxDftAlloc<B> for Module<B> {
100 fn new_vec_znx_dft(&self, cols: usize, size: usize) -> VecZnxDftOwned<B> {
101 VecZnxDftOwned::new(&self, cols, size)
102 }
103
104 fn new_vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B> {
105 VecZnxDftOwned::new_from_bytes(self, cols, size, bytes)
106 }
107
108 fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize {
109 bytes_of_vec_znx_dft(self, cols, size)
110 }
111}
112
113impl VecZnxDftOps<FFT64> for Module<FFT64> {
114 fn vec_znx_dft_add<R, A, D>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
115 where
116 R: VecZnxDftToMut<FFT64>,
117 A: VecZnxDftToRef<FFT64>,
118 D: VecZnxDftToRef<FFT64>,
119 {
120 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
121 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
122 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
123
124 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
125
126 unsafe {
127 (0..min_size).for_each(|j| {
128 vec_znx_dft::vec_dft_add(
129 self.ptr,
130 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
131 1,
132 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
133 1,
134 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
135 1,
136 );
137 });
138 }
139 (min_size..res_mut.size()).for_each(|j| {
140 res_mut.zero_at(res_col, j);
141 })
142 }
143
144 fn vec_znx_dft_add_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
145 where
146 R: VecZnxDftToMut<FFT64>,
147 A: VecZnxDftToRef<FFT64>,
148 {
149 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
150 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
151
152 let min_size: usize = res_mut.size().min(a_ref.size());
153
154 unsafe {
155 (0..min_size).for_each(|j| {
156 vec_znx_dft::vec_dft_add(
157 self.ptr,
158 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
159 1,
160 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
161 1,
162 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
163 1,
164 );
165 });
166 }
167 }
168
169 fn vec_znx_dft_sub<R, A, D>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
170 where
171 R: VecZnxDftToMut<FFT64>,
172 A: VecZnxDftToRef<FFT64>,
173 D: VecZnxDftToRef<FFT64>,
174 {
175 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
176 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
177 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
178
179 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
180
181 unsafe {
182 (0..min_size).for_each(|j| {
183 vec_znx_dft::vec_dft_sub(
184 self.ptr,
185 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
186 1,
187 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
188 1,
189 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
190 1,
191 );
192 });
193 }
194 (min_size..res_mut.size()).for_each(|j| {
195 res_mut.zero_at(res_col, j);
196 })
197 }
198
199 fn vec_znx_dft_sub_ab_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
200 where
201 R: VecZnxDftToMut<FFT64>,
202 A: VecZnxDftToRef<FFT64>,
203 {
204 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
205 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
206
207 let min_size: usize = res_mut.size().min(a_ref.size());
208
209 unsafe {
210 (0..min_size).for_each(|j| {
211 vec_znx_dft::vec_dft_sub(
212 self.ptr,
213 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
214 1,
215 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
216 1,
217 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
218 1,
219 );
220 });
221 }
222 }
223
224 fn vec_znx_dft_sub_ba_inplace<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize)
225 where
226 R: VecZnxDftToMut<FFT64>,
227 A: VecZnxDftToRef<FFT64>,
228 {
229 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
230 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
231
232 let min_size: usize = res_mut.size().min(a_ref.size());
233
234 unsafe {
235 (0..min_size).for_each(|j| {
236 vec_znx_dft::vec_dft_sub(
237 self.ptr,
238 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
239 1,
240 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
241 1,
242 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
243 1,
244 );
245 });
246 }
247 }
248
249 fn vec_znx_dft_copy<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
250 where
251 R: VecZnxDftToMut<FFT64>,
252 A: VecZnxDftToRef<FFT64>,
253 {
254 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
255 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
256
257 let steps: usize = (a_ref.size() + step - 1) / step;
258 let min_steps: usize = min(res_mut.size(), steps);
259
260 (0..min_steps).for_each(|j| {
261 let limb: usize = offset + j * step;
262 if limb < a_ref.size() {
263 res_mut
264 .at_mut(res_col, j)
265 .copy_from_slice(a_ref.at(a_col, limb));
266 }
267 });
268 (min_steps..res_mut.size()).for_each(|j| {
269 res_mut.zero_at(res_col, j);
270 })
271 }
272
273 fn vec_znx_idft_tmp_a<R, A>(&self, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
274 where
275 R: VecZnxBigToMut<FFT64>,
276 A: VecZnxDftToMut<FFT64>,
277 {
278 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
279 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
280
281 let min_size: usize = min(res_mut.size(), a_mut.size());
282
283 unsafe {
284 (0..min_size).for_each(|j| {
285 vec_znx_dft::vec_znx_idft_tmp_a(
286 self.ptr,
287 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
288 1 as u64,
289 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
290 1 as u64,
291 )
292 });
293 (min_size..res_mut.size()).for_each(|j| {
294 res_mut.zero_at(res_col, j);
295 })
296 }
297 }
298
299 fn vec_znx_idft_consume<D>(&self, mut a: VecZnxDft<D, FFT64>) -> VecZnxBig<D, FFT64>
300 where
301 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>,
302 {
303 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
304
305 unsafe {
306 // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
307 (0..a_mut.size()).for_each(|j| {
308 (0..a_mut.cols()).for_each(|i| {
309 vec_znx_dft::vec_znx_idft_tmp_a(
310 self.ptr,
311 a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
312 1 as u64,
313 a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
314 1 as u64,
315 )
316 });
317 });
318 }
319
320 a.into_big()
321 }
322
323 fn vec_znx_idft_tmp_bytes(&self) -> usize {
324 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(self.ptr) as usize }
325 }
326
327 fn vec_znx_dft<R, A>(&self, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
328 where
329 R: VecZnxDftToMut<FFT64>,
330 A: VecZnxToRef,
331 {
332 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
333 let a_ref: crate::VecZnx<&[u8]> = a.to_ref();
334 let steps: usize = (a_ref.size() + step - 1) / step;
335 let min_steps: usize = min(res_mut.size(), steps);
336 unsafe {
337 (0..min_steps).for_each(|j| {
338 let limb: usize = offset + j * step;
339 if limb < a_ref.size() {
340 vec_znx_dft::vec_znx_dft(
341 self.ptr,
342 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
343 1 as u64,
344 a_ref.at_ptr(a_col, limb),
345 1 as u64,
346 a_ref.sl() as u64,
347 )
348 }
349 });
350 (min_steps..res_mut.size()).for_each(|j| {
351 res_mut.zero_at(res_col, j);
352 });
353 }
354 }
355
356 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
357 fn vec_znx_idft<R, A>(&self, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch)
358 where
359 R: VecZnxBigToMut<FFT64>,
360 A: VecZnxDftToRef<FFT64>,
361 {
362 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
363 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
364
365 let (tmp_bytes, _) = scratch.tmp_slice(self.vec_znx_idft_tmp_bytes());
366
367 let min_size: usize = min(res_mut.size(), a_ref.size());
368
369 unsafe {
370 (0..min_size).for_each(|j| {
371 vec_znx_dft::vec_znx_idft(
372 self.ptr,
373 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
374 1 as u64,
375 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
376 1 as u64,
377 tmp_bytes.as_mut_ptr(),
378 )
379 });
380 (min_size..res_mut.size()).for_each(|j| {
381 res_mut.zero_at(res_col, j);
382 });
383 }
384 }
385}