port avx512 implementation of compare256

Author: folkertdev

.github/workflows/checks.yaml

@@ -488,6 +488,34 @@ jobs:
         run: "cargo +nightly miri nextest run -j4 -p test-libz-rs-sys --target ${{ matrix.target }} null::"
         env:
           RUSTFLAGS: "-Ctarget-feature=+avx2,+bmi2,+bmi1"
+      - 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
+        working-directory: libz-rs-sys-cdylib
+        run: "cargo +nightly miri nextest run -j4 -p libz-rs-sys-cdylib --target ${{ matrix.target }} --features=gz"
+        env:
+          MIRIFLAGS: "-Zmiri-tree-borrows -Zmiri-disable-isolation"
+
+  miri-avx512:
+    name: "Miri avx512"
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        include:
+          - target: x86_64-unknown-linux-gnu
+    env:
+      QUICKCHECK_TESTS: 10
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Miri
+        run: |
+          rustup target add ${{ matrix.target }}
+          rustup toolchain install nightly --component miri
+          cargo +nightly miri setup
+      - name: Install cargo-nextest
+        uses: taiki-e/install-action@d12e869b89167df346dd0ff65da342d1fb1202fb # v2.53.2
+        with:
+          tool: cargo-nextest
       - name: Test avx512 crc32 implementation
         run: "cargo +nightly miri nextest run -j4 -p zlib-rs --target ${{ matrix.target }} --features=vpclmulqdq crc32::"
         env:
@@ -496,13 +524,10 @@ jobs:
         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
-        working-directory: libz-rs-sys-cdylib
-        run: "cargo +nightly miri nextest run -j4 -p libz-rs-sys-cdylib --target ${{ matrix.target }} --features=gz"
+      - name: Test avx512 compare256 implementation
+        run: "cargo +nightly miri nextest run -j4 -p zlib-rs --target ${{ matrix.target }} --features=avx512 compare256"
         env:
-          MIRIFLAGS: "-Zmiri-tree-borrows -Zmiri-disable-isolation"
+          RUSTFLAGS: "-Ctarget-feature=+avx512vl,+avx512bw"
 
   run-flate2-test-suite:
     name: run flate2 test suite

zlib-rs/src/deflate/compare256.rs

@@ -9,6 +9,12 @@ pub fn compare256_slice(src0: &[u8], src1: &[u8]) -> usize {
 }
 
 fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
+    #[cfg(feature = "avx512")]
+    #[cfg(target_arch = "x86_64")]
+    if cfg!(target_feature = "avx512vl") && cfg!(target_feature = "avx512bw") {
+        return unsafe { avx512::compare256(src0, src1) };
+    }
+
     #[cfg(target_arch = "x86_64")]
     if crate::cpu_features::is_enabled_avx2_and_bmi2() {
         return unsafe { avx2::compare256(src0, src1) };
@@ -263,6 +269,93 @@ mod avx2 {
     }
 }
 
+#[cfg(feature = "avx512")]
+#[cfg(target_arch = "x86_64")]
+mod avx512 {
+    use core::arch::x86_64::{
+        _mm512_cmpeq_epu8_mask, _mm512_loadu_si512, _mm_cmpeq_epu8_mask, _mm_loadu_si128,
+    };
+
+    /// # Safety
+    ///
+    /// Behavior is undefined if the `avx` target feature is not enabled
+    #[target_feature(enable = "avx512vl")]
+    #[target_feature(enable = "avx512bw")]
+    pub unsafe fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
+        // First do a 16byte round before increasing to 64bytes, this reduces the
+        // penalty for the short matches, and those are usually the most common ones.
+        // This requires us to overlap on the last round, giving a small penalty
+        // on matches of 192+ bytes (Still faster than AVX2 though).
+
+        unsafe {
+            // 16 bytes
+            let xmm_src0_0 = _mm_loadu_si128(src0.as_ptr().cast());
+            let xmm_src1_0 = _mm_loadu_si128(src1.as_ptr().cast());
+            let mask_0 = u32::from(_mm_cmpeq_epu8_mask(xmm_src0_0, xmm_src1_0)); // zero-extended to use __builtin_ctz
+            if mask_0 != 0x0000FFFF {
+                // There is potential for using __builtin_ctzg/__builtin_ctzs/_tzcnt_u16/__tzcnt_u16 here
+                let match_byte = mask_0.trailing_ones();
+                return match_byte as usize;
+            }
+
+            // 64 bytes
+            let zmm_src0_1 = _mm512_loadu_si512(src0[16..].as_ptr().cast());
+            let zmm_src1_1 = _mm512_loadu_si512(src1[16..].as_ptr().cast());
+            let mask_1 = _mm512_cmpeq_epu8_mask(zmm_src0_1, zmm_src1_1);
+            if mask_1 != 0xFFFFFFFFFFFFFFFF {
+                let match_byte = mask_1.trailing_ones();
+                return 16 + match_byte as usize;
+            }
+
+            // 64 bytes
+            let zmm_src0_2 = _mm512_loadu_si512(src0[80..].as_ptr().cast());
+            let zmm_src1_2 = _mm512_loadu_si512(src1[80..].as_ptr().cast());
+            let mask_2 = _mm512_cmpeq_epu8_mask(zmm_src0_2, zmm_src1_2);
+            if mask_2 != 0xFFFFFFFFFFFFFFFF {
+                let match_byte = mask_2.trailing_ones();
+                return 80 + match_byte as usize;
+            }
+
+            // 64 bytes
+            let zmm_src0_3 = _mm512_loadu_si512(src0[144..].as_ptr().cast());
+            let zmm_src1_3 = _mm512_loadu_si512(src1[144..].as_ptr().cast());
+            let mask_3 = _mm512_cmpeq_epu8_mask(zmm_src0_3, zmm_src1_3);
+            if mask_3 != 0xFFFFFFFFFFFFFFFF {
+                let match_byte = mask_3.trailing_ones();
+                return 144 + match_byte as usize;
+            }
+
+            // 64 bytes (overlaps the previous 16 bytes for fast tail processing)
+            let zmm_src0_4 = _mm512_loadu_si512(src0[192..].as_ptr().cast());
+            let zmm_src1_4 = _mm512_loadu_si512(src1[192..].as_ptr().cast());
+            let mask_4 = _mm512_cmpeq_epu8_mask(zmm_src0_4, zmm_src1_4);
+            if mask_4 != 0xFFFFFFFFFFFFFFFF {
+                let match_byte = mask_4.trailing_ones();
+                return 192 + match_byte as usize;
+            }
+        }
+
+        256
+    }
+
+    #[test]
+    fn test_compare256() {
+        if true {
+            let str1 = [b'a'; super::MAX_COMPARE_SIZE];
+            let mut str2 = [b'a'; super::MAX_COMPARE_SIZE];
+
+            for i in 0..str1.len() {
+                str2[i] = 0;
+
+                let match_len = unsafe { compare256(&str1, &str2) };
+                assert_eq!(match_len, i);
+
+                str2[i] = b'a';
+            }
+        }
+    }
+}
+
 #[cfg(target_arch = "wasm32")]
 mod wasm32 {
     use core::arch::wasm32::{u8x16_bitmask, u8x16_eq, v128, v128_load};