poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/cplx/cplx_fft_ref.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/cplx/cplx_fft_ref.cRaw
1#include <memory.h>
2#include <stdio.h>
3
4#include "../commons_private.h"
5#include "cplx_fft_internal.h"
6#include "cplx_fft_private.h"
7
8/** @brief (a,b) <- (a+omega.b,a-omega.b) */
9void ctwiddle(CPLX a, CPLX b, const CPLX om) {
10 double re = om[0] * b[0] - om[1] * b[1];
11 double im = om[0] * b[1] + om[1] * b[0];
12 b[0] = a[0] - re;
13 b[1] = a[1] - im;
14 a[0] += re;
15 a[1] += im;
16}
17
18/** @brief (a,b) <- (a+i.omega.b,a-i.omega.b) */
19void citwiddle(CPLX a, CPLX b, const CPLX om) {
20 double re = -om[1] * b[0] - om[0] * b[1];
21 double im = -om[1] * b[1] + om[0] * b[0];
22 b[0] = a[0] - re;
23 b[1] = a[1] - im;
24 a[0] += re;
25 a[1] += im;
26}
27
28/**
29 * @brief FFT modulo X^16-omega^2 (in registers)
30 * @param data contains 16 complexes
31 * @param omega 8 complexes in this order:
32 * omega,alpha,beta,j.beta,gamma,j.gamma,k.gamma,kj.gamma
33 * alpha = sqrt(omega), beta = sqrt(alpha), gamma = sqrt(beta)
34 * j = sqrt(i), k=sqrt(j)
35 */
36void cplx_fft16_ref(void* data, const void* omega) {
37 CPLX* d = data;
38 const CPLX* om = omega;
39 // first pass
40 for (uint64_t i = 0; i < 8; ++i) {
41 ctwiddle(d[0 + i], d[8 + i], om[0]);
42 }
43 //
44 ctwiddle(d[0], d[4], om[1]);
45 ctwiddle(d[1], d[5], om[1]);
46 ctwiddle(d[2], d[6], om[1]);
47 ctwiddle(d[3], d[7], om[1]);
48 citwiddle(d[8], d[12], om[1]);
49 citwiddle(d[9], d[13], om[1]);
50 citwiddle(d[10], d[14], om[1]);
51 citwiddle(d[11], d[15], om[1]);
52 //
53 ctwiddle(d[0], d[2], om[2]);
54 ctwiddle(d[1], d[3], om[2]);
55 citwiddle(d[4], d[6], om[2]);
56 citwiddle(d[5], d[7], om[2]);
57 ctwiddle(d[8], d[10], om[3]);
58 ctwiddle(d[9], d[11], om[3]);
59 citwiddle(d[12], d[14], om[3]);
60 citwiddle(d[13], d[15], om[3]);
61 //
62 ctwiddle(d[0], d[1], om[4]);
63 citwiddle(d[2], d[3], om[4]);
64 ctwiddle(d[4], d[5], om[5]);
65 citwiddle(d[6], d[7], om[5]);
66 ctwiddle(d[8], d[9], om[6]);
67 citwiddle(d[10], d[11], om[6]);
68 ctwiddle(d[12], d[13], om[7]);
69 citwiddle(d[14], d[15], om[7]);
70}
71
72double cos_2pix(double x) { return m_accurate_cos(2 * M_PI * x); }
73double sin_2pix(double x) { return m_accurate_sin(2 * M_PI * x); }
74void cplx_set_e2pix(CPLX res, double x) {
75 res[0] = cos_2pix(x);
76 res[1] = sin_2pix(x);
77}
78
79void cplx_fft16_precomp(const double entry_pwr, CPLX** omg) {
80 static const double j_pow = 1. / 8.;
81 static const double k_pow = 1. / 16.;
82 const double pom = entry_pwr / 2.;
83 const double pom_2 = entry_pwr / 4.;
84 const double pom_4 = entry_pwr / 8.;
85 const double pom_8 = entry_pwr / 16.;
86 cplx_set_e2pix((*omg)[0], pom);
87 cplx_set_e2pix((*omg)[1], pom_2);
88 cplx_set_e2pix((*omg)[2], pom_4);
89 cplx_set_e2pix((*omg)[3], pom_4 + j_pow);
90 cplx_set_e2pix((*omg)[4], pom_8);
91 cplx_set_e2pix((*omg)[5], pom_8 + j_pow);
92 cplx_set_e2pix((*omg)[6], pom_8 + k_pow);
93 cplx_set_e2pix((*omg)[7], pom_8 + j_pow + k_pow);
94 *omg += 8;
95}
96
97/**
98 * @brief h twiddles-fft on the same omega
99 * (also called merge-fft)merges 2 times h evaluations of even/odd polynomials into 2h evaluations of a sigle polynomial
100 * Input: P_0(z),...,P_{h-1}(z),P_h(z),...,P_{2h-1}(z)
101 * Output: Q_0(y),...,Q_{h-1}(y),Q_0(-y),...,Q_{h-1}(-y)
102 * where Q_i(X)=P_i(X^2)+X.P_{h+i}(X^2) and y^2 = z
103 * @param h number of "coefficients" h >= 1
104 * @param data 2h complex coefficients interleaved and 256b aligned
105 * @param powom y represented as (yre,yim)
106 */
107void cplx_twiddle_fft_ref(int32_t h, CPLX* data, const CPLX powom) {
108 CPLX* d0 = data;
109 CPLX* d1 = data + h;
110 for (uint64_t i = 0; i < h; ++i) {
111 ctwiddle(d0[i], d1[i], powom);
112 }
113}
114
115void cplx_bitwiddle_fft_ref(int32_t h, CPLX* data, const CPLX powom[2]) {
116 CPLX* d0 = data;
117 CPLX* d1 = data + h;
118 CPLX* d2 = data + 2 * h;
119 CPLX* d3 = data + 3 * h;
120 for (uint64_t i = 0; i < h; ++i) {
121 ctwiddle(d0[i], d2[i], powom[0]);
122 ctwiddle(d1[i], d3[i], powom[0]);
123 }
124 for (uint64_t i = 0; i < h; ++i) {
125 ctwiddle(d0[i], d1[i], powom[1]);
126 citwiddle(d2[i], d3[i], powom[1]);
127 }
128}
129
130/**
131 * Input: P_0(z),P_1(z)
132 * Output: Q(y),Q(-y)
133 * where Q(X)=P_0(X^2)+X.P_1(X^2) and y^2 = z
134 * @param data 2 complexes coefficients interleaved and 256b aligned
135 * @param powom (z,-z) interleaved: (zre,zim,-zre,-zim)
136 */
137void merge_fft_last_ref(CPLX* data, const CPLX powom) {
138 CPLX prod;
139 cplx_mul(prod, data[1], powom);
140 cplx_sub(data[1], data[0], prod);
141 cplx_add(data[0], data[0], prod);
142}
143
144void cplx_fft_ref_bfs_2(CPLX* dat, const CPLX** omg, uint32_t m) {
145 CPLX* data = (CPLX*)dat;
146 CPLX* const dend = data + m;
147 for (int32_t h = m / 2; h >= 2; h >>= 1) {
148 for (CPLX* d = data; d < dend; d += 2 * h) {
149 if (memcmp((*omg)[0], (*omg)[1], 8) != 0) abort();
150 cplx_twiddle_fft_ref(h, d, **omg);
151 *omg += 2;
152 }
153#if 0
154 printf("after merge %d: ", h);
155 for (uint64_t ii=0; ii<nn/2; ++ii) {
156 printf("%.6lf %.6lf ",data[ii][0],data[ii][1]);
157 }
158 printf("\n");
159#endif
160 }
161 for (CPLX* d = data; d < dend; d += 2) {
162 // TODO see if encoding changes
163 if ((*omg)[0][0] != -(*omg)[1][0]) abort();
164 if ((*omg)[0][1] != -(*omg)[1][1]) abort();
165 merge_fft_last_ref(d, **omg);
166 *omg += 2;
167 }
168#if 0
169 printf("after last: ");
170 for (uint64_t ii=0; ii<nn/2; ++ii) {
171 printf("%.6lf %.6lf ",data[ii][0],data[ii][1]);
172 }
173 printf("\n");
174#endif
175}
176
177void cplx_fft_ref_bfs_16(CPLX* dat, const CPLX** omg, uint32_t m) {
178 CPLX* data = (CPLX*)dat;
179 CPLX* const dend = data + m;
180 uint32_t mm = m;
181 uint32_t log2m = log2(m);
182 if (log2m % 2 == 1) {
183 cplx_twiddle_fft_ref(mm / 2, data, **omg);
184 *omg += 2;
185 mm >>= 1;
186 }
187 while (mm > 16) {
188 uint32_t h = mm / 4;
189 for (CPLX* d = data; d < dend; d += mm) {
190 cplx_bitwiddle_fft_ref(h, d, *omg);
191 *omg += 2;
192 }
193 mm = h;
194 }
195 for (CPLX* d = data; d < dend; d += 16) {
196 cplx_fft16_ref(d, *omg);
197 *omg += 8;
198 }
199#if 0
200 printf("after last: ");
201 for (uint64_t ii=0; ii<nn/2; ++ii) {
202 printf("%.6lf %.6lf ",data[ii][0],data[ii][1]);
203 }
204 printf("\n");
205#endif
206}
207
208/** @brief fft modulo X^m-exp(i.2pi.entry+pwr) -- reference code */
209void cplx_fft_naive(const uint32_t m, const double entry_pwr, CPLX* data) {
210 if (m == 1) return;
211 const double pom = entry_pwr / 2.;
212 const uint32_t h = m / 2;
213 // apply the twiddle factors
214 CPLX cpom;
215 cplx_set_e2pix(cpom, pom);
216 for (uint64_t i = 0; i < h; ++i) {
217 ctwiddle(data[i], data[i + h], cpom);
218 }
219 // do the recursive calls
220 cplx_fft_naive(h, pom, data);
221 cplx_fft_naive(h, pom + 0.5, data + h);
222}
223
224/** @brief fills omega for cplx_fft_bfs_16 modulo X^m-exp(i.2.pi.entry_pwr) */
225void fill_cplx_fft_omegas_bfs_16(const double entry_pwr, CPLX** omg, uint32_t m) {
226 uint32_t mm = m;
227 uint32_t log2m = log2(m);
228 double ss = entry_pwr;
229 if (log2m % 2 == 1) {
230 uint32_t h = mm / 2;
231 double pom = ss / 2.;
232 for (uint32_t i = 0; i < m / mm; i++) {
233 cplx_set_e2pix(omg[0][0], pom + fracrevbits(i) / 2.);
234 cplx_set(omg[0][1], omg[0][0]);
235 *omg += 2;
236 }
237 mm = h;
238 ss = pom;
239 }
240 while (mm > 16) {
241 double pom = ss / 4.;
242 uint32_t h = mm / 4;
243 for (uint32_t i = 0; i < m / mm; i++) {
244 double om = pom + fracrevbits(i) / 4.;
245 cplx_set_e2pix(omg[0][0], 2. * om);
246 cplx_set_e2pix(omg[0][1], om);
247 *omg += 2;
248 }
249 mm = h;
250 ss = pom;
251 }
252 {
253 // mm=16
254 for (uint32_t i = 0; i < m / 16; i++) {
255 cplx_fft16_precomp(ss + fracrevbits(i), omg);
256 }
257 }
258}
259
260/** @brief fills omega for cplx_fft_bfs_2 modulo X^m-exp(i.2.pi.entry_pwr) */
261void fill_cplx_fft_omegas_bfs_2(const double entry_pwr, CPLX** omg, uint32_t m) {
262 double pom = entry_pwr / 2.;
263 for (int32_t h = m / 2; h >= 2; h >>= 1) {
264 for (uint32_t i = 0; i < m / (2 * h); i++) {
265 cplx_set_e2pix(omg[0][0], pom + fracrevbits(i) / 2.);
266 cplx_set(omg[0][1], omg[0][0]);
267 *omg += 2;
268 }
269 pom /= 2;
270 }
271 {
272 // h=1
273 for (uint32_t i = 0; i < m / 2; i++) {
274 cplx_set_e2pix((*omg)[0], pom + fracrevbits(i) / 2.);
275 cplx_neg((*omg)[1], (*omg)[0]);
276 *omg += 2;
277 }
278 }
279}
280
281/** @brief fills omega for cplx_fft_rec modulo X^m-exp(i.2.pi.entry_pwr) */
282void fill_cplx_fft_omegas_rec_16(const double entry_pwr, CPLX** omg, uint32_t m) {
283 // note that the cases below are for recursive calls only!
284 // externally, this function shall only be called with m>=4096
285 if (m == 1) return;
286 if (m <= 8) return fill_cplx_fft_omegas_bfs_2(entry_pwr, omg, m);
287 if (m <= 2048) return fill_cplx_fft_omegas_bfs_16(entry_pwr, omg, m);
288 double pom = entry_pwr / 2.;
289 cplx_set_e2pix((*omg)[0], pom);
290 cplx_set_e2pix((*omg)[1], pom);
291 *omg += 2;
292 fill_cplx_fft_omegas_rec_16(pom, omg, m / 2);
293 fill_cplx_fft_omegas_rec_16(pom + 0.5, omg, m / 2);
294}
295
296void cplx_fft_ref_rec_16(CPLX* dat, const CPLX** omg, uint32_t m) {
297 if (m == 1) return;
298 if (m <= 8) return cplx_fft_ref_bfs_2(dat, omg, m);
299 if (m <= 2048) return cplx_fft_ref_bfs_16(dat, omg, m);
300 const uint32_t h = m / 2;
301 if (memcmp((*omg)[0], (*omg)[1], 8) != 0) abort();
302 cplx_twiddle_fft_ref(h, dat, **omg);
303 *omg += 2;
304 cplx_fft_ref_rec_16(dat, omg, h);
305 cplx_fft_ref_rec_16(dat + h, omg, h);
306}
307
308void cplx_fft_ref(const CPLX_FFT_PRECOMP* precomp, void* d) {
309 CPLX* data = (CPLX*)d;
310 const int32_t m = precomp->m;
311 const CPLX* omg = (CPLX*)precomp->powomegas;
312 if (m == 1) return;
313 if (m <= 8) return cplx_fft_ref_bfs_2(data, &omg, m);
314 if (m <= 2048) return cplx_fft_ref_bfs_16(data, &omg, m);
315 cplx_fft_ref_rec_16(data, &omg, m);
316}
317
318EXPORT CPLX_FFT_PRECOMP* new_cplx_fft_precomp(uint32_t m, uint32_t num_buffers) {
319 const uint64_t OMG_SPACE = ceilto64b((2 * m) * sizeof(CPLX));
320 const uint64_t BUF_SIZE = ceilto64b(m * sizeof(CPLX));
321 void* reps = malloc(sizeof(CPLX_FFT_PRECOMP) + 63 // padding
322 + OMG_SPACE // tables //TODO 16?
323 + num_buffers * BUF_SIZE // buffers
324 );
325 uint64_t aligned_addr = ceilto64b((uint64_t)(reps) + sizeof(CPLX_FFT_PRECOMP));
326 CPLX_FFT_PRECOMP* r = (CPLX_FFT_PRECOMP*)reps;
327 r->m = m;
328 r->buf_size = BUF_SIZE;
329 r->powomegas = (double*)aligned_addr;
330 r->aligned_buffers = (void*)(aligned_addr + OMG_SPACE);
331 // fill in powomegas
332 CPLX* omg = (CPLX*)r->powomegas;
333 if (m <= 8) {
334 fill_cplx_fft_omegas_bfs_2(0.25, &omg, m);
335 } else if (m <= 2048) {
336 fill_cplx_fft_omegas_bfs_16(0.25, &omg, m);
337 } else {
338 fill_cplx_fft_omegas_rec_16(0.25, &omg, m);
339 }
340 if (((uint64_t)omg) - aligned_addr > OMG_SPACE) abort();
341 // dispatch the right implementation
342 {
343 if (m <= 4) {
344 // currently, we do not have any acceletated
345 // implementation for m<=4
346 r->function = cplx_fft_ref;
347 } else if (CPU_SUPPORTS("fma")) {
348 r->function = cplx_fft_avx2_fma;
349 } else {
350 r->function = cplx_fft_ref;
351 }
352 }
353 return reps;
354}
355
356EXPORT void* cplx_fft_precomp_get_buffer(const CPLX_FFT_PRECOMP* tables, uint32_t buffer_index) {
357 return (uint8_t*)tables->aligned_buffers + buffer_index * tables->buf_size;
358}
359
360EXPORT void cplx_fft_simple(uint32_t m, void* data) {
361 static CPLX_FFT_PRECOMP* p[31] = {0};
362 CPLX_FFT_PRECOMP** f = p + log2m(m);
363 if (!*f) *f = new_cplx_fft_precomp(m, 0);
364 (*f)->function(*f, data);
365}
366
367EXPORT void cplx_fft(const CPLX_FFT_PRECOMP* tables, void* data) { tables->function(tables, data); }