poulpy/poulpy-backend/src/cpu_fft64_avx/module.rs
poulpy-backend/src/cpu_fft64_avx/module.rsRaw
1use std::ptr::NonNull;
2
3use poulpy_hal::{
4 layouts::{Backend, Module},
5 oep::ModuleNewImpl,
6 reference::{
7 fft64::{
8 reim::{
9 ReimAdd, ReimAddInplace, ReimAddMul, ReimCopy, ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimIFFTTable, ReimMul,
10 ReimMulInplace, ReimNegate, ReimNegateInplace, ReimSub, ReimSubABInplace, ReimSubBAInplace, ReimToZnx,
11 ReimToZnxInplace, ReimZero, reim_copy_ref, reim_zero_ref,
12 },
13 reim4::{
14 Reim4Extract1Blk, Reim4Mat1ColProd, Reim4Mat2Cols2ndColProd, Reim4Mat2ColsProd, Reim4Save1Blk, Reim4Save2Blks,
15 },
16 },
17 znx::{
18 ZnxAdd, ZnxAddInplace, ZnxAutomorphism, ZnxCopy, ZnxNegate, ZnxNegateInplace, ZnxNormalizeFinalStep,
19 ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStep, ZnxNormalizeFirstStepCarryOnly, ZnxNormalizeFirstStepInplace,
20 ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly, ZnxNormalizeMiddleStepInplace, ZnxRotate, ZnxSub,
21 ZnxSubABInplace, ZnxSubBAInplace, ZnxSwitchRing, ZnxZero, znx_copy_ref, znx_rotate, znx_zero_ref,
22 },
23 },
24};
25
26use crate::cpu_fft64_avx::{
27 FFT64Avx,
28 reim::{
29 ReimFFTAvx, ReimIFFTAvx, reim_add_avx2_fma, reim_add_inplace_avx2_fma, reim_addmul_avx2_fma, reim_from_znx_i64_bnd50_fma,
30 reim_mul_avx2_fma, reim_mul_inplace_avx2_fma, reim_negate_avx2_fma, reim_negate_inplace_avx2_fma,
31 reim_sub_ab_inplace_avx2_fma, reim_sub_avx2_fma, reim_sub_ba_inplace_avx2_fma, reim_to_znx_i64_inplace_bnd63_avx2_fma,
32 },
33 reim_to_znx_i64_bnd63_avx2_fma,
34 reim4::{
35 reim4_extract_1blk_from_reim_avx, reim4_save_1blk_to_reim_avx, reim4_save_2blk_to_reim_avx,
36 reim4_vec_mat1col_product_avx, reim4_vec_mat2cols_2ndcol_product_avx, reim4_vec_mat2cols_product_avx,
37 },
38 znx_avx::{
39 znx_add_avx, znx_add_inplace_avx, znx_automorphism_avx, znx_negate_avx, znx_negate_inplace_avx,
40 znx_normalize_final_step_avx, znx_normalize_final_step_inplace_avx, znx_normalize_first_step_avx,
41 znx_normalize_first_step_carry_only_avx, znx_normalize_first_step_inplace_avx, znx_normalize_middle_step_avx,
42 znx_normalize_middle_step_carry_only_avx, znx_normalize_middle_step_inplace_avx, znx_sub_ab_inplace_avx, znx_sub_avx,
43 znx_sub_ba_inplace_avx, znx_switch_ring_avx,
44 },
45};
46
47#[repr(C)]
48pub struct FFT64AvxHandle {
49 table_fft: ReimFFTTable<f64>,
50 table_ifft: ReimIFFTTable<f64>,
51}
52
53impl Backend for FFT64Avx {
54 type ScalarPrep = f64;
55 type ScalarBig = i64;
56 type Handle = FFT64AvxHandle;
57 unsafe fn destroy(handle: NonNull<Self::Handle>) {
58 unsafe {
59 drop(Box::from_raw(handle.as_ptr()));
60 }
61 }
62
63 fn layout_big_word_count() -> usize {
64 1
65 }
66
67 fn layout_prep_word_count() -> usize {
68 1
69 }
70}
71
72unsafe impl ModuleNewImpl<Self> for FFT64Avx {
73 fn new_impl(n: u64) -> Module<Self> {
74 if !std::arch::is_x86_feature_detected!("avx")
75 || !std::arch::is_x86_feature_detected!("avx2")
76 || !std::arch::is_x86_feature_detected!("fma")
77 {
78 panic!("arch must support avx2, avx and fma")
79 }
80
81 let handle: FFT64AvxHandle = FFT64AvxHandle {
82 table_fft: ReimFFTTable::new(n as usize >> 1),
83 table_ifft: ReimIFFTTable::new(n as usize >> 1),
84 };
85 // Leak Box to get a stable NonNull pointer
86 let ptr: NonNull<FFT64AvxHandle> = NonNull::from(Box::leak(Box::new(handle)));
87 unsafe { Module::from_nonnull(ptr, n) }
88 }
89}
90
91pub trait FFT64ModuleHandle {
92 fn get_fft_table(&self) -> &ReimFFTTable<f64>;
93 fn get_ifft_table(&self) -> &ReimIFFTTable<f64>;
94}
95
96impl FFT64ModuleHandle for Module<FFT64Avx> {
97 fn get_fft_table(&self) -> &ReimFFTTable<f64> {
98 let h: &FFT64AvxHandle = unsafe { &*self.ptr() };
99 &h.table_fft
100 }
101 fn get_ifft_table(&self) -> &ReimIFFTTable<f64> {
102 let h: &FFT64AvxHandle = unsafe { &*self.ptr() };
103 &h.table_ifft
104 }
105}
106
107impl ZnxAdd for FFT64Avx {
108 #[inline(always)]
109 fn znx_add(res: &mut [i64], a: &[i64], b: &[i64]) {
110 unsafe {
111 znx_add_avx(res, a, b);
112 }
113 }
114}
115
116impl ZnxAddInplace for FFT64Avx {
117 #[inline(always)]
118 fn znx_add_inplace(res: &mut [i64], a: &[i64]) {
119 unsafe {
120 znx_add_inplace_avx(res, a);
121 }
122 }
123}
124
125impl ZnxSub for FFT64Avx {
126 #[inline(always)]
127 fn znx_sub(res: &mut [i64], a: &[i64], b: &[i64]) {
128 unsafe {
129 znx_sub_avx(res, a, b);
130 }
131 }
132}
133
134impl ZnxSubABInplace for FFT64Avx {
135 #[inline(always)]
136 fn znx_sub_ab_inplace(res: &mut [i64], a: &[i64]) {
137 unsafe {
138 znx_sub_ab_inplace_avx(res, a);
139 }
140 }
141}
142
143impl ZnxSubBAInplace for FFT64Avx {
144 #[inline(always)]
145 fn znx_sub_ba_inplace(res: &mut [i64], a: &[i64]) {
146 unsafe {
147 znx_sub_ba_inplace_avx(res, a);
148 }
149 }
150}
151
152impl ZnxAutomorphism for FFT64Avx {
153 #[inline(always)]
154 fn znx_automorphism(p: i64, res: &mut [i64], a: &[i64]) {
155 unsafe {
156 znx_automorphism_avx(p, res, a);
157 }
158 }
159}
160
161impl ZnxCopy for FFT64Avx {
162 #[inline(always)]
163 fn znx_copy(res: &mut [i64], a: &[i64]) {
164 znx_copy_ref(res, a);
165 }
166}
167
168impl ZnxNegate for FFT64Avx {
169 #[inline(always)]
170 fn znx_negate(res: &mut [i64], src: &[i64]) {
171 unsafe {
172 znx_negate_avx(res, src);
173 }
174 }
175}
176
177impl ZnxNegateInplace for FFT64Avx {
178 #[inline(always)]
179 fn znx_negate_inplace(res: &mut [i64]) {
180 unsafe {
181 znx_negate_inplace_avx(res);
182 }
183 }
184}
185
186impl ZnxRotate for FFT64Avx {
187 #[inline(always)]
188 fn znx_rotate(p: i64, res: &mut [i64], src: &[i64]) {
189 znx_rotate::<Self>(p, res, src);
190 }
191}
192
193impl ZnxZero for FFT64Avx {
194 #[inline(always)]
195 fn znx_zero(res: &mut [i64]) {
196 znx_zero_ref(res);
197 }
198}
199
200impl ZnxSwitchRing for FFT64Avx {
201 #[inline(always)]
202 fn znx_switch_ring(res: &mut [i64], a: &[i64]) {
203 unsafe {
204 znx_switch_ring_avx(res, a);
205 }
206 }
207}
208
209impl ZnxNormalizeFinalStep for FFT64Avx {
210 #[inline(always)]
211 fn znx_normalize_final_step(basek: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
212 unsafe {
213 znx_normalize_final_step_avx(basek, lsh, x, a, carry);
214 }
215 }
216}
217
218impl ZnxNormalizeFinalStepInplace for FFT64Avx {
219 #[inline(always)]
220 fn znx_normalize_final_step_inplace(basek: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
221 unsafe {
222 znx_normalize_final_step_inplace_avx(basek, lsh, x, carry);
223 }
224 }
225}
226
227impl ZnxNormalizeFirstStep for FFT64Avx {
228 #[inline(always)]
229 fn znx_normalize_first_step(basek: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
230 unsafe {
231 znx_normalize_first_step_avx(basek, lsh, x, a, carry);
232 }
233 }
234}
235
236impl ZnxNormalizeFirstStepCarryOnly for FFT64Avx {
237 #[inline(always)]
238 fn znx_normalize_first_step_carry_only(basek: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
239 unsafe {
240 znx_normalize_first_step_carry_only_avx(basek, lsh, x, carry);
241 }
242 }
243}
244
245impl ZnxNormalizeFirstStepInplace for FFT64Avx {
246 #[inline(always)]
247 fn znx_normalize_first_step_inplace(basek: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
248 unsafe {
249 znx_normalize_first_step_inplace_avx(basek, lsh, x, carry);
250 }
251 }
252}
253
254impl ZnxNormalizeMiddleStep for FFT64Avx {
255 #[inline(always)]
256 fn znx_normalize_middle_step(basek: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
257 unsafe {
258 znx_normalize_middle_step_avx(basek, lsh, x, a, carry);
259 }
260 }
261}
262
263impl ZnxNormalizeMiddleStepCarryOnly for FFT64Avx {
264 #[inline(always)]
265 fn znx_normalize_middle_step_carry_only(basek: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
266 unsafe {
267 znx_normalize_middle_step_carry_only_avx(basek, lsh, x, carry);
268 }
269 }
270}
271
272impl ZnxNormalizeMiddleStepInplace for FFT64Avx {
273 #[inline(always)]
274 fn znx_normalize_middle_step_inplace(basek: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
275 unsafe {
276 znx_normalize_middle_step_inplace_avx(basek, lsh, x, carry);
277 }
278 }
279}
280
281impl ReimDFTExecute<ReimFFTTable<f64>, f64> for FFT64Avx {
282 #[inline(always)]
283 fn reim_dft_execute(table: &ReimFFTTable<f64>, data: &mut [f64]) {
284 ReimFFTAvx::reim_dft_execute(table, data);
285 }
286}
287
288impl ReimDFTExecute<ReimIFFTTable<f64>, f64> for FFT64Avx {
289 #[inline(always)]
290 fn reim_dft_execute(table: &ReimIFFTTable<f64>, data: &mut [f64]) {
291 ReimIFFTAvx::reim_dft_execute(table, data);
292 }
293}
294
295impl ReimFromZnx for FFT64Avx {
296 #[inline(always)]
297 fn reim_from_znx(res: &mut [f64], a: &[i64]) {
298 unsafe {
299 reim_from_znx_i64_bnd50_fma(res, a);
300 }
301 }
302}
303
304impl ReimToZnx for FFT64Avx {
305 #[inline(always)]
306 fn reim_to_znx(res: &mut [i64], divisor: f64, a: &[f64]) {
307 unsafe {
308 reim_to_znx_i64_bnd63_avx2_fma(res, divisor, a);
309 }
310 }
311}
312
313impl ReimToZnxInplace for FFT64Avx {
314 #[inline(always)]
315 fn reim_to_znx_inplace(res: &mut [f64], divisor: f64) {
316 unsafe {
317 reim_to_znx_i64_inplace_bnd63_avx2_fma(res, divisor);
318 }
319 }
320}
321
322impl ReimAdd for FFT64Avx {
323 #[inline(always)]
324 fn reim_add(res: &mut [f64], a: &[f64], b: &[f64]) {
325 unsafe {
326 reim_add_avx2_fma(res, a, b);
327 }
328 }
329}
330
331impl ReimAddInplace for FFT64Avx {
332 #[inline(always)]
333 fn reim_add_inplace(res: &mut [f64], a: &[f64]) {
334 unsafe {
335 reim_add_inplace_avx2_fma(res, a);
336 }
337 }
338}
339
340impl ReimSub for FFT64Avx {
341 #[inline(always)]
342 fn reim_sub(res: &mut [f64], a: &[f64], b: &[f64]) {
343 unsafe {
344 reim_sub_avx2_fma(res, a, b);
345 }
346 }
347}
348
349impl ReimSubABInplace for FFT64Avx {
350 #[inline(always)]
351 fn reim_sub_ab_inplace(res: &mut [f64], a: &[f64]) {
352 unsafe {
353 reim_sub_ab_inplace_avx2_fma(res, a);
354 }
355 }
356}
357
358impl ReimSubBAInplace for FFT64Avx {
359 #[inline(always)]
360 fn reim_sub_ba_inplace(res: &mut [f64], a: &[f64]) {
361 unsafe {
362 reim_sub_ba_inplace_avx2_fma(res, a);
363 }
364 }
365}
366
367impl ReimNegate for FFT64Avx {
368 #[inline(always)]
369 fn reim_negate(res: &mut [f64], a: &[f64]) {
370 unsafe {
371 reim_negate_avx2_fma(res, a);
372 }
373 }
374}
375
376impl ReimNegateInplace for FFT64Avx {
377 #[inline(always)]
378 fn reim_negate_inplace(res: &mut [f64]) {
379 unsafe {
380 reim_negate_inplace_avx2_fma(res);
381 }
382 }
383}
384
385impl ReimMul for FFT64Avx {
386 #[inline(always)]
387 fn reim_mul(res: &mut [f64], a: &[f64], b: &[f64]) {
388 unsafe {
389 reim_mul_avx2_fma(res, a, b);
390 }
391 }
392}
393
394impl ReimMulInplace for FFT64Avx {
395 #[inline(always)]
396 fn reim_mul_inplace(res: &mut [f64], a: &[f64]) {
397 unsafe {
398 reim_mul_inplace_avx2_fma(res, a);
399 }
400 }
401}
402
403impl ReimAddMul for FFT64Avx {
404 #[inline(always)]
405 fn reim_addmul(res: &mut [f64], a: &[f64], b: &[f64]) {
406 unsafe {
407 reim_addmul_avx2_fma(res, a, b);
408 }
409 }
410}
411
412impl ReimCopy for FFT64Avx {
413 #[inline(always)]
414 fn reim_copy(res: &mut [f64], a: &[f64]) {
415 reim_copy_ref(res, a);
416 }
417}
418
419impl ReimZero for FFT64Avx {
420 #[inline(always)]
421 fn reim_zero(res: &mut [f64]) {
422 reim_zero_ref(res);
423 }
424}
425
426impl Reim4Extract1Blk for FFT64Avx {
427 #[inline(always)]
428 fn reim4_extract_1blk(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
429 unsafe {
430 reim4_extract_1blk_from_reim_avx(m, rows, blk, dst, src);
431 }
432 }
433}
434
435impl Reim4Save1Blk for FFT64Avx {
436 #[inline(always)]
437 fn reim4_save_1blk<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
438 unsafe {
439 reim4_save_1blk_to_reim_avx::<OVERWRITE>(m, blk, dst, src);
440 }
441 }
442}
443
444impl Reim4Save2Blks for FFT64Avx {
445 #[inline(always)]
446 fn reim4_save_2blks<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
447 unsafe {
448 reim4_save_2blk_to_reim_avx::<OVERWRITE>(m, blk, dst, src);
449 }
450 }
451}
452
453impl Reim4Mat1ColProd for FFT64Avx {
454 #[inline(always)]
455 fn reim4_mat1col_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
456 unsafe {
457 reim4_vec_mat1col_product_avx(nrows, dst, u, v);
458 }
459 }
460}
461
462impl Reim4Mat2ColsProd for FFT64Avx {
463 #[inline(always)]
464 fn reim4_mat2cols_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
465 unsafe {
466 reim4_vec_mat2cols_product_avx(nrows, dst, u, v);
467 }
468 }
469}
470
471impl Reim4Mat2Cols2ndColProd for FFT64Avx {
472 #[inline(always)]
473 fn reim4_mat2cols_2ndcol_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
474 unsafe {
475 reim4_vec_mat2cols_2ndcol_product_avx(nrows, dst, u, v);
476 }
477 }
478}