Skip to main content

round4

Function round4 

Source
unsafe fn round4(value: &[f64; 4]) -> [i64; 4]
Examples found in repository?
crates/competitive/src/math/mint_fft_convolve.rs (line 233)
199pub unsafe fn convolve_u64_avx2(a: Vec<u64>, b: Vec<u64>) -> Vec<u64> {
200    let len = a.len() + b.len() - 1;
201    let n = len.next_power_of_two() / 2;
202    let a_parts = if a.iter().all(|&value| value <= u32::MAX as u64) {
203        3
204    } else {
205        5
206    };
207    let mut fa = split_u64_coefficients(&a, n);
208    drop(a);
209    let b_parts = if b.iter().all(|&value| value <= u32::MAX as u64) {
210        3
211    } else {
212        5
213    };
214    let mut fb = split_u64_coefficients(&b, n);
215    drop(b);
216    for part in 0..3 {
217        fft_soa(&mut fa[part]);
218        fft_soa(&mut fb[part]);
219    }
220    for part in &mut fa[3..a_parts] {
221        fft_soa(part);
222    }
223    for part in &mut fb[3..b_parts] {
224        fft_soa(part);
225    }
226    dot_u64_soa(&mut fa, &fb);
227    drop(fb);
228    for part in &mut fa {
229        ifft_soa(part);
230    }
231    let mut result = vec![0; len];
232    for (block, _) in fa[0].iter().enumerate() {
233        let real: [[i64; 4]; 5] = std::array::from_fn(|part| round4(&fa[part][block].re));
234        let imag: [[i64; 4]; 5] = std::array::from_fn(|part| round4(&fa[part][block].im));
235        for lane in 0..4 {
236            let i = block * 4 + lane;
237            if i < len {
238                result[i] = (real[0][lane] as u64)
239                    .wrapping_add((real[1][lane] as u64) << 13)
240                    .wrapping_add((real[2][lane] as u64) << 26)
241                    .wrapping_add((real[3][lane] as u64) << 39)
242                    .wrapping_add((real[4][lane] as u64) << 52);
243            }
244            if i + n < len {
245                result[i + n] = (imag[0][lane] as u64)
246                    .wrapping_add((imag[1][lane] as u64) << 13)
247                    .wrapping_add((imag[2][lane] as u64) << 26)
248                    .wrapping_add((imag[3][lane] as u64) << 39)
249                    .wrapping_add((imag[4][lane] as u64) << 52);
250            }
251        }
252    }
253    result
254}