poulpy/poulpy-cpu-avx/src/reim/ifft_avx2_fma.rs
poulpy-cpu-avx/src/reim/ifft_avx2_fma.rsRaw
1use std::arch::x86_64::{
2 __m128d, __m256d, _mm_load_pd, _mm256_add_pd, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_mul_pd,
3 _mm256_permute2f128_pd, _mm256_set_m128d, _mm256_storeu_pd, _mm256_sub_pd, _mm256_unpackhi_pd, _mm256_unpacklo_pd,
4};
5
6use crate::reim::{as_arr, as_arr_mut};
7
8#[target_feature(enable = "avx2,fma")]
9pub(crate) fn ifft_avx2_fma(m: usize, omg: &[f64], data: &mut [f64]) {
10 if m < 16 {
11 use poulpy_hal::reference::fft64::reim::ifft_ref;
12 ifft_ref(m, omg, data);
13 return;
14 }
15
16 assert!(data.len() == 2 * m);
17 let (re, im) = data.split_at_mut(m);
18
19 if m == 16 {
20 ifft16_avx2_fma(as_arr_mut::<16, f64>(re), as_arr_mut::<16, f64>(im), as_arr::<16, f64>(omg))
21 } else if m <= 2048 {
22 ifft_bfs_16_avx2_fma(m, re, im, omg, 0);
23 } else {
24 ifft_rec_16_avx2_fma(m, re, im, omg, 0);
25 }
26}
27
28unsafe extern "sysv64" {
29 unsafe fn ifft16_avx2_fma_asm(re: *mut f64, im: *mut f64, omg: *const f64);
30}
31
32#[target_feature(enable = "avx2,fma")]
33fn ifft16_avx2_fma(re: &mut [f64; 16], im: &mut [f64; 16], omg: &[f64; 16]) {
34 unsafe {
35 ifft16_avx2_fma_asm(re.as_mut_ptr(), im.as_mut_ptr(), omg.as_ptr());
36 }
37}
38
39#[target_feature(enable = "avx2,fma")]
40fn ifft_rec_16_avx2_fma(m: usize, re: &mut [f64], im: &mut [f64], omg: &[f64], mut pos: usize) -> usize {
41 if m <= 2048 {
42 return ifft_bfs_16_avx2_fma(m, re, im, omg, pos);
43 };
44 let h: usize = m >> 1;
45 pos = ifft_rec_16_avx2_fma(h, re, im, omg, pos);
46 pos = ifft_rec_16_avx2_fma(h, &mut re[h..], &mut im[h..], omg, pos);
47 inv_twiddle_ifft_avx2_fma(h, re, im, *as_arr::<2, f64>(&omg[pos..]));
48 pos += 2;
49 pos
50}
51
52#[target_feature(enable = "avx2,fma")]
53fn ifft_bfs_16_avx2_fma(m: usize, re: &mut [f64], im: &mut [f64], omg: &[f64], mut pos: usize) -> usize {
54 let log_m: usize = (usize::BITS - (m - 1).leading_zeros()) as usize;
55
56 for off in (0..m).step_by(16) {
57 ifft16_avx2_fma(
58 as_arr_mut::<16, f64>(&mut re[off..]),
59 as_arr_mut::<16, f64>(&mut im[off..]),
60 as_arr::<16, f64>(&omg[pos..]),
61 );
62 pos += 16;
63 }
64
65 let mut h: usize = 16;
66 let m_half: usize = m >> 1;
67
68 while h < m_half {
69 let mm: usize = h << 2;
70 for off in (0..m).step_by(mm) {
71 inv_bitwiddle_ifft_avx2_fma(h, &mut re[off..], &mut im[off..], as_arr::<4, f64>(&omg[pos..]));
72 pos += 4;
73 }
74 h = mm;
75 }
76
77 if !log_m.is_multiple_of(2) {
78 inv_twiddle_ifft_avx2_fma(h, re, im, *as_arr::<2, f64>(&omg[pos..]));
79 pos += 2;
80 }
81
82 pos
83}
84
85#[target_feature(enable = "avx2,fma")]
86fn inv_twiddle_ifft_avx2_fma(h: usize, re: &mut [f64], im: &mut [f64], omg: [f64; 2]) {
87 unsafe {
88 let omx: __m128d = _mm_load_pd(omg.as_ptr());
89 let omra: __m256d = _mm256_set_m128d(omx, omx);
90 let omi: __m256d = _mm256_unpackhi_pd(omra, omra);
91 let omr: __m256d = _mm256_unpacklo_pd(omra, omra);
92 let mut r0: *mut f64 = re.as_mut_ptr();
93 let mut r1: *mut f64 = re.as_mut_ptr().add(h);
94 let mut i0: *mut f64 = im.as_mut_ptr();
95 let mut i1: *mut f64 = im.as_mut_ptr().add(h);
96 for _ in (0..h).step_by(4) {
97 let mut ur0: __m256d = _mm256_loadu_pd(r0);
98 let mut ur1: __m256d = _mm256_loadu_pd(r1);
99 let mut ui0: __m256d = _mm256_loadu_pd(i0);
100 let mut ui1: __m256d = _mm256_loadu_pd(i1);
101 let tra = _mm256_sub_pd(ur0, ur1);
102 let tia = _mm256_sub_pd(ui0, ui1);
103 ur0 = _mm256_add_pd(ur0, ur1);
104 ui0 = _mm256_add_pd(ui0, ui1);
105 ur1 = _mm256_mul_pd(omi, tia);
106 ui1 = _mm256_mul_pd(omi, tra);
107 ur1 = _mm256_fmsub_pd(omr, tra, ur1);
108 ui1 = _mm256_fmadd_pd(omr, tia, ui1);
109 _mm256_storeu_pd(r0, ur0);
110 _mm256_storeu_pd(r1, ur1);
111 _mm256_storeu_pd(i0, ui0);
112 _mm256_storeu_pd(i1, ui1);
113
114 r0 = r0.add(4);
115 r1 = r1.add(4);
116 i0 = i0.add(4);
117 i1 = i1.add(4);
118 }
119 }
120}
121
122#[target_feature(enable = "avx2,fma")]
123fn inv_bitwiddle_ifft_avx2_fma(h: usize, re: &mut [f64], im: &mut [f64], omg: &[f64; 4]) {
124 unsafe {
125 let mut r0: *mut f64 = re.as_mut_ptr();
126 let mut r1: *mut f64 = re.as_mut_ptr().add(h);
127 let mut r2: *mut f64 = re.as_mut_ptr().add(2 * h);
128 let mut r3: *mut f64 = re.as_mut_ptr().add(3 * h);
129 let mut i0: *mut f64 = im.as_mut_ptr();
130 let mut i1: *mut f64 = im.as_mut_ptr().add(h);
131 let mut i2: *mut f64 = im.as_mut_ptr().add(2 * h);
132 let mut i3: *mut f64 = im.as_mut_ptr().add(3 * h);
133 let om0: __m256d = _mm256_loadu_pd(omg.as_ptr());
134 let omb: __m256d = _mm256_permute2f128_pd(om0, om0, 0x11);
135 let oma: __m256d = _mm256_permute2f128_pd(om0, om0, 0x00);
136 let omai: __m256d = _mm256_unpackhi_pd(oma, oma);
137 let omar: __m256d = _mm256_unpacklo_pd(oma, oma);
138 let ombi: __m256d = _mm256_unpackhi_pd(omb, omb);
139 let ombr: __m256d = _mm256_unpacklo_pd(omb, omb);
140 for _ in (0..h).step_by(4) {
141 let mut ur0: __m256d = _mm256_loadu_pd(r0);
142 let mut ur1: __m256d = _mm256_loadu_pd(r1);
143 let mut ur2: __m256d = _mm256_loadu_pd(r2);
144 let mut ur3: __m256d = _mm256_loadu_pd(r3);
145 let mut ui0: __m256d = _mm256_loadu_pd(i0);
146 let mut ui1: __m256d = _mm256_loadu_pd(i1);
147 let mut ui2: __m256d = _mm256_loadu_pd(i2);
148 let mut ui3: __m256d = _mm256_loadu_pd(i3);
149
150 let mut tra: __m256d = _mm256_sub_pd(ur0, ur1);
151 let mut trb: __m256d = _mm256_sub_pd(ur2, ur3);
152 let mut tia: __m256d = _mm256_sub_pd(ui0, ui1);
153 let mut tib: __m256d = _mm256_sub_pd(ui2, ui3);
154 ur0 = _mm256_add_pd(ur0, ur1);
155 ur2 = _mm256_add_pd(ur2, ur3);
156 ui0 = _mm256_add_pd(ui0, ui1);
157 ui2 = _mm256_add_pd(ui2, ui3);
158 ur1 = _mm256_mul_pd(omai, tia);
159 ur3 = _mm256_mul_pd(omar, tib);
160 ui1 = _mm256_mul_pd(omai, tra);
161 ui3 = _mm256_mul_pd(omar, trb);
162 ur1 = _mm256_fmsub_pd(omar, tra, ur1);
163 ur3 = _mm256_fmadd_pd(omai, trb, ur3);
164 ui1 = _mm256_fmadd_pd(omar, tia, ui1);
165 ui3 = _mm256_fmsub_pd(omai, tib, ui3);
166
167 tra = _mm256_sub_pd(ur0, ur2);
168 trb = _mm256_sub_pd(ur1, ur3);
169 tia = _mm256_sub_pd(ui0, ui2);
170 tib = _mm256_sub_pd(ui1, ui3);
171 ur0 = _mm256_add_pd(ur0, ur2);
172 ur1 = _mm256_add_pd(ur1, ur3);
173 ui0 = _mm256_add_pd(ui0, ui2);
174 ui1 = _mm256_add_pd(ui1, ui3);
175 ur2 = _mm256_mul_pd(ombi, tia);
176 ur3 = _mm256_mul_pd(ombi, tib);
177 ui2 = _mm256_mul_pd(ombi, tra);
178 ui3 = _mm256_mul_pd(ombi, trb);
179 ur2 = _mm256_fmsub_pd(ombr, tra, ur2);
180 ur3 = _mm256_fmsub_pd(ombr, trb, ur3);
181 ui2 = _mm256_fmadd_pd(ombr, tia, ui2);
182 ui3 = _mm256_fmadd_pd(ombr, tib, ui3);
183
184 _mm256_storeu_pd(r0, ur0);
185 _mm256_storeu_pd(r1, ur1);
186 _mm256_storeu_pd(r2, ur2);
187 _mm256_storeu_pd(r3, ur3);
188 _mm256_storeu_pd(i0, ui0);
189 _mm256_storeu_pd(i1, ui1);
190 _mm256_storeu_pd(i2, ui2);
191 _mm256_storeu_pd(i3, ui3);
192
193 r0 = r0.add(4);
194 r1 = r1.add(4);
195 r2 = r2.add(4);
196 r3 = r3.add(4);
197 i0 = i0.add(4);
198 i1 = i1.add(4);
199 i2 = i2.add(4);
200 i3 = i3.add(4);
201 }
202 }
203}
204
205#[test]
206fn test_ifft_avx2_fma() {
207 use super::*;
208
209 #[target_feature(enable = "avx2,fma")]
210 fn internal(log_m: usize) {
211 use poulpy_hal::reference::fft64::reim::ReimIFFTRef;
212
213 let m: usize = 1 << log_m;
214
215 let table: ReimIFFTTable<f64> = ReimIFFTTable::<f64>::new(m);
216
217 let mut values_0: Vec<f64> = vec![0f64; m << 1];
218 let scale: f64 = 1.0f64 / m as f64;
219 values_0.iter_mut().enumerate().for_each(|(i, x)| *x = (i + 1) as f64 * scale);
220
221 let mut values_1: Vec<f64> = vec![0f64; m << 1];
222 values_1.iter_mut().zip(values_0.iter()).for_each(|(y, x)| *y = *x);
223
224 ReimIFFTAvx::reim_dft_execute(&table, &mut values_0);
225 ReimIFFTRef::reim_dft_execute(&table, &mut values_1);
226
227 let max_diff: f64 = 1.0 / ((1u64 << (53 - log_m - 1)) as f64);
228
229 for i in 0..m * 2 {
230 let diff: f64 = (values_0[i] - values_1[i]).abs();
231 assert!(diff <= max_diff, "{} -> {}-{} = {}", i, values_0[i], values_1[i], diff)
232 }
233 }
234
235 if std::is_x86_feature_detected!("avx2") {
236 for log_m in 0..16 {
237 unsafe { internal(log_m) }
238 }
239 } else {
240 eprintln!("skipping: CPU lacks avx2");
241 }
242}