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