poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/cplx/cplx_ifft_avx2_fma.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/cplx/cplx_ifft_avx2_fma.cRaw
1#include <immintrin.h>
2#include <math.h>
3#include <string.h>
4
5#include "cplx_fft_internal.h"
6#include "cplx_fft_private.h"
7
8typedef double D4MEM[4];
9typedef double D2MEM[2];
10
11/**
12 * @brief complex ifft via bfs strategy (for m between 2 and 8)
13 * @param dat the data to run the algorithm on
14 * @param omg precomputed tables (must have been filled with fill_omega)
15 * @param m ring dimension of the FFT (modulo X^m-i)
16 */
17void cplx_ifft_avx2_fma_bfs_2(D4MEM* dat, const D2MEM** omga, uint32_t m) {
18 double* data = (double*)dat;
19 D4MEM* const finaldd = (D4MEM*)(data + 2 * m);
20 {
21 // loop with h = 1
22 // we do not do any particular optimization in this loop,
23 // since this function is only called for small dimensions
24 D4MEM* dd = (D4MEM*)data;
25 do {
26 /*
27 BEGIN_TEMPLATE
28 const __m256d ab% = _mm256_loadu_pd(dd[0+2*%]);
29 const __m256d cd% = _mm256_loadu_pd(dd[1+2*%]);
30 const __m256d ac% = _mm256_permute2f128_pd(ab%, cd%, 0b100000);
31 const __m256d bd% = _mm256_permute2f128_pd(ab%, cd%, 0b110001);
32 const __m256d sum% = _mm256_add_pd(ac%, bd%);
33 const __m256d diff% = _mm256_sub_pd(ac%, bd%);
34 const __m256d diffbar% = _mm256_shuffle_pd(diff%, diff%, 5);
35 const __m256d om% = _mm256_load_pd((*omg)[0+%]);
36 const __m256d omre% = _mm256_unpacklo_pd(om%, om%);
37 const __m256d omim% = _mm256_unpackhi_pd(om%, om%);
38 const __m256d t1% = _mm256_mul_pd(diffbar%, omim%);
39 const __m256d t2% = _mm256_fmaddsub_pd(diff%, omre%, t1%);
40 const __m256d newab% = _mm256_permute2f128_pd(sum%, t2%, 0b100000);
41 const __m256d newcd% = _mm256_permute2f128_pd(sum%, t2%, 0b110001);
42 _mm256_storeu_pd(dd[0+2*%], newab%);
43 _mm256_storeu_pd(dd[1+2*%], newcd%);
44 dd += 2*@;
45 *omg += 2*@;
46 END_TEMPLATE
47 */
48 // BEGIN_INTERLEAVE 1
49 const __m256d ab0 = _mm256_loadu_pd(dd[0 + 2 * 0]);
50 const __m256d cd0 = _mm256_loadu_pd(dd[1 + 2 * 0]);
51 const __m256d ac0 = _mm256_permute2f128_pd(ab0, cd0, 0b100000);
52 const __m256d bd0 = _mm256_permute2f128_pd(ab0, cd0, 0b110001);
53 const __m256d sum0 = _mm256_add_pd(ac0, bd0);
54 const __m256d diff0 = _mm256_sub_pd(ac0, bd0);
55 const __m256d diffbar0 = _mm256_shuffle_pd(diff0, diff0, 5);
56 const __m256d om0 = _mm256_load_pd((*omga)[0 + 0]);
57 const __m256d omre0 = _mm256_unpacklo_pd(om0, om0);
58 const __m256d omim0 = _mm256_unpackhi_pd(om0, om0);
59 const __m256d t10 = _mm256_mul_pd(diffbar0, omim0);
60 const __m256d t20 = _mm256_fmaddsub_pd(diff0, omre0, t10);
61 const __m256d newab0 = _mm256_permute2f128_pd(sum0, t20, 0b100000);
62 const __m256d newcd0 = _mm256_permute2f128_pd(sum0, t20, 0b110001);
63 _mm256_storeu_pd(dd[0 + 2 * 0], newab0);
64 _mm256_storeu_pd(dd[1 + 2 * 0], newcd0);
65 dd += 2 * 1;
66 *omga += 2 * 1;
67 // END_INTERLEAVE
68 } while (dd < finaldd);
69#if 0
70 printf("c after first: ");
71 for (uint64_t ii=0; ii<nn/2; ++ii) {
72 printf("%.6lf %.6lf ",ddata[ii][0],ddata[ii][1]);
73 }
74 printf("\n");
75#endif
76 }
77 // general case
78 const uint32_t ms2 = m >> 1;
79 for (uint32_t _2nblock = 2; _2nblock <= ms2; _2nblock <<= 1) {
80 // _2nblock = h in ref code
81 uint32_t nblock = _2nblock >> 1; // =h/2 in ref code
82 D4MEM* dd = (D4MEM*)data;
83 do {
84 const __m256d om = _mm256_load_pd((*omga)[0]);
85 const __m256d omre = _mm256_unpacklo_pd(om, om);
86 const __m256d omim = _mm256_addsub_pd(_mm256_setzero_pd(), _mm256_unpackhi_pd(om, om));
87 D4MEM* const ddend = (dd + nblock);
88 D4MEM* ddmid = ddend;
89 do {
90 const __m256d a = _mm256_loadu_pd(dd[0]);
91 const __m256d b = _mm256_loadu_pd(ddmid[0]);
92 const __m256d newa = _mm256_add_pd(a, b);
93 _mm256_storeu_pd(dd[0], newa);
94 const __m256d diff = _mm256_sub_pd(a, b);
95 const __m256d t1 = _mm256_mul_pd(diff, omre);
96 const __m256d bardiff = _mm256_shuffle_pd(diff, diff, 5);
97 const __m256d t2 = _mm256_fmadd_pd(bardiff, omim, t1);
98 _mm256_storeu_pd(ddmid[0], t2);
99 dd += 1;
100 ddmid += 1;
101 } while (dd < ddend);
102 dd += nblock;
103 *omga += 2;
104 } while (dd < finaldd);
105 }
106}
107
108/**
109 * @brief complex fft via bfs strategy (for m >= 16)
110 * @param dat the data to run the algorithm on
111 * @param omg precomputed tables (must have been filled with fill_omega)
112 * @param m ring dimension of the FFT (modulo X^m-i)
113 */
114void cplx_ifft_avx2_fma_bfs_16(D4MEM* dat, const D2MEM** omga, uint32_t m) {
115 double* data = (double*)dat;
116 D4MEM* const finaldd = (D4MEM*)(data + 2 * m);
117 // base iteration when h = _2nblock == 8
118 {
119 D4MEM* dd = (D4MEM*)data;
120 do {
121 cplx_ifft16_avx_fma(dd, *omga);
122 dd += 8;
123 *omga += 8;
124 } while (dd < finaldd);
125 }
126 // general case
127 const uint32_t log2m = _mm_popcnt_u32(m - 1); //_popcnt32(m-1); //log2(m);
128 uint32_t h = 16;
129 if (log2m % 2 == 1) {
130 uint32_t nblock = h >> 1; // =h/2 in ref code
131 D4MEM* dd = (D4MEM*)data;
132 do {
133 const __m128d om1 = _mm_loadu_pd((*omga)[0]);
134 const __m256d om = _mm256_set_m128d(om1, om1);
135 const __m256d omre = _mm256_unpacklo_pd(om, om);
136 const __m256d omim = _mm256_addsub_pd(_mm256_setzero_pd(), _mm256_unpackhi_pd(om, om));
137 D4MEM* const ddend = (dd + nblock);
138 D4MEM* ddmid = ddend;
139 do {
140 const __m256d a = _mm256_loadu_pd(dd[0]);
141 const __m256d b = _mm256_loadu_pd(ddmid[0]);
142 const __m256d newa = _mm256_add_pd(a, b);
143 _mm256_storeu_pd(dd[0], newa);
144 const __m256d diff = _mm256_sub_pd(a, b);
145 const __m256d t1 = _mm256_mul_pd(diff, omre);
146 const __m256d bardiff = _mm256_shuffle_pd(diff, diff, 5);
147 const __m256d t2 = _mm256_fmadd_pd(bardiff, omim, t1);
148 _mm256_storeu_pd(ddmid[0], t2);
149 dd += 1;
150 ddmid += 1;
151 } while (dd < ddend);
152 dd += nblock;
153 *omga += 1;
154 } while (dd < finaldd);
155 h = 32;
156 }
157 for (; h < m; h <<= 2) {
158 // _2nblock = h in ref code
159 uint32_t nblock = h >> 1; // =h/2 in ref code
160 D4MEM* dd0 = (D4MEM*)data;
161 do {
162 const __m128d om1 = _mm_loadu_pd((*omga)[0]);
163 const __m128d al1 = _mm_loadu_pd((*omga)[1]);
164 const __m256d om = _mm256_set_m128d(om1, om1);
165 const __m256d al = _mm256_set_m128d(al1, al1);
166 const __m256d omre = _mm256_unpacklo_pd(om, om);
167 const __m256d omim = _mm256_unpackhi_pd(om, om);
168 const __m256d alre = _mm256_unpacklo_pd(al, al);
169 const __m256d alim = _mm256_unpackhi_pd(al, al);
170 D4MEM* const ddend = (dd0 + nblock);
171 D4MEM* dd1 = ddend;
172 D4MEM* dd2 = dd1 + nblock;
173 D4MEM* dd3 = dd2 + nblock;
174 do {
175 __m256d u0 = _mm256_loadu_pd(dd0[0]);
176 __m256d u1 = _mm256_loadu_pd(dd1[0]);
177 __m256d u2 = _mm256_loadu_pd(dd2[0]);
178 __m256d u3 = _mm256_loadu_pd(dd3[0]);
179 __m256d u4 = _mm256_add_pd(u0, u1);
180 __m256d u5 = _mm256_sub_pd(u0, u1);
181 __m256d u6 = _mm256_add_pd(u2, u3);
182 __m256d u7 = _mm256_sub_pd(u2, u3);
183 u0 = _mm256_shuffle_pd(u5, u5, 5);
184 u2 = _mm256_shuffle_pd(u7, u7, 5);
185 u1 = _mm256_mul_pd(u0, omim);
186 u3 = _mm256_mul_pd(u2, omre);
187 u5 = _mm256_fmaddsub_pd(u5, omre, u1);
188 u7 = _mm256_fmsubadd_pd(u7, omim, u3);
189 //////
190 u0 = _mm256_add_pd(u4, u6);
191 u1 = _mm256_add_pd(u5, u7);
192 u2 = _mm256_sub_pd(u4, u6);
193 u3 = _mm256_sub_pd(u5, u7);
194 u4 = _mm256_shuffle_pd(u2, u2, 5);
195 u5 = _mm256_shuffle_pd(u3, u3, 5);
196 u6 = _mm256_mul_pd(u4, alim);
197 u7 = _mm256_mul_pd(u5, alim);
198 u2 = _mm256_fmaddsub_pd(u2, alre, u6);
199 u3 = _mm256_fmaddsub_pd(u3, alre, u7);
200 ///////
201 _mm256_storeu_pd(dd0[0], u0);
202 _mm256_storeu_pd(dd1[0], u1);
203 _mm256_storeu_pd(dd2[0], u2);
204 _mm256_storeu_pd(dd3[0], u3);
205 dd0 += 1;
206 dd1 += 1;
207 dd2 += 1;
208 dd3 += 1;
209 } while (dd0 < ddend);
210 dd0 += 3 * nblock;
211 *omga += 2;
212 } while (dd0 < finaldd);
213 }
214}
215
216/**
217 * @brief complex ifft via dfs recursion (for m >= 16)
218 * @param dat the data to run the algorithm on
219 * @param omg precomputed tables (must have been filled with fill_omega)
220 * @param m ring dimension of the FFT (modulo X^m-i)
221 */
222void cplx_ifft_avx2_fma_rec_16(D4MEM* dat, const D2MEM** omga, uint32_t m) {
223 if (m <= 8) return cplx_ifft_avx2_fma_bfs_2(dat, omga, m);
224 if (m <= 2048) return cplx_ifft_avx2_fma_bfs_16(dat, omga, m);
225 const uint32_t _2nblock = m >> 1; // = h in ref code
226 const uint32_t nblock = _2nblock >> 1; // =h/2 in ref code
227 cplx_ifft_avx2_fma_rec_16(dat, omga, _2nblock);
228 cplx_ifft_avx2_fma_rec_16(dat + nblock, omga, _2nblock);
229 {
230 // final iteration
231 D4MEM* dd = dat;
232 const __m256d om = _mm256_load_pd((*omga)[0]);
233 const __m256d omre = _mm256_unpacklo_pd(om, om);
234 const __m256d omim = _mm256_addsub_pd(_mm256_setzero_pd(), _mm256_unpackhi_pd(om, om));
235 D4MEM* const ddend = (dd + nblock);
236 D4MEM* ddmid = ddend;
237 do {
238 const __m256d a = _mm256_loadu_pd(dd[0]);
239 const __m256d b = _mm256_loadu_pd(ddmid[0]);
240 const __m256d newa = _mm256_add_pd(a, b);
241 _mm256_storeu_pd(dd[0], newa);
242 const __m256d diff = _mm256_sub_pd(a, b);
243 const __m256d t1 = _mm256_mul_pd(diff, omre);
244 const __m256d bardiff = _mm256_shuffle_pd(diff, diff, 5);
245 const __m256d t2 = _mm256_fmadd_pd(bardiff, omim, t1);
246 _mm256_storeu_pd(ddmid[0], t2);
247 dd += 1;
248 ddmid += 1;
249 } while (dd < ddend);
250 *omga += 2;
251 }
252}
253
254/**
255 * @brief complex ifft via best strategy (for m>=1)
256 * @param dat the data to run the algorithm on: m complex numbers
257 * @param omg precomputed tables (must have been filled with fill_omega)
258 * @param m ring dimension of the FFT (modulo X^m-i)
259 */
260EXPORT void cplx_ifft_avx2_fma(const CPLX_IFFT_PRECOMP* precomp, void* d) {
261 const uint32_t m = precomp->m;
262 const D2MEM* omg = (D2MEM*)precomp->powomegas;
263 if (m <= 1) return;
264 if (m <= 8) return cplx_ifft_avx2_fma_bfs_2(d, &omg, m);
265 if (m <= 2048) return cplx_ifft_avx2_fma_bfs_16(d, &omg, m);
266 cplx_ifft_avx2_fma_rec_16(d, &omg, m);
267}