Function batch_ntt_simd_backend 

fn batch_ntt_simd_backend(len: usize, width: usize) -> SimdBackend

Examples found in repository

crates/competitive/src/math/number_theoretic_transform.rs (line 246)

fn ntt_batch<M>(a: &mut [MInt<M>], width: usize)
where
    M: Montgomery32NttModulus,
{
    #[cfg(target_arch = "x86_64")]
    {
        match batch_ntt_simd_backend(a.len(), width) {
            SimdBackend::Avx512 => {
                // SAFETY: backend detection checked all required AVX-512 features.
                unsafe { ntt_simd::ntt_batch_avx512(a, width) };
                return;
            }
            SimdBackend::Avx2 => {
                // SAFETY: backend detection checked AVX2.
                unsafe { ntt_simd::ntt_batch_avx2::<_, false>(a, width) };
                return;
            }
            SimdBackend::Scalar => {}
        }
    }
    ntt_batch_scalar(a, width);
}

fn ntt_batch_scalar<M>(a: &mut [MInt<M>], width: usize)
where
    M: Montgomery32NttModulus,
{
    let n = a.len() / width;
    if n <= 1 {
        return;
    }
    let mut v = n / 2;
    if n.trailing_zeros() & 1 == 1 {
        let (l, r) = a.split_at_mut(v * width);
        for (x0, x1) in l.iter_mut().zip(r) {
            let a0 = *x0;
            let a1 = *x1;
            *x0 = a0 + a1;
            *x1 = a0 - a1;
        }
        v >>= 1;
    }
    let imag = MInt::<M>::new_unchecked(M::INFO.root[2]);
    while v > 1 {
        let mut w1 = MInt::<M>::one();
        let mut w2 = w1;
        let mut w3 = w1;
        for (s, a) in a.chunks_exact_mut((v << 1) * width).enumerate() {
            let (l, r) = a.split_at_mut(v * width);
            let (ll, lr) = l.split_at_mut((v >> 1) * width);
            let (rl, rr) = r.split_at_mut((v >> 1) * width);
            for (((x0, x1), x2), x3) in ll.iter_mut().zip(lr).zip(rl).zip(rr) {
                let a0 = *x0;
                let a1 = *x1 * w1;
                let a2 = *x2 * w2;
                let a3 = *x3 * w3;
                let a0pa2 = a0 + a2;
                let a0na2 = a0 - a2;
                let a1pa3 = a1 + a3;
                let a1na3imag = (a1 - a3) * imag;
                *x0 = a0pa2 + a1pa3;
                *x1 = a0pa2 - a1pa3;
                *x2 = a0na2 + a1na3imag;
                *x3 = a0na2 - a1na3imag;
            }
            let rate = &M::INFO.rate3_packed[s.trailing_ones() as usize];
            w1 *= MInt::<M>::new_unchecked(rate[1]);
            w2 *= MInt::<M>::new_unchecked(rate[3]);
            w3 *= MInt::<M>::new_unchecked(rate[5]);
        }
        v >>= 2;
    }
}

fn intt_scalar<M>(a: &mut [MInt<M>])
where
    M: Montgomery32NttModulus,
{
    intt_batch_scalar(a, 1);
}

fn intt_batch<M>(a: &mut [MInt<M>], width: usize)
where
    M: Montgomery32NttModulus,
{
    #[cfg(target_arch = "x86_64")]
    {
        match batch_ntt_simd_backend(a.len(), width) {
            SimdBackend::Avx512 => {
                // SAFETY: backend detection checked all required AVX-512 features.
                unsafe { ntt_simd::intt_batch_avx512::<_, false>(a, width) };
                return;
            }
            SimdBackend::Avx2 => {
                // SAFETY: backend detection checked AVX2.
                unsafe { ntt_simd::intt_batch_avx2::<_, false>(a, width) };
                return;
            }
            SimdBackend::Scalar => {}
        }
    }
    intt_batch_scalar(a, width);
}