add avx512 adler32 implementation

Author: folkertdev

.github/workflows/checks.yaml

@@ -484,6 +484,10 @@ jobs:
         run: "cargo +nightly miri nextest run -j4 -p zlib-rs --target ${{ matrix.target }} --features=vpclmulqdq crc32::"
         env:
           RUSTFLAGS: "-Ctarget-feature=+vpclmulqdq,+avx512f"
+      - name: Test avx512 adler32 implementation
+        run: "cargo +nightly miri nextest run -j4 -p zlib-rs --target ${{ matrix.target }} --features=avx512 adler32::"
+        env:
+          RUSTFLAGS: "-Ctarget-feature=+avx2,+bmi2,+bmi1,+avx512f,+avx512bw"
       - name: Test allocator with miri
         run: "cargo +nightly miri nextest run -j4 -p zlib-rs --target ${{ matrix.target }} allocate::"
       - name: Test gz logic with miri

Cargo.lock

@@ -24,7 +24,8 @@ __internal-fuzz = ["arbitrary"]
 __internal-fuzz-disable-checksum = [] # disable checksum validation on inflate
 __internal-test = ["quickcheck"]
 ZLIB_DEBUG = []
-vpclmulqdq = [] # use avx512 to speed up crc32. Only stable from 1.89.0 onwards
+vpclmulqdq = [] # use avx512 to speed up crc32. Only stable from 1.89.0 onwards.
+avx512 = ["vpclmulqdq"] # use avx512 to speed up crc32 and adler32. Only stable from 1.89.0 onwards.
 
 
 [dependencies]
@@ -33,5 +34,5 @@ quickcheck = { workspace = true, optional = true }
 
 [dev-dependencies]
 crc32fast = "1.3.2"
-memoffset = "0.9.1"
+# memoffset = "0.9.1"
 quickcheck.workspace = true

zlib-rs/Cargo.toml

@@ -1,12 +1,21 @@
 #[cfg(target_arch = "x86_64")]
 mod avx2;
+#[cfg(feature = "avx512")]
+#[cfg(target_arch = "x86_64")]
+mod avx512;
 mod generic;
 #[cfg(target_arch = "aarch64")]
 mod neon;
 #[cfg(any(target_arch = "wasm32", target_arch = "wasm64"))]
 mod wasm;
 
 pub fn adler32(start_checksum: u32, data: &[u8]) -> u32 {
+    #[cfg(feature = "avx512")]
+    #[cfg(target_arch = "x86_64")]
+    if cfg!(all(target_feature = "avx512f", target_feature = "avx512bw")) {
+        return unsafe { avx512::adler32_avx512(start_checksum, data) };
+    }
+
     #[cfg(target_arch = "x86_64")]
     if crate::cpu_features::is_enabled_avx2_and_bmi2() {
         return avx2::adler32_avx2(start_checksum, data);

zlib-rs/src/adler32.rs

@@ -27,7 +27,7 @@ const DOT3V: __m256i = __m256i_literal([
     1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
 ]);
 
-const ZERO: __m256i = __m256i_literal([0; 32]);
+const ZERO: __m256i = __m256i_literal([0u8; 32]);
 
 /// 32 bit horizontal sum, adapted from Agner Fog's vector library.
 #[target_feature(enable = "avx2")]

zlib-rs/src/adler32/avx2.rs

@@ -0,0 +1,142 @@
+use core::arch::x86_64::{
+    __m512i, _mm256_add_epi32, _mm256_castsi256_si128, _mm256_extracti128_si256, _mm512_add_epi32,
+    _mm512_castsi512_si256, _mm512_extracti64x4_epi64, _mm512_madd_epi16, _mm512_maddubs_epi16,
+    _mm512_permutexvar_epi32, _mm512_sad_epu8, _mm512_set1_epi16, _mm512_setr_epi32,
+    _mm512_slli_epi32, _mm512_zextsi128_si512, _mm_add_epi32, _mm_cvtsi128_si32, _mm_cvtsi32_si128,
+    _mm_shuffle_epi32, _mm_unpackhi_epi64,
+};
+
+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([
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26,
+    27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
+    51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64,
+]);
+
+const ZERO: __m512i = __m512i_literal([0u8; 64]);
+
+#[target_feature(enable = "avx512f")]
+#[target_feature(enable = "avx512bw")]
+pub fn adler32_avx512(adler: u32, src: &[u8]) -> u32 {
+    assert!(cfg!(target_feature = "avx512f"));
+    assert!(cfg!(target_feature = "avx512bw"));
+    // SAFETY: the assertion above ensures this code is not executed unless the CPU has AVX2.
+    unsafe { adler32_avx512_help(adler, src) }
+}
+
+#[target_feature(enable = "avx512f")]
+#[target_feature(enable = "avx512bw")]
+unsafe fn adler32_avx512_help(adler: u32, src: &[u8]) -> u32 {
+    if src.is_empty() {
+        return adler;
+    }
+
+    // SAFETY: [u8; 32] safely transmutes into __m256i.
+    let (before, middle, after) = unsafe { src.align_to::<__m512i>() };
+
+    let adler = if !before.is_empty() {
+        super::avx2::adler32_avx2(adler, src)
+    } else {
+        adler
+    };
+
+    let mut adler1 = (adler >> 16) & 0xffff;
+    let mut adler0 = adler & 0xffff;
+
+    // use largest step possible (without causing overflow)
+    for chunk in middle.chunks(NMAX as usize / 64) {
+        (adler0, adler1) = unsafe { helper_64_bytes(adler0, adler1, chunk) };
+    }
+
+    if after.is_empty() {
+        adler0 | (adler1 << 16)
+    } else {
+        super::avx2::adler32_avx2(adler0 | (adler1 << 16), src)
+    }
+}
+
+#[target_feature(enable = "avx2")]
+unsafe fn helper_64_bytes(mut adler0: u32, mut adler1: u32, src: &[__m512i]) -> (u32, u32) {
+    unsafe {
+        let mut vs1 = _mm512_zextsi128_si512(_mm_cvtsi32_si128(adler0 as i32));
+        let mut vs2 = _mm512_zextsi128_si512(_mm_cvtsi32_si128(adler1 as i32));
+
+        let mut vs1_0 = vs1;
+        let mut vs3 = ZERO;
+
+        let dot3v = _mm512_set1_epi16(1);
+
+        for vbuf in src.iter().copied() {
+            let vs1_sad = _mm512_sad_epu8(vbuf, ZERO);
+            let v_short_sum2 = _mm512_maddubs_epi16(vbuf, DOT2V);
+            vs1 = _mm512_add_epi32(vs1_sad, vs1);
+            vs3 = _mm512_add_epi32(vs3, vs1_0);
+            let vsum2 = _mm512_madd_epi16(v_short_sum2, dot3v);
+            vs2 = _mm512_add_epi32(vsum2, vs2);
+            vs1_0 = vs1;
+        }
+
+        /* Defer the multiplication with 32 to outside of the loop */
+        vs3 = _mm512_slli_epi32(vs3, 6);
+        vs2 = _mm512_add_epi32(vs2, vs3);
+
+        adler0 = partial_hsum(vs1) % BASE;
+        adler1 = _mm512_reduce_add_epu32(vs2) % BASE;
+
+        (adler0, adler1)
+    }
+}
+
+#[inline(always)]
+unsafe fn _mm512_reduce_add_epu32(x: __m512i) -> u32 {
+    unsafe {
+        let a = _mm512_extracti64x4_epi64(x, 1);
+        let b = _mm512_extracti64x4_epi64(x, 0);
+
+        let a_plus_b = _mm256_add_epi32(a, b);
+        let c = _mm256_extracti128_si256(a_plus_b, 1);
+        let d = _mm256_extracti128_si256(a_plus_b, 0);
+        let c_plus_d = _mm_add_epi32(c, d);
+
+        let sum1 = _mm_unpackhi_epi64(c_plus_d, c_plus_d);
+        let sum2 = _mm_add_epi32(sum1, c_plus_d);
+        let sum3 = _mm_shuffle_epi32(sum2, 0x01);
+        let sum4 = _mm_add_epi32(sum2, sum3);
+
+        _mm_cvtsi128_si32(sum4) as u32
+    }
+}
+
+#[inline(always)]
+unsafe fn partial_hsum(x: __m512i) -> u32 {
+    unsafe {
+        // We need a permutation vector to extract every other integer. The
+        // rest are going to be zeros. Marking this const so the compiler stands
+        // a better chance of keeping this resident in a register through entire
+        // loop execution. We certainly have enough zmm registers (32) */
+        let perm_vec: __m512i =
+            _mm512_setr_epi32(0, 2, 4, 6, 8, 10, 12, 14, 1, 1, 1, 1, 1, 1, 1, 1);
+
+        let non_zero = _mm512_permutexvar_epi32(perm_vec, x);
+
+        /* From here, it's a simple 256 bit wide reduction sum */
+        let non_zero_avx = _mm512_castsi512_si256(non_zero);
+
+        /* See Agner Fog's vectorclass for a decent reference. Essentially, phadd is
+         * pretty slow, much slower than the longer instruction sequence below */
+        let sum1 = _mm_add_epi32(
+            _mm256_extracti128_si256(non_zero_avx, 1),
+            _mm256_castsi256_si128(non_zero_avx),
+        );
+        let sum2 = _mm_add_epi32(sum1, _mm_unpackhi_epi64(sum1, sum1));
+        let sum3 = _mm_add_epi32(sum2, _mm_shuffle_epi32(sum2, 1));
+
+        _mm_cvtsi128_si32(sum3) as u32
+    }
+}