poulpy/poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_fft_ref.c
poulpy-backend/src/cpu_spqlios/spqlios-arithmetic/spqlios/reim/reim_fft_ref.cRaw
1#include "../commons_private.h"
2#include "reim_fft_internal.h"
3#include "reim_fft_private.h"
4
5EXPORT void reim_fft_simple(uint32_t m, void* data) {
6 static REIM_FFT_PRECOMP* p[31] = {0};
7 REIM_FFT_PRECOMP** f = p + log2m(m);
8 if (!*f) *f = new_reim_fft_precomp(m, 0);
9 (*f)->function(*f, data);
10}
11
12EXPORT void reim_ifft_simple(uint32_t m, void* data) {
13 static REIM_IFFT_PRECOMP* p[31] = {0};
14 REIM_IFFT_PRECOMP** f = p + log2m(m);
15 if (!*f) *f = new_reim_ifft_precomp(m, 0);
16 (*f)->function(*f, data);
17}
18
19EXPORT void reim_fftvec_add_simple(uint32_t m, void* r, const void* a, const void* b) {
20 static REIM_FFTVEC_ADD_PRECOMP* p[31] = {0};
21 REIM_FFTVEC_ADD_PRECOMP** f = p + log2m(m);
22 if (!*f) *f = new_reim_fftvec_add_precomp(m);
23 (*f)->function(*f, r, a, b);
24}
25
26EXPORT void reim_fftvec_sub_simple(uint32_t m, void* r, const void* a, const void* b) {
27 static REIM_FFTVEC_SUB_PRECOMP* p[31] = {0};
28 REIM_FFTVEC_SUB_PRECOMP** f = p + log2m(m);
29 if (!*f) *f = new_reim_fftvec_sub_precomp(m);
30 (*f)->function(*f, r, a, b);
31}
32
33EXPORT void reim_fftvec_mul_simple(uint32_t m, void* r, const void* a, const void* b) {
34 static REIM_FFTVEC_MUL_PRECOMP* p[31] = {0};
35 REIM_FFTVEC_MUL_PRECOMP** f = p + log2m(m);
36 if (!*f) *f = new_reim_fftvec_mul_precomp(m);
37 (*f)->function(*f, r, a, b);
38}
39
40EXPORT void reim_fftvec_addmul_simple(uint32_t m, void* r, const void* a, const void* b) {
41 static REIM_FFTVEC_ADDMUL_PRECOMP* p[31] = {0};
42 REIM_FFTVEC_ADDMUL_PRECOMP** f = p + log2m(m);
43 if (!*f) *f = new_reim_fftvec_addmul_precomp(m);
44 (*f)->function(*f, r, a, b);
45}
46
47void reim_ctwiddle(double* ra, double* ia, double* rb, double* ib, double omre, double omim) {
48 double newrt = *rb * omre - *ib * omim;
49 double newit = *rb * omim + *ib * omre;
50 *rb = *ra - newrt;
51 *ib = *ia - newit;
52 *ra = *ra + newrt;
53 *ia = *ia + newit;
54}
55
56// i (omre + i omim) = -omim + i omre
57void reim_citwiddle(double* ra, double* ia, double* rb, double* ib, double omre, double omim) {
58 double newrt = -*rb * omim - *ib * omre;
59 double newit = *rb * omre - *ib * omim;
60 *rb = *ra - newrt;
61 *ib = *ia - newit;
62 *ra = *ra + newrt;
63 *ia = *ia + newit;
64}
65
66void reim_fft16_ref(double* dre, double* dim, const void* pom) {
67 const double* om = (const double*)pom;
68 {
69 double omre = om[0];
70 double omim = om[1];
71 reim_ctwiddle(&dre[0], &dim[0], &dre[8], &dim[8], omre, omim);
72 reim_ctwiddle(&dre[1], &dim[1], &dre[9], &dim[9], omre, omim);
73 reim_ctwiddle(&dre[2], &dim[2], &dre[10], &dim[10], omre, omim);
74 reim_ctwiddle(&dre[3], &dim[3], &dre[11], &dim[11], omre, omim);
75 reim_ctwiddle(&dre[4], &dim[4], &dre[12], &dim[12], omre, omim);
76 reim_ctwiddle(&dre[5], &dim[5], &dre[13], &dim[13], omre, omim);
77 reim_ctwiddle(&dre[6], &dim[6], &dre[14], &dim[14], omre, omim);
78 reim_ctwiddle(&dre[7], &dim[7], &dre[15], &dim[15], omre, omim);
79 }
80 {
81 double omre = om[2];
82 double omim = om[3];
83 reim_ctwiddle(&dre[0], &dim[0], &dre[4], &dim[4], omre, omim);
84 reim_ctwiddle(&dre[1], &dim[1], &dre[5], &dim[5], omre, omim);
85 reim_ctwiddle(&dre[2], &dim[2], &dre[6], &dim[6], omre, omim);
86 reim_ctwiddle(&dre[3], &dim[3], &dre[7], &dim[7], omre, omim);
87 reim_citwiddle(&dre[8], &dim[8], &dre[12], &dim[12], omre, omim);
88 reim_citwiddle(&dre[9], &dim[9], &dre[13], &dim[13], omre, omim);
89 reim_citwiddle(&dre[10], &dim[10], &dre[14], &dim[14], omre, omim);
90 reim_citwiddle(&dre[11], &dim[11], &dre[15], &dim[15], omre, omim);
91 }
92 {
93 double omare = om[4];
94 double omaim = om[5];
95 double ombre = om[6];
96 double ombim = om[7];
97 reim_ctwiddle(&dre[0], &dim[0], &dre[2], &dim[2], omare, omaim);
98 reim_ctwiddle(&dre[1], &dim[1], &dre[3], &dim[3], omare, omaim);
99 reim_citwiddle(&dre[4], &dim[4], &dre[6], &dim[6], omare, omaim);
100 reim_citwiddle(&dre[5], &dim[5], &dre[7], &dim[7], omare, omaim);
101 reim_ctwiddle(&dre[8], &dim[8], &dre[10], &dim[10], ombre, ombim);
102 reim_ctwiddle(&dre[9], &dim[9], &dre[11], &dim[11], ombre, ombim);
103 reim_citwiddle(&dre[12], &dim[12], &dre[14], &dim[14], ombre, ombim);
104 reim_citwiddle(&dre[13], &dim[13], &dre[15], &dim[15], ombre, ombim);
105 }
106 {
107 double omare = om[8];
108 double ombre = om[9];
109 double omcre = om[10];
110 double omdre = om[11];
111 double omaim = om[12];
112 double ombim = om[13];
113 double omcim = om[14];
114 double omdim = om[15];
115 reim_ctwiddle(&dre[0], &dim[0], &dre[1], &dim[1], omare, omaim);
116 reim_citwiddle(&dre[2], &dim[2], &dre[3], &dim[3], omare, omaim);
117 reim_ctwiddle(&dre[4], &dim[4], &dre[5], &dim[5], ombre, ombim);
118 reim_citwiddle(&dre[6], &dim[6], &dre[7], &dim[7], ombre, ombim);
119 reim_ctwiddle(&dre[8], &dim[8], &dre[9], &dim[9], omcre, omcim);
120 reim_citwiddle(&dre[10], &dim[10], &dre[11], &dim[11], omcre, omcim);
121 reim_ctwiddle(&dre[12], &dim[12], &dre[13], &dim[13], omdre, omdim);
122 reim_citwiddle(&dre[14], &dim[14], &dre[15], &dim[15], omdre, omdim);
123 }
124}
125
126void fill_reim_fft16_omegas(const double entry_pwr, double** omg) {
127 const double j_pow = 1. / 8.;
128 const double k_pow = 1. / 16.;
129 const double pin = entry_pwr / 2.;
130 const double pin_2 = entry_pwr / 4.;
131 const double pin_4 = entry_pwr / 8.;
132 const double pin_8 = entry_pwr / 16.;
133 // 0 and 1 are real and imag of om
134 (*omg)[0] = cos(2. * M_PI * pin);
135 (*omg)[1] = sin(2. * M_PI * pin);
136 // 2 and 3 are real and imag of om^1/2
137 (*omg)[2] = cos(2. * M_PI * (pin_2));
138 (*omg)[3] = sin(2. * M_PI * (pin_2));
139 // (4,5) and (6,7) are real and imag of om^1/4 and j.om^1/4
140 (*omg)[4] = cos(2. * M_PI * (pin_4));
141 (*omg)[5] = sin(2. * M_PI * (pin_4));
142 (*omg)[6] = cos(2. * M_PI * (pin_4 + j_pow));
143 (*omg)[7] = sin(2. * M_PI * (pin_4 + j_pow));
144 // ((8,9,10,11),(12,13,14,15)) are 4 reals then 4 imag of om^1/8*(1,k,j,kj)
145 (*omg)[8] = cos(2. * M_PI * (pin_8));
146 (*omg)[9] = cos(2. * M_PI * (pin_8 + j_pow));
147 (*omg)[10] = cos(2. * M_PI * (pin_8 + k_pow));
148 (*omg)[11] = cos(2. * M_PI * (pin_8 + j_pow + k_pow));
149 (*omg)[12] = sin(2. * M_PI * (pin_8));
150 (*omg)[13] = sin(2. * M_PI * (pin_8 + j_pow));
151 (*omg)[14] = sin(2. * M_PI * (pin_8 + k_pow));
152 (*omg)[15] = sin(2. * M_PI * (pin_8 + j_pow + k_pow));
153 *omg += 16;
154}
155
156void reim_fft8_ref(double* dre, double* dim, const void* pom) {
157 const double* om = (const double*)pom;
158 {
159 double omre = om[0];
160 double omim = om[1];
161 reim_ctwiddle(&dre[0], &dim[0], &dre[4], &dim[4], omre, omim);
162 reim_ctwiddle(&dre[1], &dim[1], &dre[5], &dim[5], omre, omim);
163 reim_ctwiddle(&dre[2], &dim[2], &dre[6], &dim[6], omre, omim);
164 reim_ctwiddle(&dre[3], &dim[3], &dre[7], &dim[7], omre, omim);
165 }
166 {
167 double omare = om[2];
168 double omaim = om[3];
169 reim_ctwiddle(&dre[0], &dim[0], &dre[2], &dim[2], omare, omaim);
170 reim_ctwiddle(&dre[1], &dim[1], &dre[3], &dim[3], omare, omaim);
171 reim_citwiddle(&dre[4], &dim[4], &dre[6], &dim[6], omare, omaim);
172 reim_citwiddle(&dre[5], &dim[5], &dre[7], &dim[7], omare, omaim);
173 }
174 {
175 double omare = om[4];
176 double ombre = om[5];
177 double omaim = om[6];
178 double ombim = om[7];
179 reim_ctwiddle(&dre[0], &dim[0], &dre[1], &dim[1], omare, omaim);
180 reim_citwiddle(&dre[2], &dim[2], &dre[3], &dim[3], omare, omaim);
181 reim_ctwiddle(&dre[4], &dim[4], &dre[5], &dim[5], ombre, ombim);
182 reim_citwiddle(&dre[6], &dim[6], &dre[7], &dim[7], ombre, ombim);
183 }
184}
185
186void fill_reim_fft8_omegas(const double entry_pwr, double** omg) {
187 const double j_pow = 1. / 8.;
188 const double pin = entry_pwr / 2.;
189 const double pin_2 = entry_pwr / 4.;
190 const double pin_4 = entry_pwr / 8.;
191 // 0 and 1 are real and imag of om
192 (*omg)[0] = cos(2. * M_PI * pin);
193 (*omg)[1] = sin(2. * M_PI * pin);
194 // 2 and 3 are real and imag of om^1/2
195 (*omg)[2] = cos(2. * M_PI * (pin_2));
196 (*omg)[3] = sin(2. * M_PI * (pin_2));
197 // (4,5) and (6,7) are real and imag of om^1/4 and j.om^1/4
198 (*omg)[4] = cos(2. * M_PI * (pin_4));
199 (*omg)[5] = cos(2. * M_PI * (pin_4 + j_pow));
200 (*omg)[6] = sin(2. * M_PI * (pin_4));
201 (*omg)[7] = sin(2. * M_PI * (pin_4 + j_pow));
202 *omg += 8;
203}
204
205void reim_fft4_ref(double* dre, double* dim, const void* pom) {
206 const double* om = (const double*)pom;
207 {
208 double omare = om[0];
209 double omaim = om[1];
210 reim_ctwiddle(&dre[0], &dim[0], &dre[2], &dim[2], omare, omaim);
211 reim_ctwiddle(&dre[1], &dim[1], &dre[3], &dim[3], omare, omaim);
212 }
213 {
214 double omare = om[2];
215 double omaim = om[3];
216 reim_ctwiddle(&dre[0], &dim[0], &dre[1], &dim[1], omare, omaim);
217 reim_citwiddle(&dre[2], &dim[2], &dre[3], &dim[3], omare, omaim);
218 }
219}
220
221void fill_reim_fft4_omegas(const double entry_pwr, double** omg) {
222 const double pin = entry_pwr / 2.;
223 const double pin_2 = entry_pwr / 4.;
224 // 0 and 1 are real and imag of om
225 (*omg)[0] = cos(2. * M_PI * pin);
226 (*omg)[1] = sin(2. * M_PI * pin);
227 // 2 and 3 are real and imag of om^1/2
228 (*omg)[2] = cos(2. * M_PI * (pin_2));
229 (*omg)[3] = sin(2. * M_PI * (pin_2));
230 *omg += 4;
231}
232
233void reim_fft2_ref(double* dre, double* dim, const void* pom) {
234 const double* om = (const double*)pom;
235 {
236 double omare = om[0];
237 double omaim = om[1];
238 reim_ctwiddle(&dre[0], &dim[0], &dre[1], &dim[1], omare, omaim);
239 }
240}
241
242void fill_reim_fft2_omegas(const double entry_pwr, double** omg) {
243 const double pin = entry_pwr / 2.;
244 // 0 and 1 are real and imag of om
245 (*omg)[0] = cos(2. * M_PI * pin);
246 (*omg)[1] = sin(2. * M_PI * pin);
247 *omg += 2;
248}
249
250void reim_twiddle_fft_ref(uint64_t h, double* re, double* im, double om[2]) {
251 for (uint64_t i = 0; i < h; ++i) {
252 reim_ctwiddle(&re[i], &im[i], &re[h + i], &im[h + i], om[0], om[1]);
253 }
254}
255
256void reim_bitwiddle_fft_ref(uint64_t h, double* re, double* im, double om[4]) {
257 double* r0 = re;
258 double* r1 = re + h;
259 double* r2 = re + 2 * h;
260 double* r3 = re + 3 * h;
261 double* i0 = im;
262 double* i1 = im + h;
263 double* i2 = im + 2 * h;
264 double* i3 = im + 3 * h;
265 for (uint64_t i = 0; i < h; ++i) {
266 reim_ctwiddle(&r0[i], &i0[i], &r2[i], &i2[i], om[0], om[1]);
267 reim_ctwiddle(&r1[i], &i1[i], &r3[i], &i3[i], om[0], om[1]);
268 }
269 for (uint64_t i = 0; i < h; ++i) {
270 reim_ctwiddle(&r0[i], &i0[i], &r1[i], &i1[i], om[2], om[3]);
271 reim_citwiddle(&r2[i], &i2[i], &r3[i], &i3[i], om[2], om[3]);
272 }
273}
274
275void reim_fft_bfs_16_ref(uint64_t m, double* re, double* im, double** omg) {
276 uint64_t log2m = log2(m);
277 uint64_t mm = m;
278 if (log2m % 2 != 0) {
279 uint64_t h = mm >> 1;
280 // do the first twiddle iteration normally
281 reim_twiddle_fft_ref(h, re, im, *omg);
282 *omg += 2;
283 mm = h;
284 }
285 while (mm > 16) {
286 uint64_t h = mm >> 2;
287 for (uint64_t off = 0; off < m; off += mm) {
288 reim_bitwiddle_fft_ref(h, re + off, im + off, *omg);
289 *omg += 4;
290 }
291 mm = h;
292 }
293 if (mm != 16) abort(); // bug!
294 for (uint64_t off = 0; off < m; off += 16) {
295 reim_fft16_ref(re + off, im + off, *omg);
296 *omg += 16;
297 }
298}
299
300void fill_reim_fft_bfs_16_omegas(uint64_t m, double entry_pwr, double** omg) {
301 uint64_t log2m = log2(m);
302 uint64_t mm = m;
303 double ss = entry_pwr;
304 if (log2m % 2 != 0) {
305 uint64_t h = mm >> 1;
306 double s = ss / 2.;
307 // do the first twiddle iteration normally
308 (*omg)[0] = cos(2 * M_PI * s);
309 (*omg)[1] = sin(2 * M_PI * s);
310 *omg += 2;
311 mm = h;
312 ss = s;
313 }
314 while (mm > 16) {
315 uint64_t h = mm >> 2;
316 double s = ss / 4.;
317 for (uint64_t off = 0; off < m; off += mm) {
318 double rs0 = s + fracrevbits(off / mm) / 4.;
319 double rs1 = 2. * rs0;
320 (*omg)[0] = cos(2 * M_PI * rs1);
321 (*omg)[1] = sin(2 * M_PI * rs1);
322 (*omg)[2] = cos(2 * M_PI * rs0);
323 (*omg)[3] = sin(2 * M_PI * rs0);
324 *omg += 4;
325 }
326 mm = h;
327 ss = s;
328 }
329 if (mm != 16) abort(); // bug!
330 for (uint64_t off = 0; off < m; off += 16) {
331 double s = ss + fracrevbits(off / 16);
332 fill_reim_fft16_omegas(s, omg);
333 }
334}
335
336void reim_fft_rec_16_ref(uint64_t m, double* re, double* im, double** omg) {
337 if (m <= 2048) return reim_fft_bfs_16_ref(m, re, im, omg);
338 const uint32_t h = m / 2;
339 reim_twiddle_fft_ref(h, re, im, *omg);
340 *omg += 2;
341 reim_fft_rec_16_ref(h, re, im, omg);
342 reim_fft_rec_16_ref(h, re + h, im + h, omg);
343}
344
345void fill_reim_fft_rec_16_omegas(uint64_t m, double entry_pwr, double** omg) {
346 if (m <= 2048) return fill_reim_fft_bfs_16_omegas(m, entry_pwr, omg);
347 const uint64_t h = m / 2;
348 const double s = entry_pwr / 2;
349 (*omg)[0] = cos(2 * M_PI * s);
350 (*omg)[1] = sin(2 * M_PI * s);
351 *omg += 2;
352 fill_reim_fft_rec_16_omegas(h, s, omg);
353 fill_reim_fft_rec_16_omegas(h, s + 0.5, omg);
354}
355
356void reim_fft_ref(const REIM_FFT_PRECOMP* precomp, double* dat) {
357 const int32_t m = precomp->m;
358 double* omg = precomp->powomegas;
359 double* re = dat;
360 double* im = dat + m;
361 if (m <= 16) {
362 switch (m) {
363 case 1:
364 return;
365 case 2:
366 return reim_fft2_ref(re, im, omg);
367 case 4:
368 return reim_fft4_ref(re, im, omg);
369 case 8:
370 return reim_fft8_ref(re, im, omg);
371 case 16:
372 return reim_fft16_ref(re, im, omg);
373 default:
374 abort(); // m is not a power of 2
375 }
376 }
377 if (m <= 2048) return reim_fft_bfs_16_ref(m, re, im, &omg);
378 return reim_fft_rec_16_ref(m, re, im, &omg);
379}
380
381EXPORT REIM_FFT_PRECOMP* new_reim_fft_precomp(uint32_t m, uint32_t num_buffers) {
382 const uint64_t OMG_SPACE = ceilto64b(2 * m * sizeof(double));
383 const uint64_t BUF_SIZE = ceilto64b(2 * m * sizeof(double));
384 void* reps = malloc(sizeof(REIM_FFT_PRECOMP) // base
385 + 63 // padding
386 + OMG_SPACE // tables //TODO 16?
387 + num_buffers * BUF_SIZE // buffers
388 );
389 uint64_t aligned_addr = ceilto64b((uint64_t)(reps) + sizeof(REIM_FFT_PRECOMP));
390 REIM_FFT_PRECOMP* r = (REIM_FFT_PRECOMP*)reps;
391 r->m = m;
392 r->buf_size = BUF_SIZE;
393 r->powomegas = (double*)aligned_addr;
394 r->aligned_buffers = (void*)(aligned_addr + OMG_SPACE);
395 // fill in powomegas
396 double* omg = (double*)r->powomegas;
397 if (m <= 16) {
398 switch (m) {
399 case 1:
400 break;
401 case 2:
402 fill_reim_fft2_omegas(0.25, &omg);
403 break;
404 case 4:
405 fill_reim_fft4_omegas(0.25, &omg);
406 break;
407 case 8:
408 fill_reim_fft8_omegas(0.25, &omg);
409 break;
410 case 16:
411 fill_reim_fft16_omegas(0.25, &omg);
412 break;
413 default:
414 abort(); // m is not a power of 2
415 }
416 } else if (m <= 2048) {
417 fill_reim_fft_bfs_16_omegas(m, 0.25, &omg);
418 } else {
419 fill_reim_fft_rec_16_omegas(m, 0.25, &omg);
420 }
421 if (((uint64_t)omg) - aligned_addr > OMG_SPACE) abort();
422 // dispatch the right implementation
423 {
424 if (CPU_SUPPORTS("fma")) {
425 r->function = reim_fft_avx2_fma;
426 } else {
427 r->function = reim_fft_ref;
428 }
429 }
430 return reps;
431}
432
433void reim_naive_fft(uint64_t m, double entry_pwr, double* re, double* im) {
434 if (m == 1) return;
435 // twiddle
436 const uint64_t h = m / 2;
437 const double s = entry_pwr / 2.;
438 const double sre = cos(2 * M_PI * s);
439 const double sim = sin(2 * M_PI * s);
440 for (uint64_t j = 0; j < h; ++j) {
441 double pre = re[h + j] * sre - im[h + j] * sim;
442 double pim = im[h + j] * sre + re[h + j] * sim;
443 re[h + j] = re[j] - pre;
444 im[h + j] = im[j] - pim;
445 re[j] += pre;
446 im[j] += pim;
447 }
448 reim_naive_fft(h, s, re, im);
449 reim_naive_fft(h, s + 0.5, re + h, im + h);
450}