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