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>) -> usize {
61 unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(module.ptr()) 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_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
78 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
79
80 let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_dft_to_vec_znx_big_tmp_bytes());
81
82 let min_size: usize = res_mut.size().min(a_ref.size());
83
84 unsafe {
85 (0..min_size).for_each(|j| {
86 vec_znx_dft::vec_znx_idft(
87 module.ptr(),
88 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
89 1 as u64,
90 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
91 1 as u64,
92 tmp_bytes.as_mut_ptr(),
93 )
94 });
95 (min_size..res_mut.size()).for_each(|j| {
96 res_mut.zero_at(res_col, j);
97 });
98 }
99 }
100}
101
102unsafe impl VecZnxDftToVecZnxBigTmpAImpl<FFT64> for FFT64 {
103 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)
104 where
105 R: VecZnxBigToMut<FFT64>,
106 A: VecZnxDftToMut<FFT64>,
107 {
108 let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
109 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
110
111 let min_size: usize = res_mut.size().min(a_mut.size());
112
113 unsafe {
114 (0..min_size).for_each(|j| {
115 vec_znx_dft::vec_znx_idft_tmp_a(
116 module.ptr(),
117 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
118 1 as u64,
119 a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
120 1 as u64,
121 )
122 });
123 (min_size..res_mut.size()).for_each(|j| {
124 res_mut.zero_at(res_col, j);
125 })
126 }
127 }
128}
129
130unsafe impl VecZnxDftToVecZnxBigConsumeImpl<FFT64> for FFT64 {
131 fn vec_znx_dft_to_vec_znx_big_consume_impl<D: Data>(module: &Module<FFT64>, mut a: VecZnxDft<D, FFT64>) -> VecZnxBig<D, FFT64>
132 where
133 VecZnxDft<D, FFT64>: VecZnxDftToMut<FFT64>,
134 {
135 let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
136
137 unsafe {
138 // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
139 (0..a_mut.size()).for_each(|j| {
140 (0..a_mut.cols()).for_each(|i| {
141 vec_znx_dft::vec_znx_idft_tmp_a(
142 module.ptr(),
143 a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
144 1 as u64,
145 a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
146 1 as u64,
147 )
148 });
149 });
150 }
151
152 a.into_big()
153 }
154}
155
156unsafe impl VecZnxDftFromVecZnxImpl<FFT64> for FFT64 {
157 fn vec_znx_dft_from_vec_znx_impl<R, A>(
158 module: &Module<FFT64>,
159 step: usize,
160 offset: usize,
161 res: &mut R,
162 res_col: usize,
163 a: &A,
164 a_col: usize,
165 ) where
166 R: VecZnxDftToMut<FFT64>,
167 A: VecZnxToRef,
168 {
169 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
170 let a_ref: VecZnx<&[u8]> = a.to_ref();
171 let steps: usize = a_ref.size().div_ceil(step);
172 let min_steps: usize = res_mut.size().min(steps);
173 unsafe {
174 (0..min_steps).for_each(|j| {
175 let limb: usize = offset + j * step;
176 if limb < a_ref.size() {
177 vec_znx_dft::vec_znx_dft(
178 module.ptr(),
179 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
180 1 as u64,
181 a_ref.at_ptr(a_col, limb),
182 1 as u64,
183 a_ref.sl() as u64,
184 )
185 }
186 });
187 (min_steps..res_mut.size()).for_each(|j| {
188 res_mut.zero_at(res_col, j);
189 });
190 }
191 }
192}
193
194unsafe impl VecZnxDftAddImpl<FFT64> for FFT64 {
195 fn vec_znx_dft_add_impl<R, A, D>(
196 module: &Module<FFT64>,
197 res: &mut R,
198 res_col: usize,
199 a: &A,
200 a_col: usize,
201 b: &D,
202 b_col: usize,
203 ) where
204 R: VecZnxDftToMut<FFT64>,
205 A: VecZnxDftToRef<FFT64>,
206 D: VecZnxDftToRef<FFT64>,
207 {
208 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
209 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
210 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
211
212 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
213
214 unsafe {
215 (0..min_size).for_each(|j| {
216 vec_znx_dft::vec_dft_add(
217 module.ptr(),
218 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
219 1,
220 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
221 1,
222 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
223 1,
224 );
225 });
226 }
227 (min_size..res_mut.size()).for_each(|j| {
228 res_mut.zero_at(res_col, j);
229 })
230 }
231}
232
233unsafe impl VecZnxDftAddInplaceImpl<FFT64> for FFT64 {
234 fn vec_znx_dft_add_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
235 where
236 R: VecZnxDftToMut<FFT64>,
237 A: VecZnxDftToRef<FFT64>,
238 {
239 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
240 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
241
242 let min_size: usize = res_mut.size().min(a_ref.size());
243
244 unsafe {
245 (0..min_size).for_each(|j| {
246 vec_znx_dft::vec_dft_add(
247 module.ptr(),
248 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
249 1,
250 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
251 1,
252 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
253 1,
254 );
255 });
256 }
257 }
258}
259
260unsafe impl VecZnxDftSubImpl<FFT64> for FFT64 {
261 fn vec_znx_dft_sub_impl<R, A, D>(
262 module: &Module<FFT64>,
263 res: &mut R,
264 res_col: usize,
265 a: &A,
266 a_col: usize,
267 b: &D,
268 b_col: usize,
269 ) where
270 R: VecZnxDftToMut<FFT64>,
271 A: VecZnxDftToRef<FFT64>,
272 D: VecZnxDftToRef<FFT64>,
273 {
274 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
275 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
276 let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
277
278 let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
279
280 unsafe {
281 (0..min_size).for_each(|j| {
282 vec_znx_dft::vec_dft_sub(
283 module.ptr(),
284 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
285 1,
286 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
287 1,
288 b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
289 1,
290 );
291 });
292 }
293 (min_size..res_mut.size()).for_each(|j| {
294 res_mut.zero_at(res_col, j);
295 })
296 }
297}
298
299unsafe impl VecZnxDftSubABInplaceImpl<FFT64> for FFT64 {
300 fn vec_znx_dft_sub_ab_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
301 where
302 R: VecZnxDftToMut<FFT64>,
303 A: VecZnxDftToRef<FFT64>,
304 {
305 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
306 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
307
308 let min_size: usize = res_mut.size().min(a_ref.size());
309
310 unsafe {
311 (0..min_size).for_each(|j| {
312 vec_znx_dft::vec_dft_sub(
313 module.ptr(),
314 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
315 1,
316 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
317 1,
318 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
319 1,
320 );
321 });
322 }
323 }
324}
325
326unsafe impl VecZnxDftSubBAInplaceImpl<FFT64> for FFT64 {
327 fn vec_znx_dft_sub_ba_inplace_impl<R, A>(module: &Module<FFT64>, res: &mut R, res_col: usize, a: &A, a_col: usize)
328 where
329 R: VecZnxDftToMut<FFT64>,
330 A: VecZnxDftToRef<FFT64>,
331 {
332 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
333 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
334
335 let min_size: usize = res_mut.size().min(a_ref.size());
336
337 unsafe {
338 (0..min_size).for_each(|j| {
339 vec_znx_dft::vec_dft_sub(
340 module.ptr(),
341 res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
342 1,
343 a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
344 1,
345 res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
346 1,
347 );
348 });
349 }
350 }
351}
352
353unsafe impl VecZnxDftCopyImpl<FFT64> for FFT64 {
354 fn vec_znx_dft_copy_impl<R, A>(
355 _module: &Module<FFT64>,
356 step: usize,
357 offset: usize,
358 res: &mut R,
359 res_col: usize,
360 a: &A,
361 a_col: usize,
362 ) where
363 R: VecZnxDftToMut<FFT64>,
364 A: VecZnxDftToRef<FFT64>,
365 {
366 let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
367 let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
368
369 let steps: usize = a_ref.size().div_ceil(step);
370 let min_steps: usize = res_mut.size().min(steps);
371
372 (0..min_steps).for_each(|j| {
373 let limb: usize = offset + j * step;
374 if limb < a_ref.size() {
375 res_mut
376 .at_mut(res_col, j)
377 .copy_from_slice(a_ref.at(a_col, limb));
378 }
379 });
380 (min_steps..res_mut.size()).for_each(|j| {
381 res_mut.zero_at(res_col, j);
382 })
383 }
384}
385
386unsafe impl VecZnxDftZeroImpl<FFT64> for FFT64 {
387 fn vec_znx_dft_zero_impl<R>(_module: &Module<FFT64>, res: &mut R)
388 where
389 R: VecZnxDftToMut<FFT64>,
390 {
391 res.to_mut().data.fill(0);
392 }
393}
394
395impl<D: DataRef> fmt::Display for VecZnxDft<D, FFT64> {
396 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
397 writeln!(
398 f,
399 "VecZnxDft(n={}, cols={}, size={})",
400 self.n, self.cols, self.size
401 )?;
402
403 for col in 0..self.cols {
404 writeln!(f, "Column {}:", col)?;
405 for size in 0..self.size {
406 let coeffs = self.at(col, size);
407 write!(f, " Size {}: [", size)?;
408
409 let max_show = 100;
410 let show_count = coeffs.len().min(max_show);
411
412 for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
413 if i > 0 {
414 write!(f, ", ")?;
415 }
416 write!(f, "{}", coeff)?;
417 }
418
419 if coeffs.len() > max_show {
420 write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
421 }
422
423 writeln!(f, "]")?;
424 }
425 }
426 Ok(())
427 }
428}