poulpy/poulpy-cpu-avx/src/module.rs
poulpy-cpu-avx/src/module.rsRaw
1use std::ptr::NonNull;
2
3use poulpy_hal::{
4 layouts::{Backend, Module},
5 oep::ModuleNewImpl,
6 reference::{
7 fft64::{
8 convolution::{
9 I64ConvolutionByConst1Coeff, I64ConvolutionByConst2Coeffs, I64Extract1BlkContiguous, I64Save1BlkContiguous,
10 },
11 reim::{
12 ReimAdd, ReimAddInplace, ReimAddMul, ReimCopy, ReimDFTExecute, ReimFFTTable, ReimFromZnx, ReimIFFTTable, ReimMul,
13 ReimMulInplace, ReimNegate, ReimNegateInplace, ReimSub, ReimSubInplace, ReimSubNegateInplace, ReimToZnx,
14 ReimToZnxInplace, ReimZero, reim_copy_ref, reim_zero_ref,
15 },
16 reim4::{
17 Reim4Convolution1Coeff, Reim4Convolution2Coeffs, Reim4ConvolutionByRealConst1Coeff,
18 Reim4ConvolutionByRealConst2Coeffs, Reim4Extract1BlkContiguous, Reim4Mat1ColProd, Reim4Mat2Cols2ndColProd,
19 Reim4Mat2ColsProd, Reim4Save1Blk, Reim4Save1BlkContiguous, Reim4Save2Blks,
20 },
21 },
22 znx::{
23 ZnxAdd, ZnxAddInplace, ZnxAutomorphism, ZnxCopy, ZnxExtractDigitAddMul, ZnxMulAddPowerOfTwo, ZnxMulPowerOfTwo,
24 ZnxMulPowerOfTwoInplace, ZnxNegate, ZnxNegateInplace, ZnxNormalizeDigit, ZnxNormalizeFinalStep,
25 ZnxNormalizeFinalStepInplace, ZnxNormalizeFirstStep, ZnxNormalizeFirstStepCarryOnly, ZnxNormalizeFirstStepInplace,
26 ZnxNormalizeMiddleStep, ZnxNormalizeMiddleStepCarryOnly, ZnxNormalizeMiddleStepInplace, ZnxRotate, ZnxSub,
27 ZnxSubInplace, ZnxSubNegateInplace, ZnxSwitchRing, ZnxZero, znx_copy_ref, znx_rotate, znx_zero_ref,
28 },
29 },
30};
31
32use crate::{
33 FFT64Avx,
34 convolution::{
35 i64_convolution_by_const_1coeff_avx, i64_convolution_by_real_const_2coeffs_avx, i64_extract_1blk_contiguous_avx,
36 i64_save_1blk_contiguous_avx,
37 },
38 reim::{
39 ReimFFTAvx, ReimIFFTAvx, reim_add_avx2_fma, reim_add_inplace_avx2_fma, reim_addmul_avx2_fma, reim_from_znx_i64_bnd50_fma,
40 reim_mul_avx2_fma, reim_mul_inplace_avx2_fma, reim_negate_avx2_fma, reim_negate_inplace_avx2_fma, reim_sub_avx2_fma,
41 reim_sub_inplace_avx2_fma, reim_sub_negate_inplace_avx2_fma, reim_to_znx_i64_inplace_bnd63_avx2_fma,
42 },
43 reim_to_znx_i64_bnd63_avx2_fma,
44 reim4::{
45 reim4_convolution_1coeff_avx, reim4_convolution_2coeffs_avx, reim4_convolution_by_real_const_1coeff_avx,
46 reim4_convolution_by_real_const_2coeffs_avx, reim4_extract_1blk_from_reim_contiguous_avx, reim4_save_1blk_to_reim_avx,
47 reim4_save_1blk_to_reim_contiguous_avx, reim4_save_2blk_to_reim_avx, reim4_vec_mat1col_product_avx,
48 reim4_vec_mat2cols_2ndcol_product_avx, reim4_vec_mat2cols_product_avx,
49 },
50 znx_avx::{
51 znx_add_avx, znx_add_inplace_avx, znx_automorphism_avx, znx_extract_digit_addmul_avx, znx_mul_add_power_of_two_avx,
52 znx_mul_power_of_two_avx, znx_mul_power_of_two_inplace_avx, znx_negate_avx, znx_negate_inplace_avx,
53 znx_normalize_digit_avx, znx_normalize_final_step_avx, znx_normalize_final_step_inplace_avx,
54 znx_normalize_first_step_avx, znx_normalize_first_step_carry_only_avx, znx_normalize_first_step_inplace_avx,
55 znx_normalize_middle_step_avx, znx_normalize_middle_step_carry_only_avx, znx_normalize_middle_step_inplace_avx,
56 znx_sub_avx, znx_sub_inplace_avx, znx_sub_negate_inplace_avx, znx_switch_ring_avx,
57 },
58};
59
60#[repr(C)]
61pub struct FFT64AvxHandle {
62 table_fft: ReimFFTTable<f64>,
63 table_ifft: ReimIFFTTable<f64>,
64}
65
66impl Backend for FFT64Avx {
67 type ScalarPrep = f64;
68 type ScalarBig = i64;
69 type Handle = FFT64AvxHandle;
70 unsafe fn destroy(handle: NonNull<Self::Handle>) {
71 unsafe {
72 drop(Box::from_raw(handle.as_ptr()));
73 }
74 }
75
76 fn layout_big_word_count() -> usize {
77 1
78 }
79
80 fn layout_prep_word_count() -> usize {
81 1
82 }
83}
84
85unsafe impl ModuleNewImpl<Self> for FFT64Avx {
86 fn new_impl(n: u64) -> Module<Self> {
87 if !std::arch::is_x86_feature_detected!("avx")
88 || !std::arch::is_x86_feature_detected!("avx2")
89 || !std::arch::is_x86_feature_detected!("fma")
90 {
91 panic!("arch must support avx2, avx and fma")
92 }
93
94 let handle: FFT64AvxHandle = FFT64AvxHandle {
95 table_fft: ReimFFTTable::new(n as usize >> 1),
96 table_ifft: ReimIFFTTable::new(n as usize >> 1),
97 };
98 // Leak Box to get a stable NonNull pointer
99 let ptr: NonNull<FFT64AvxHandle> = NonNull::from(Box::leak(Box::new(handle)));
100 unsafe { Module::from_nonnull(ptr, n) }
101 }
102}
103
104pub trait FFT64ModuleHandle {
105 fn get_fft_table(&self) -> &ReimFFTTable<f64>;
106 fn get_ifft_table(&self) -> &ReimIFFTTable<f64>;
107}
108
109impl FFT64ModuleHandle for Module<FFT64Avx> {
110 fn get_fft_table(&self) -> &ReimFFTTable<f64> {
111 let h: &FFT64AvxHandle = unsafe { &*self.ptr() };
112 &h.table_fft
113 }
114 fn get_ifft_table(&self) -> &ReimIFFTTable<f64> {
115 let h: &FFT64AvxHandle = unsafe { &*self.ptr() };
116 &h.table_ifft
117 }
118}
119
120impl ZnxAdd for FFT64Avx {
121 #[inline(always)]
122 fn znx_add(res: &mut [i64], a: &[i64], b: &[i64]) {
123 unsafe {
124 znx_add_avx(res, a, b);
125 }
126 }
127}
128
129impl ZnxAddInplace for FFT64Avx {
130 #[inline(always)]
131 fn znx_add_inplace(res: &mut [i64], a: &[i64]) {
132 unsafe {
133 znx_add_inplace_avx(res, a);
134 }
135 }
136}
137
138impl ZnxSub for FFT64Avx {
139 #[inline(always)]
140 fn znx_sub(res: &mut [i64], a: &[i64], b: &[i64]) {
141 unsafe {
142 znx_sub_avx(res, a, b);
143 }
144 }
145}
146
147impl ZnxSubInplace for FFT64Avx {
148 #[inline(always)]
149 fn znx_sub_inplace(res: &mut [i64], a: &[i64]) {
150 unsafe {
151 znx_sub_inplace_avx(res, a);
152 }
153 }
154}
155
156impl ZnxSubNegateInplace for FFT64Avx {
157 #[inline(always)]
158 fn znx_sub_negate_inplace(res: &mut [i64], a: &[i64]) {
159 unsafe {
160 znx_sub_negate_inplace_avx(res, a);
161 }
162 }
163}
164
165impl ZnxAutomorphism for FFT64Avx {
166 #[inline(always)]
167 fn znx_automorphism(p: i64, res: &mut [i64], a: &[i64]) {
168 unsafe {
169 znx_automorphism_avx(p, res, a);
170 }
171 }
172}
173
174impl ZnxCopy for FFT64Avx {
175 #[inline(always)]
176 fn znx_copy(res: &mut [i64], a: &[i64]) {
177 znx_copy_ref(res, a);
178 }
179}
180
181impl ZnxNegate for FFT64Avx {
182 #[inline(always)]
183 fn znx_negate(res: &mut [i64], src: &[i64]) {
184 unsafe {
185 znx_negate_avx(res, src);
186 }
187 }
188}
189
190impl ZnxNegateInplace for FFT64Avx {
191 #[inline(always)]
192 fn znx_negate_inplace(res: &mut [i64]) {
193 unsafe {
194 znx_negate_inplace_avx(res);
195 }
196 }
197}
198
199impl ZnxMulAddPowerOfTwo for FFT64Avx {
200 #[inline(always)]
201 fn znx_muladd_power_of_two(k: i64, res: &mut [i64], a: &[i64]) {
202 unsafe {
203 znx_mul_add_power_of_two_avx(k, res, a);
204 }
205 }
206}
207
208impl ZnxMulPowerOfTwo for FFT64Avx {
209 #[inline(always)]
210 fn znx_mul_power_of_two(k: i64, res: &mut [i64], a: &[i64]) {
211 unsafe {
212 znx_mul_power_of_two_avx(k, res, a);
213 }
214 }
215}
216
217impl ZnxMulPowerOfTwoInplace for FFT64Avx {
218 #[inline(always)]
219 fn znx_mul_power_of_two_inplace(k: i64, res: &mut [i64]) {
220 unsafe {
221 znx_mul_power_of_two_inplace_avx(k, res);
222 }
223 }
224}
225
226impl ZnxRotate for FFT64Avx {
227 #[inline(always)]
228 fn znx_rotate(p: i64, res: &mut [i64], src: &[i64]) {
229 znx_rotate::<Self>(p, res, src);
230 }
231}
232
233impl ZnxZero for FFT64Avx {
234 #[inline(always)]
235 fn znx_zero(res: &mut [i64]) {
236 znx_zero_ref(res);
237 }
238}
239
240impl ZnxSwitchRing for FFT64Avx {
241 #[inline(always)]
242 fn znx_switch_ring(res: &mut [i64], a: &[i64]) {
243 unsafe {
244 znx_switch_ring_avx(res, a);
245 }
246 }
247}
248
249impl ZnxNormalizeFinalStep for FFT64Avx {
250 #[inline(always)]
251 fn znx_normalize_final_step(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
252 unsafe {
253 znx_normalize_final_step_avx(base2k, lsh, x, a, carry);
254 }
255 }
256}
257
258impl ZnxNormalizeFinalStepInplace for FFT64Avx {
259 #[inline(always)]
260 fn znx_normalize_final_step_inplace(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
261 unsafe {
262 znx_normalize_final_step_inplace_avx(base2k, lsh, x, carry);
263 }
264 }
265}
266
267impl ZnxNormalizeFirstStep for FFT64Avx {
268 #[inline(always)]
269 fn znx_normalize_first_step(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
270 unsafe {
271 znx_normalize_first_step_avx(base2k, lsh, x, a, carry);
272 }
273 }
274}
275
276impl ZnxNormalizeFirstStepCarryOnly for FFT64Avx {
277 #[inline(always)]
278 fn znx_normalize_first_step_carry_only(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
279 unsafe {
280 znx_normalize_first_step_carry_only_avx(base2k, lsh, x, carry);
281 }
282 }
283}
284
285impl ZnxNormalizeFirstStepInplace for FFT64Avx {
286 #[inline(always)]
287 fn znx_normalize_first_step_inplace(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
288 unsafe {
289 znx_normalize_first_step_inplace_avx(base2k, lsh, x, carry);
290 }
291 }
292}
293
294impl ZnxNormalizeMiddleStep for FFT64Avx {
295 #[inline(always)]
296 fn znx_normalize_middle_step(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
297 unsafe {
298 znx_normalize_middle_step_avx(base2k, lsh, x, a, carry);
299 }
300 }
301}
302
303impl ZnxNormalizeMiddleStepCarryOnly for FFT64Avx {
304 #[inline(always)]
305 fn znx_normalize_middle_step_carry_only(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
306 unsafe {
307 znx_normalize_middle_step_carry_only_avx(base2k, lsh, x, carry);
308 }
309 }
310}
311
312impl ZnxNormalizeMiddleStepInplace for FFT64Avx {
313 #[inline(always)]
314 fn znx_normalize_middle_step_inplace(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
315 unsafe {
316 znx_normalize_middle_step_inplace_avx(base2k, lsh, x, carry);
317 }
318 }
319}
320
321impl ZnxExtractDigitAddMul for FFT64Avx {
322 #[inline(always)]
323 fn znx_extract_digit_addmul(base2k: usize, lsh: usize, res: &mut [i64], src: &mut [i64]) {
324 unsafe {
325 znx_extract_digit_addmul_avx(base2k, lsh, res, src);
326 }
327 }
328}
329
330impl ZnxNormalizeDigit for FFT64Avx {
331 #[inline(always)]
332 fn znx_normalize_digit(base2k: usize, res: &mut [i64], src: &mut [i64]) {
333 unsafe {
334 znx_normalize_digit_avx(base2k, res, src);
335 }
336 }
337}
338
339impl ReimDFTExecute<ReimFFTTable<f64>, f64> for FFT64Avx {
340 #[inline(always)]
341 fn reim_dft_execute(table: &ReimFFTTable<f64>, data: &mut [f64]) {
342 ReimFFTAvx::reim_dft_execute(table, data);
343 }
344}
345
346impl ReimDFTExecute<ReimIFFTTable<f64>, f64> for FFT64Avx {
347 #[inline(always)]
348 fn reim_dft_execute(table: &ReimIFFTTable<f64>, data: &mut [f64]) {
349 ReimIFFTAvx::reim_dft_execute(table, data);
350 }
351}
352
353impl ReimFromZnx for FFT64Avx {
354 #[inline(always)]
355 fn reim_from_znx(res: &mut [f64], a: &[i64]) {
356 unsafe {
357 reim_from_znx_i64_bnd50_fma(res, a);
358 }
359 }
360}
361
362impl ReimToZnx for FFT64Avx {
363 #[inline(always)]
364 fn reim_to_znx(res: &mut [i64], divisor: f64, a: &[f64]) {
365 unsafe {
366 reim_to_znx_i64_bnd63_avx2_fma(res, divisor, a);
367 }
368 }
369}
370
371impl ReimToZnxInplace for FFT64Avx {
372 #[inline(always)]
373 fn reim_to_znx_inplace(res: &mut [f64], divisor: f64) {
374 unsafe {
375 reim_to_znx_i64_inplace_bnd63_avx2_fma(res, divisor);
376 }
377 }
378}
379
380impl ReimAdd for FFT64Avx {
381 #[inline(always)]
382 fn reim_add(res: &mut [f64], a: &[f64], b: &[f64]) {
383 unsafe {
384 reim_add_avx2_fma(res, a, b);
385 }
386 }
387}
388
389impl ReimAddInplace for FFT64Avx {
390 #[inline(always)]
391 fn reim_add_inplace(res: &mut [f64], a: &[f64]) {
392 unsafe {
393 reim_add_inplace_avx2_fma(res, a);
394 }
395 }
396}
397
398impl ReimSub for FFT64Avx {
399 #[inline(always)]
400 fn reim_sub(res: &mut [f64], a: &[f64], b: &[f64]) {
401 unsafe {
402 reim_sub_avx2_fma(res, a, b);
403 }
404 }
405}
406
407impl ReimSubInplace for FFT64Avx {
408 #[inline(always)]
409 fn reim_sub_inplace(res: &mut [f64], a: &[f64]) {
410 unsafe {
411 reim_sub_inplace_avx2_fma(res, a);
412 }
413 }
414}
415
416impl ReimSubNegateInplace for FFT64Avx {
417 #[inline(always)]
418 fn reim_sub_negate_inplace(res: &mut [f64], a: &[f64]) {
419 unsafe {
420 reim_sub_negate_inplace_avx2_fma(res, a);
421 }
422 }
423}
424
425impl ReimNegate for FFT64Avx {
426 #[inline(always)]
427 fn reim_negate(res: &mut [f64], a: &[f64]) {
428 unsafe {
429 reim_negate_avx2_fma(res, a);
430 }
431 }
432}
433
434impl ReimNegateInplace for FFT64Avx {
435 #[inline(always)]
436 fn reim_negate_inplace(res: &mut [f64]) {
437 unsafe {
438 reim_negate_inplace_avx2_fma(res);
439 }
440 }
441}
442
443impl ReimMul for FFT64Avx {
444 #[inline(always)]
445 fn reim_mul(res: &mut [f64], a: &[f64], b: &[f64]) {
446 unsafe {
447 reim_mul_avx2_fma(res, a, b);
448 }
449 }
450}
451
452impl ReimMulInplace for FFT64Avx {
453 #[inline(always)]
454 fn reim_mul_inplace(res: &mut [f64], a: &[f64]) {
455 unsafe {
456 reim_mul_inplace_avx2_fma(res, a);
457 }
458 }
459}
460
461impl ReimAddMul for FFT64Avx {
462 #[inline(always)]
463 fn reim_addmul(res: &mut [f64], a: &[f64], b: &[f64]) {
464 unsafe {
465 reim_addmul_avx2_fma(res, a, b);
466 }
467 }
468}
469
470impl ReimCopy for FFT64Avx {
471 #[inline(always)]
472 fn reim_copy(res: &mut [f64], a: &[f64]) {
473 reim_copy_ref(res, a);
474 }
475}
476
477impl ReimZero for FFT64Avx {
478 #[inline(always)]
479 fn reim_zero(res: &mut [f64]) {
480 reim_zero_ref(res);
481 }
482}
483
484impl Reim4Convolution1Coeff for FFT64Avx {
485 #[inline(always)]
486 fn reim4_convolution_1coeff(k: usize, dst: &mut [f64; 8], a: &[f64], a_size: usize, b: &[f64], b_size: usize) {
487 unsafe {
488 reim4_convolution_1coeff_avx(k, dst, a, a_size, b, b_size);
489 }
490 }
491}
492
493impl Reim4Convolution2Coeffs for FFT64Avx {
494 #[inline(always)]
495 fn reim4_convolution_2coeffs(k: usize, dst: &mut [f64; 16], a: &[f64], a_size: usize, b: &[f64], b_size: usize) {
496 unsafe {
497 reim4_convolution_2coeffs_avx(k, dst, a, a_size, b, b_size);
498 }
499 }
500}
501
502impl Reim4ConvolutionByRealConst1Coeff for FFT64Avx {
503 #[inline(always)]
504 fn reim4_convolution_by_real_const_1coeff(k: usize, dst: &mut [f64; 8], a: &[f64], a_size: usize, b: &[f64]) {
505 unsafe {
506 reim4_convolution_by_real_const_1coeff_avx(k, dst, a, a_size, b);
507 }
508 }
509}
510
511impl Reim4ConvolutionByRealConst2Coeffs for FFT64Avx {
512 #[inline(always)]
513 fn reim4_convolution_by_real_const_2coeffs(k: usize, dst: &mut [f64; 16], a: &[f64], a_size: usize, b: &[f64]) {
514 unsafe {
515 reim4_convolution_by_real_const_2coeffs_avx(k, dst, a, a_size, b);
516 }
517 }
518}
519
520impl Reim4Extract1BlkContiguous for FFT64Avx {
521 #[inline(always)]
522 fn reim4_extract_1blk_contiguous(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
523 unsafe {
524 reim4_extract_1blk_from_reim_contiguous_avx(m, rows, blk, dst, src);
525 }
526 }
527}
528
529impl Reim4Save1BlkContiguous for FFT64Avx {
530 fn reim4_save_1blk_contiguous(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
531 unsafe {
532 reim4_save_1blk_to_reim_contiguous_avx(m, rows, blk, dst, src);
533 }
534 }
535}
536
537impl Reim4Save1Blk for FFT64Avx {
538 #[inline(always)]
539 fn reim4_save_1blk<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
540 unsafe {
541 reim4_save_1blk_to_reim_avx::<OVERWRITE>(m, blk, dst, src);
542 }
543 }
544}
545
546impl Reim4Save2Blks for FFT64Avx {
547 #[inline(always)]
548 fn reim4_save_2blks<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
549 unsafe {
550 reim4_save_2blk_to_reim_avx::<OVERWRITE>(m, blk, dst, src);
551 }
552 }
553}
554
555impl Reim4Mat1ColProd for FFT64Avx {
556 #[inline(always)]
557 fn reim4_mat1col_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
558 unsafe {
559 reim4_vec_mat1col_product_avx(nrows, dst, u, v);
560 }
561 }
562}
563
564impl Reim4Mat2ColsProd for FFT64Avx {
565 #[inline(always)]
566 fn reim4_mat2cols_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
567 unsafe {
568 reim4_vec_mat2cols_product_avx(nrows, dst, u, v);
569 }
570 }
571}
572
573impl Reim4Mat2Cols2ndColProd for FFT64Avx {
574 #[inline(always)]
575 fn reim4_mat2cols_2ndcol_prod(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
576 unsafe {
577 reim4_vec_mat2cols_2ndcol_product_avx(nrows, dst, u, v);
578 }
579 }
580}
581
582impl I64ConvolutionByConst1Coeff for FFT64Avx {
583 #[inline(always)]
584 fn i64_convolution_by_const_1coeff(k: usize, dst: &mut [i64; 8], a: &[i64], a_size: usize, b: &[i64]) {
585 unsafe {
586 i64_convolution_by_const_1coeff_avx(k, dst, a, a_size, b);
587 }
588 }
589}
590
591impl I64ConvolutionByConst2Coeffs for FFT64Avx {
592 #[inline(always)]
593 fn i64_convolution_by_const_2coeffs(k: usize, dst: &mut [i64; 16], a: &[i64], a_size: usize, b: &[i64]) {
594 unsafe {
595 i64_convolution_by_real_const_2coeffs_avx(k, dst, a, a_size, b);
596 }
597 }
598}
599
600impl I64Save1BlkContiguous for FFT64Avx {
601 #[inline(always)]
602 fn i64_save_1blk_contiguous(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]) {
603 unsafe {
604 i64_save_1blk_contiguous_avx(n, offset, rows, blk, dst, src);
605 }
606 }
607}
608
609impl I64Extract1BlkContiguous for FFT64Avx {
610 #[inline(always)]
611 fn i64_extract_1blk_contiguous(n: usize, offset: usize, rows: usize, blk: usize, dst: &mut [i64], src: &[i64]) {
612 unsafe {
613 i64_extract_1blk_contiguous_avx(n, offset, rows, blk, dst, src);
614 }
615 }
616}