arnaucube
arnaucube/poulpyPublic Self-hosted repository
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}