poulpy/poulpy-hal/src/reference/fft64/reim/fft_ref.rs
poulpy-hal/src/reference/fft64/reim/fft_ref.rsRaw
1use std::fmt::Debug;
2
3use rand_distr::num_traits::{Float, FloatConst};
4
5use crate::reference::fft64::reim::{as_arr, as_arr_mut};
6
7#[inline(always)]
8pub fn fft_ref<R: Float + FloatConst + Debug>(m: usize, omg: &[R], data: &mut [R]) {
9 assert!(data.len() == 2 * m);
10 let (re, im) = data.split_at_mut(m);
11
12 if m <= 16 {
13 match m {
14 1 => {}
15 2 => fft2_ref(
16 as_arr_mut::<2, R>(re),
17 as_arr_mut::<2, R>(im),
18 *as_arr::<2, R>(omg),
19 ),
20 4 => fft4_ref(
21 as_arr_mut::<4, R>(re),
22 as_arr_mut::<4, R>(im),
23 *as_arr::<4, R>(omg),
24 ),
25 8 => fft8_ref(
26 as_arr_mut::<8, R>(re),
27 as_arr_mut::<8, R>(im),
28 *as_arr::<8, R>(omg),
29 ),
30 16 => fft16_ref(
31 as_arr_mut::<16, R>(re),
32 as_arr_mut::<16, R>(im),
33 *as_arr::<16, R>(omg),
34 ),
35 _ => {}
36 }
37 } else if m <= 2048 {
38 fft_bfs_16_ref(m, re, im, omg, 0);
39 } else {
40 fft_rec_16_ref(m, re, im, omg, 0);
41 }
42}
43
44#[inline(always)]
45fn fft_rec_16_ref<R: Float + FloatConst + Debug>(m: usize, re: &mut [R], im: &mut [R], omg: &[R], mut pos: usize) -> usize {
46 if m <= 2048 {
47 return fft_bfs_16_ref(m, re, im, omg, pos);
48 };
49
50 let h = m >> 1;
51 twiddle_fft_ref(h, re, im, as_arr::<2, R>(&omg[pos..]));
52 pos += 2;
53 pos = fft_rec_16_ref(h, re, im, omg, pos);
54 pos = fft_rec_16_ref(h, &mut re[h..], &mut im[h..], omg, pos);
55 pos
56}
57
58#[inline(always)]
59fn cplx_twiddle<R: Float + FloatConst>(ra: &mut R, ia: &mut R, rb: &mut R, ib: &mut R, omg_re: R, omg_im: R) {
60 let dr: R = *rb * omg_re - *ib * omg_im;
61 let di: R = *rb * omg_im + *ib * omg_re;
62 *rb = *ra - dr;
63 *ib = *ia - di;
64 *ra = *ra + dr;
65 *ia = *ia + di;
66}
67
68#[inline(always)]
69fn cplx_i_twiddle<R: Float + FloatConst>(ra: &mut R, ia: &mut R, rb: &mut R, ib: &mut R, omg_re: R, omg_im: R) {
70 let dr: R = *rb * omg_im + *ib * omg_re;
71 let di: R = *rb * omg_re - *ib * omg_im;
72 *rb = *ra + dr;
73 *ib = *ia - di;
74 *ra = *ra - dr;
75 *ia = *ia + di;
76}
77
78#[inline(always)]
79fn fft2_ref<R: Float + FloatConst>(re: &mut [R; 2], im: &mut [R; 2], omg: [R; 2]) {
80 let [ra, rb] = re;
81 let [ia, ib] = im;
82 let [romg, iomg] = omg;
83 cplx_twiddle(ra, ia, rb, ib, romg, iomg);
84}
85
86#[inline(always)]
87fn fft4_ref<R: Float + FloatConst>(re: &mut [R; 4], im: &mut [R; 4], omg: [R; 4]) {
88 let [re_0, re_1, re_2, re_3] = re;
89 let [im_0, im_1, im_2, im_3] = im;
90
91 {
92 let omg_0 = omg[0];
93 let omg_1 = omg[1];
94 cplx_twiddle(re_0, im_0, re_2, im_2, omg_0, omg_1);
95 cplx_twiddle(re_1, im_1, re_3, im_3, omg_0, omg_1);
96 }
97
98 {
99 let omg_0 = omg[2];
100 let omg_1 = omg[3];
101 cplx_twiddle(re_0, im_0, re_1, im_1, omg_0, omg_1);
102 cplx_i_twiddle(re_2, im_2, re_3, im_3, omg_0, omg_1);
103 }
104}
105
106#[inline(always)]
107fn fft8_ref<R: Float + FloatConst>(re: &mut [R; 8], im: &mut [R; 8], omg: [R; 8]) {
108 let [re_0, re_1, re_2, re_3, re_4, re_5, re_6, re_7] = re;
109 let [im_0, im_1, im_2, im_3, im_4, im_5, im_6, im_7] = im;
110
111 {
112 let omg_0 = omg[0];
113 let omg_1 = omg[1];
114 cplx_twiddle(re_0, im_0, re_4, im_4, omg_0, omg_1);
115 cplx_twiddle(re_1, im_1, re_5, im_5, omg_0, omg_1);
116 cplx_twiddle(re_2, im_2, re_6, im_6, omg_0, omg_1);
117 cplx_twiddle(re_3, im_3, re_7, im_7, omg_0, omg_1);
118 }
119
120 {
121 let omg_2 = omg[2];
122 let omg_3 = omg[3];
123 cplx_twiddle(re_0, im_0, re_2, im_2, omg_2, omg_3);
124 cplx_twiddle(re_1, im_1, re_3, im_3, omg_2, omg_3);
125 cplx_i_twiddle(re_4, im_4, re_6, im_6, omg_2, omg_3);
126 cplx_i_twiddle(re_5, im_5, re_7, im_7, omg_2, omg_3);
127 }
128
129 {
130 let omg_4 = omg[4];
131 let omg_5 = omg[5];
132 let omg_6 = omg[6];
133 let omg_7 = omg[7];
134 cplx_twiddle(re_0, im_0, re_1, im_1, omg_4, omg_6);
135 cplx_i_twiddle(re_2, im_2, re_3, im_3, omg_4, omg_6);
136 cplx_twiddle(re_4, im_4, re_5, im_5, omg_5, omg_7);
137 cplx_i_twiddle(re_6, im_6, re_7, im_7, omg_5, omg_7);
138 }
139}
140
141#[inline(always)]
142fn fft16_ref<R: Float + FloatConst + Debug>(re: &mut [R; 16], im: &mut [R; 16], omg: [R; 16]) {
143 let [
144 re_0,
145 re_1,
146 re_2,
147 re_3,
148 re_4,
149 re_5,
150 re_6,
151 re_7,
152 re_8,
153 re_9,
154 re_10,
155 re_11,
156 re_12,
157 re_13,
158 re_14,
159 re_15,
160 ] = re;
161 let [
162 im_0,
163 im_1,
164 im_2,
165 im_3,
166 im_4,
167 im_5,
168 im_6,
169 im_7,
170 im_8,
171 im_9,
172 im_10,
173 im_11,
174 im_12,
175 im_13,
176 im_14,
177 im_15,
178 ] = im;
179
180 {
181 let omg_0: R = omg[0];
182 let omg_1: R = omg[1];
183 cplx_twiddle(re_0, im_0, re_8, im_8, omg_0, omg_1);
184 cplx_twiddle(re_1, im_1, re_9, im_9, omg_0, omg_1);
185 cplx_twiddle(re_2, im_2, re_10, im_10, omg_0, omg_1);
186 cplx_twiddle(re_3, im_3, re_11, im_11, omg_0, omg_1);
187
188 cplx_twiddle(re_4, im_4, re_12, im_12, omg_0, omg_1);
189 cplx_twiddle(re_5, im_5, re_13, im_13, omg_0, omg_1);
190 cplx_twiddle(re_6, im_6, re_14, im_14, omg_0, omg_1);
191 cplx_twiddle(re_7, im_7, re_15, im_15, omg_0, omg_1);
192 }
193
194 {
195 let omg_2: R = omg[2];
196 let omg_3: R = omg[3];
197 cplx_twiddle(re_0, im_0, re_4, im_4, omg_2, omg_3);
198 cplx_twiddle(re_1, im_1, re_5, im_5, omg_2, omg_3);
199 cplx_twiddle(re_2, im_2, re_6, im_6, omg_2, omg_3);
200 cplx_twiddle(re_3, im_3, re_7, im_7, omg_2, omg_3);
201
202 cplx_i_twiddle(re_8, im_8, re_12, im_12, omg_2, omg_3);
203 cplx_i_twiddle(re_9, im_9, re_13, im_13, omg_2, omg_3);
204 cplx_i_twiddle(re_10, im_10, re_14, im_14, omg_2, omg_3);
205 cplx_i_twiddle(re_11, im_11, re_15, im_15, omg_2, omg_3);
206 }
207
208 {
209 let omg_0: R = omg[4];
210 let omg_1: R = omg[5];
211 let omg_2: R = omg[6];
212 let omg_3: R = omg[7];
213 cplx_twiddle(re_0, im_0, re_2, im_2, omg_0, omg_1);
214 cplx_twiddle(re_1, im_1, re_3, im_3, omg_0, omg_1);
215 cplx_twiddle(re_8, im_8, re_10, im_10, omg_2, omg_3);
216 cplx_twiddle(re_9, im_9, re_11, im_11, omg_2, omg_3);
217
218 cplx_i_twiddle(re_4, im_4, re_6, im_6, omg_0, omg_1);
219 cplx_i_twiddle(re_5, im_5, re_7, im_7, omg_0, omg_1);
220 cplx_i_twiddle(re_12, im_12, re_14, im_14, omg_2, omg_3);
221 cplx_i_twiddle(re_13, im_13, re_15, im_15, omg_2, omg_3);
222 }
223
224 {
225 let omg_0: R = omg[8];
226 let omg_1: R = omg[9];
227 let omg_2: R = omg[10];
228 let omg_3: R = omg[11];
229 let omg_4: R = omg[12];
230 let omg_5: R = omg[13];
231 let omg_6: R = omg[14];
232 let omg_7: R = omg[15];
233 cplx_twiddle(re_0, im_0, re_1, im_1, omg_0, omg_4);
234 cplx_twiddle(re_4, im_4, re_5, im_5, omg_1, omg_5);
235 cplx_twiddle(re_8, im_8, re_9, im_9, omg_2, omg_6);
236 cplx_twiddle(re_12, im_12, re_13, im_13, omg_3, omg_7);
237
238 cplx_i_twiddle(re_2, im_2, re_3, im_3, omg_0, omg_4);
239 cplx_i_twiddle(re_6, im_6, re_7, im_7, omg_1, omg_5);
240 cplx_i_twiddle(re_10, im_10, re_11, im_11, omg_2, omg_6);
241 cplx_i_twiddle(re_14, im_14, re_15, im_15, omg_3, omg_7);
242 }
243}
244
245#[inline(always)]
246fn fft_bfs_16_ref<R: Float + FloatConst + Debug>(m: usize, re: &mut [R], im: &mut [R], omg: &[R], mut pos: usize) -> usize {
247 let log_m: usize = (usize::BITS - (m - 1).leading_zeros()) as usize;
248 let mut mm: usize = m;
249
250 if !log_m.is_multiple_of(2) {
251 let h: usize = mm >> 1;
252 twiddle_fft_ref(h, re, im, as_arr::<2, R>(&omg[pos..]));
253 pos += 2;
254 mm = h
255 }
256
257 while mm > 16 {
258 let h: usize = mm >> 2;
259 for off in (0..m).step_by(mm) {
260 bitwiddle_fft_ref(
261 h,
262 &mut re[off..],
263 &mut im[off..],
264 as_arr::<4, R>(&omg[pos..]),
265 );
266 pos += 4;
267 }
268 mm = h
269 }
270
271 for off in (0..m).step_by(16) {
272 fft16_ref(
273 as_arr_mut::<16, R>(&mut re[off..]),
274 as_arr_mut::<16, R>(&mut im[off..]),
275 *as_arr::<16, R>(&omg[pos..]),
276 );
277 pos += 16;
278 }
279
280 pos
281}
282
283#[inline(always)]
284fn twiddle_fft_ref<R: Float + FloatConst>(h: usize, re: &mut [R], im: &mut [R], omg: &[R; 2]) {
285 let romg = omg[0];
286 let iomg = omg[1];
287
288 let (re_lhs, re_rhs) = re.split_at_mut(h);
289 let (im_lhs, im_rhs) = im.split_at_mut(h);
290
291 for i in 0..h {
292 cplx_twiddle(
293 &mut re_lhs[i],
294 &mut im_lhs[i],
295 &mut re_rhs[i],
296 &mut im_rhs[i],
297 romg,
298 iomg,
299 );
300 }
301}
302
303#[inline(always)]
304fn bitwiddle_fft_ref<R: Float + FloatConst>(h: usize, re: &mut [R], im: &mut [R], omg: &[R; 4]) {
305 let (r0, r2) = re.split_at_mut(2 * h);
306 let (r0, r1) = r0.split_at_mut(h);
307 let (r2, r3) = r2.split_at_mut(h);
308
309 let (i0, i2) = im.split_at_mut(2 * h);
310 let (i0, i1) = i0.split_at_mut(h);
311 let (i2, i3) = i2.split_at_mut(h);
312
313 let omg_0: R = omg[0];
314 let omg_1: R = omg[1];
315 let omg_2: R = omg[2];
316 let omg_3: R = omg[3];
317
318 for i in 0..h {
319 cplx_twiddle(&mut r0[i], &mut i0[i], &mut r2[i], &mut i2[i], omg_0, omg_1);
320 cplx_twiddle(&mut r1[i], &mut i1[i], &mut r3[i], &mut i3[i], omg_0, omg_1);
321 }
322
323 for i in 0..h {
324 cplx_twiddle(&mut r0[i], &mut i0[i], &mut r1[i], &mut i1[i], omg_2, omg_3);
325 cplx_i_twiddle(&mut r2[i], &mut i2[i], &mut r3[i], &mut i3[i], omg_2, omg_3);
326 }
327}