poulpy/backend/src/implementation/cpu_spqlios/vec_znx_dft_fft64.rs
backend/src/implementation/cpu_spqlios/vec_znx_dft_fft64.rsRaw
1use std::fmt;
2
3use crate::{
4 hal::{
5 api::{TakeSlice, VecZnxDftToVecZnxBigTmpBytes, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
6 layouts::{
7 Data, DataRef, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftBytesOf, VecZnxDftOwned,
8 VecZnxDftToMut, VecZnxDftToRef, VecZnxToRef,
9 },
10 oep::{
11 VecZnxDftAddImpl, VecZnxDftAddInplaceImpl, VecZnxDftAllocBytesImpl, VecZnxDftAllocImpl, VecZnxDftCopyImpl,
12 VecZnxDftFromBytesImpl, VecZnxDftFromVecZnxImpl, VecZnxDftSubABInplaceImpl, VecZnxDftSubBAInplaceImpl,
13 VecZnxDftSubImpl, VecZnxDftToVecZnxBigConsumeImpl, VecZnxDftToVecZnxBigImpl, VecZnxDftToVecZnxBigTmpAImpl,
14 VecZnxDftToVecZnxBigTmpBytesImpl, VecZnxDftZeroImpl,
15 },
16 },
17 implementation::cpu_spqlios::{
18 ffi::{vec_znx_big, vec_znx_dft},
19 module_fft64::FFT64,
20 },
21};
22
23const VEC_ZNX_DFT_FFT64_WORDSIZE: usize = 1;
24
25impl<D: Data> ZnxSliceSize for VecZnxDft<D, FFT64> {
26 fn sl(&self) -> usize {
27 VEC_ZNX_DFT_FFT64_WORDSIZE * self.n() * self.cols()
28 }
29}
30
31impl<D: Data> VecZnxDftBytesOf for VecZnxDft<D, FFT64> {
32 fn bytes_of(n: usize, cols: usize, size: usize) -> usize {
33 VEC_ZNX_DFT_FFT64_WORDSIZE * n * cols * size * size_of::<f64>()
34 }
35}
36
37impl<D: DataRef> ZnxView for VecZnxDft<D, FFT64> {
38 type Scalar = f64;
39}
40
41unsafe impl VecZnxDftFromBytesImpl<FFT64> for FFT64 {
42 fn vec_znx_dft_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<FFT64> {
43 VecZnxDft::<Vec<u8>, FFT64>::from_bytes(n, cols, size, bytes)
44 }
45}
46
47unsafe impl VecZnxDftAllocBytesImpl<FFT64> for FFT64 {
48 fn vec_znx_dft_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
49 VecZnxDft::<Vec<u8>, FFT64>::bytes_of(n, cols, size)
50 }
51}
52
53unsafe impl VecZnxDftAllocImpl<FFT64> for FFT64 {
54 fn vec_znx_dft_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxDftOwned<FFT64> {
55 VecZnxDftOwned::alloc(n, cols, size)
56 }
57}
58
59unsafe impl VecZnxDftToVecZnxBigTmpBytesImpl<FFT64> for FFT64 {
60 fn vec_znx_dft_to_vec_znx_big_tmp_bytes_impl(module: &Module<FFT64>, n: usize) -> usize {
61 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(module.ptr(), n as u64) as usize }
62 }
63}
64
65unsafe impl VecZnxDftToVecZnxBigImpl<FFT64> for FFT64 {
66 fn vec_znx_dft_to_vec_znx_big_impl<R, A>(
67 module: &Module<FFT64>,
68 res: &mut R,
69 res_col: usize,
70 a: &A,
71 a_col: usize,
72 scratch: &mut Scratch<FFT64>,
73 ) where
74 R: VecZnxBigToMut<FFT64>,
75 A: VecZnxDftToRef<FFT64>,
76 {
77 let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
78 let a: VecZnxDft<&[u8], FFT64> = a.to_ref();
79
80 #[cfg(debug_assertions)]
81 {
82 assert_eq!(res.n(), a.n())
83 }
84
85 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_dft_to_vec_znx_big_tmp_bytes(a.n()));
86
87 let min_size: usize = res.size().min(a.size());
88
89 unsafe {
90 (0..min_size).for_each(|j| {
91 vec_znx_dft::vec_znx_idft(
92 module.ptr(),
93 res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
94 1 as u64,
95 a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
96 1 as u64,
97 tmp_bytes.as_mut_ptr(),
98 )
99 });
100 (min_size..res.size()).for_each(|j| {
101 res.zero_at(res_col, j);
102 });
103 }
104 }
105}
106
107unsafe impl VecZnxDftToVecZnxBigTmpAImpl<FFT64> for FFT64 {
108 fn vec_znx_dft_to_vec_znx_big_tmp_a_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
109 where
110 R: VecZnxBigToMut<FFT64>,
111 A: VecZnxDftToMut<FFT64>,
112 {
113 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
114 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
115
116 let min_size: usize = res_mut.size().min(a_mut.size());
117
118 unsafe {
119 (0..min_size).for_each(|j| {
120 vec_znx_dft::vec_znx_idft_tmp_a(
121 module.ptr(),
122 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
123 1 as u64,
124 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
125 1 as u64,
126 )
127 });
128 (min_size..res_mut.size()).for_each(|j| {
129 res_mut.zero_at(res_col, j);
130 })
131 }
132 }
133}
134
135unsafe impl VecZnxDftToVecZnxBigConsumeImpl<FFT64> for FFT64 {
136 fn vec_znx_dft_to_vec_znx_big_consume_impl<D: Data>(module: &Module<FFT64>, mut a: VecZnxDft<D, FFT64>) -> VecZnxBig<D, FFT64>
137 where
138 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>,
139 {
140 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
141
142 unsafe {
143 // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
144 (0..a_mut.size()).for_each(|j| {
145 (0..a_mut.cols()).for_each(|i| {
146 vec_znx_dft::vec_znx_idft_tmp_a(
147 module.ptr(),
148 a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
149 1 as u64,
150 a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
151 1 as u64,
152 )
153 });
154 });
155 }
156
157 a.into_big()
158 }
159}
160
161unsafe impl VecZnxDftFromVecZnxImpl<FFT64> for FFT64 {
162 fn vec_znx_dft_from_vec_znx_impl<R, A>(
163 module: &Module<FFT64>,
164 step: usize,
165 offset: usize,
166 res: &mut R,
167 res_col: usize,
168 a: &A,
169 a_col: usize,
170 ) where
171 R: VecZnxDftToMut<FFT64>,
172 A: VecZnxToRef,
173 {
174 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
175 let a_ref: VecZnx<&[u8]> = a.to_ref();
176 let steps: usize = a_ref.size().div_ceil(step);
177 let min_steps: usize = res_mut.size().min(steps);
178 unsafe {
179 (0..min_steps).for_each(|j| {
180 let limb: usize = offset + j * step;
181 if limb < a_ref.size() {
182 vec_znx_dft::vec_znx_dft(
183 module.ptr(),
184 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
185 1 as u64,
186 a_ref.at_ptr(a_col, limb),
187 1 as u64,
188 a_ref.sl() as u64,
189 )
190 }
191 });
192 (min_steps..res_mut.size()).for_each(|j| {
193 res_mut.zero_at(res_col, j);
194 });
195 }
196 }
197}
198
199unsafe impl VecZnxDftAddImpl<FFT64> for FFT64 {
200 fn vec_znx_dft_add_impl<R, A, D>(
201 module: &Module<FFT64>,
202 res: &mut R,
203 res_col: usize,
204 a: &A,
205 a_col: usize,
206 b: &D,
207 b_col: usize,
208 ) where
209 R: VecZnxDftToMut<FFT64>,
210 A: VecZnxDftToRef<FFT64>,
211 D: VecZnxDftToRef<FFT64>,
212 {
213 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
214 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
215 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
216
217 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
218
219 unsafe {
220 (0..min_size).for_each(|j| {
221 vec_znx_dft::vec_dft_add(
222 module.ptr(),
223 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
224 1,
225 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
226 1,
227 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
228 1,
229 );
230 });
231 }
232 (min_size..res_mut.size()).for_each(|j| {
233 res_mut.zero_at(res_col, j);
234 })
235 }
236}
237
238unsafe impl VecZnxDftAddInplaceImpl<FFT64> for FFT64 {
239 fn vec_znx_dft_add_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
240 where
241 R: VecZnxDftToMut<FFT64>,
242 A: VecZnxDftToRef<FFT64>,
243 {
244 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
245 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
246
247 let min_size: usize = res_mut.size().min(a_ref.size());
248
249 unsafe {
250 (0..min_size).for_each(|j| {
251 vec_znx_dft::vec_dft_add(
252 module.ptr(),
253 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
254 1,
255 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
256 1,
257 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
258 1,
259 );
260 });
261 }
262 }
263}
264
265unsafe impl VecZnxDftSubImpl<FFT64> for FFT64 {
266 fn vec_znx_dft_sub_impl<R, A, D>(
267 module: &Module<FFT64>,
268 res: &mut R,
269 res_col: usize,
270 a: &A,
271 a_col: usize,
272 b: &D,
273 b_col: usize,
274 ) where
275 R: VecZnxDftToMut<FFT64>,
276 A: VecZnxDftToRef<FFT64>,
277 D: VecZnxDftToRef<FFT64>,
278 {
279 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
280 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
281 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
282
283 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
284
285 unsafe {
286 (0..min_size).for_each(|j| {
287 vec_znx_dft::vec_dft_sub(
288 module.ptr(),
289 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
290 1,
291 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
292 1,
293 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
294 1,
295 );
296 });
297 }
298 (min_size..res_mut.size()).for_each(|j| {
299 res_mut.zero_at(res_col, j);
300 })
301 }
302}
303
304unsafe impl VecZnxDftSubABInplaceImpl<FFT64> for FFT64 {
305 fn vec_znx_dft_sub_ab_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
306 where
307 R: VecZnxDftToMut<FFT64>,
308 A: VecZnxDftToRef<FFT64>,
309 {
310 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
311 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
312
313 let min_size: usize = res_mut.size().min(a_ref.size());
314
315 unsafe {
316 (0..min_size).for_each(|j| {
317 vec_znx_dft::vec_dft_sub(
318 module.ptr(),
319 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
320 1,
321 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
322 1,
323 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
324 1,
325 );
326 });
327 }
328 }
329}
330
331unsafe impl VecZnxDftSubBAInplaceImpl<FFT64> for FFT64 {
332 fn vec_znx_dft_sub_ba_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
333 where
334 R: VecZnxDftToMut<FFT64>,
335 A: VecZnxDftToRef<FFT64>,
336 {
337 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
338 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
339
340 let min_size: usize = res_mut.size().min(a_ref.size());
341
342 unsafe {
343 (0..min_size).for_each(|j| {
344 vec_znx_dft::vec_dft_sub(
345 module.ptr(),
346 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
347 1,
348 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
349 1,
350 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
351 1,
352 );
353 });
354 }
355 }
356}
357
358unsafe impl VecZnxDftCopyImpl<FFT64> for FFT64 {
359 fn vec_znx_dft_copy_impl<R, A>(
360 _module: &Module<FFT64>,
361 step: usize,
362 offset: usize,
363 res: &mut R,
364 res_col: usize,
365 a: &A,
366 a_col: usize,
367 ) where
368 R: VecZnxDftToMut<FFT64>,
369 A: VecZnxDftToRef<FFT64>,
370 {
371 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
372 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
373
374 let steps: usize = a_ref.size().div_ceil(step);
375 let min_steps: usize = res_mut.size().min(steps);
376
377 (0..min_steps).for_each(|j| {
378 let limb: usize = offset + j * step;
379 if limb < a_ref.size() {
380 res_mut
381 .at_mut(res_col, j)
382 .copy_from_slice(a_ref.at(a_col, limb));
383 }
384 });
385 (min_steps..res_mut.size()).for_each(|j| {
386 res_mut.zero_at(res_col, j);
387 })
388 }
389}
390
391unsafe impl VecZnxDftZeroImpl<FFT64> for FFT64 {
392 fn vec_znx_dft_zero_impl<R>(_module: &Module<FFT64>, res: &mut R)
393 where
394 R: VecZnxDftToMut<FFT64>,
395 {
396 res.to_mut().data.fill(0);
397 }
398}
399
400impl<D: DataRef> fmt::Display for VecZnxDft<D, FFT64> {
401 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
402 writeln!(
403 f,
404 "VecZnxDft(n={}, cols={}, size={})",
405 self.n, self.cols, self.size
406 )?;
407
408 for col in 0..self.cols {
409 writeln!(f, "Column {}:", col)?;
410 for size in 0..self.size {
411 let coeffs = self.at(col, size);
412 write!(f, " Size {}: [", size)?;
413
414 let max_show = 100;
415 let show_count = coeffs.len().min(max_show);
416
417 for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
418 if i > 0 {
419 write!(f, ", ")?;
420 }
421 write!(f, "{}", coeff)?;
422 }
423
424 if coeffs.len() > max_show {
425 write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
426 }
427
428 writeln!(f, "]")?;
429 }
430 }
431 Ok(())
432 }
433}