port avx512 implementation of compare256

Author: folkertdev

zlib-rs/src/deflate/compare256.rs

@@ -9,6 +9,10 @@ pub fn compare256_slice(src0: &[u8], src1: &[u8]) -> usize {
 }
 
 fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
+    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 +267,92 @@ mod avx2 {
     }
 }
 
+#[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};