zlib-rs 0.6.3

use core::arch::x86_64::{
    __m512i, _mm512_add_epi32, _mm512_dpbusd_epi32, _mm512_loadu_si512, _mm512_sad_epu8,
    _mm512_setzero_si512, _mm512_slli_epi32, _mm512_zextsi128_si512, _mm_cvtsi32_si128,
};

use super::avx512::{_mm512_reduce_add_epu32, partial_hsum};
use crate::adler32::{BASE, NMAX};

const fn __m512i_literal(bytes: [u8; 64]) -> __m512i {
    // SAFETY: any valid [u8; 64] represents a valid __m512i
    unsafe { core::mem::transmute(bytes) }
}

const DOT2V: __m512i = __m512i_literal({
    let mut arr = [0; 64];

    // generates [64, 63, ..., 2, 1]
    let mut i = 64;
    while i > 0 {
        i -= 1;
        arr[i] = (64 - i) as u8;
    }

    arr
});

pub fn adler32_avx512(adler: u32, src: &[u8]) -> u32 {
    assert!(cfg!(target_feature = "avx512bw"));
    assert!(cfg!(target_feature = "avx512vnni"));
    // SAFETY: the assertion above ensures this code is not executed unless the CPU has avx512.
    unsafe { adler32_avx512_vnni(adler, src) }
}

#[target_feature(enable = "avx512bw")]
#[target_feature(enable = "avx512vnni")]
pub(super) fn adler32_avx512_vnni(mut adler: u32, mut src: &[u8]) -> u32 {
    if src.is_empty() {
        return adler;
    }

    let mut adler0;
    let mut adler1;
    adler1 = (adler >> 16) & 0xffff;
    adler0 = adler & 0xffff;

    unsafe {
        'rem_peel: loop {
            if src.len() < 32 {
                return super::avx2::adler32_avx2(adler, src);
            }

            if src.len() < 64 {
                return super::avx2::adler32_avx2(adler, src);
            }

            let dot2v = DOT2V;

            let zero = _mm512_setzero_si512();
            let mut vs1;
            let mut vs2;

            while src.len() >= 64 {
                vs1 = _mm512_zextsi128_si512(_mm_cvtsi32_si128(adler0 as i32));
                vs2 = _mm512_zextsi128_si512(_mm_cvtsi32_si128(adler1 as i32));
                let mut k: usize = Ord::min(src.len(), NMAX as usize);
                k -= k % 64;
                let mut vs1_0 = vs1;
                let mut vs3 = _mm512_setzero_si512();
                /* We might get a tad bit more ILP here if we sum to a second register in the loop */
                let mut vs2_1 = _mm512_setzero_si512();
                let mut vbuf0;
                let mut vbuf1;

                /* Remainder peeling */
                if (k % 128) != 0 {
                    vbuf1 = _mm512_loadu_si512(src.as_ptr().cast::<__m512i>());

                    src = &src[64..];
                    k -= 64;

                    let vs1_sad = _mm512_sad_epu8(vbuf1, zero);
                    vs1 = _mm512_add_epi32(vs1, vs1_sad);
                    vs3 = _mm512_add_epi32(vs3, vs1_0);
                    vs2 = _mm512_dpbusd_epi32(vs2, vbuf1, dot2v);
                    vs1_0 = vs1;
                }

                /* Manually unrolled this loop by 2 for an decent amount of ILP */
                while k >= 128 {
                    /*
                       vs1 = adler + sum(c[i])
                       vs2 = sum2 + 64 vs1 + sum( (64-i+1) c[i] )
                    */
                    vbuf0 = _mm512_loadu_si512(src.as_ptr().cast::<__m512i>());
                    vbuf1 = _mm512_loadu_si512(src.as_ptr().cast::<__m512i>().add(1));
                    src = &src[128..];
                    k -= 128;

                    let mut vs1_sad = _mm512_sad_epu8(vbuf0, zero);
                    vs1 = _mm512_add_epi32(vs1, vs1_sad);
                    vs3 = _mm512_add_epi32(vs3, vs1_0);
                    /* multiply-add, resulting in 16 ints. Fuse with sum stage from prior versions, as we now have the dp
                     * instructions to eliminate them */
                    vs2 = _mm512_dpbusd_epi32(vs2, vbuf0, dot2v);

                    vs3 = _mm512_add_epi32(vs3, vs1);
                    vs1_sad = _mm512_sad_epu8(vbuf1, zero);
                    vs1 = _mm512_add_epi32(vs1, vs1_sad);
                    vs2_1 = _mm512_dpbusd_epi32(vs2_1, vbuf1, dot2v);
                    vs1_0 = vs1;
                }

                vs3 = _mm512_slli_epi32(vs3, 6);
                vs2 = _mm512_add_epi32(vs2, vs3);
                vs2 = _mm512_add_epi32(vs2, vs2_1);

                adler0 = partial_hsum(vs1) % BASE;
                adler1 = _mm512_reduce_add_epu32(vs2) % BASE;
            }

            adler = adler0 | (adler1 << 16);

            /* Process tail (len < 64). */
            if !src.is_empty() {
                continue 'rem_peel;
            }

            return adler;
        }
    }
}

#[cfg(test)]
#[cfg(target_feature = "avx512vnni")]
#[cfg(target_feature = "avx512bw")]
mod test {
    use super::*;
    use core::arch::x86_64::__m256i;

    #[test]
    fn empty_input() {
        let avx512 = unsafe { adler32_avx512(0, &[]) };
        let rust = crate::adler32::generic::adler32_rust(0, &[]);

        assert_eq!(rust, avx512);
    }

    #[test]
    fn zero_chunks() {
        let input = &[
            1u8, 39, 76, 148, 0, 58, 0, 14, 255, 59, 1, 229, 1, 83, 5, 84, 207, 152, 188,
        ];
        let avx512 = unsafe { adler32_avx512(0, input) };
        let rust = crate::adler32::generic::adler32_rust(0, input);

        assert_eq!(rust, avx512);
    }

    #[test]
    fn one_chunk() {
        let input: [u8; 85] = core::array::from_fn(|i| i as u8);
        let avx512 = unsafe { adler32_avx512(0, &input) };
        let rust = crate::adler32::generic::adler32_rust(0, &input);

        assert_eq!(rust, avx512);
    }

    #[test]
    fn foobar() {
        let v = vec![
            0u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 27, 255, 209, 220, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0,
        ];
        let start = 0;
        let avx512 = unsafe { adler32_avx512(start, &v) };
        let rust = crate::adler32::generic::adler32_rust(start, &v);

        assert_eq!(rust, avx512);
    }

    quickcheck::quickcheck! {
        fn adler32_avx512_is_adler32_rust(v: Vec<u8>, start: u32) -> bool {
            let avx512 = unsafe { adler32_avx512(start, &v) };
            let rust = crate::adler32::generic::adler32_rust(start, &v);

            rust == avx512
        }
    }

    const INPUT: [u8; 128] = {
        let mut array = [0; 128];
        let mut i = 0;
        while i < array.len() {
            array[i] = i as u8;
            i += 1;
        }

        array
    };

    #[test]
    fn start_alignment() {
        // SIMD algorithm is sensitive to alignment;
        for i in 0..16 {
            for start in [crate::ADLER32_INITIAL_VALUE as u32, 42] {
                let avx512 = unsafe { adler32_avx512(start, &INPUT[i..]) };
                let rust = crate::adler32::generic::adler32_rust(start, &INPUT[i..]);

                assert_eq!(avx512, rust, "offset = {i}, start = {start}");
            }
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn large_input() {
        const DEFAULT: &[u8] = include_bytes!("../deflate/test-data/paper-100k.pdf");

        let avx512 = unsafe { adler32_avx512(42, DEFAULT) };
        let rust = crate::adler32::generic::adler32_rust(42, DEFAULT);

        assert_eq!(avx512, rust);
    }
}