Source code
Revision control
Copy as Markdown
Other Tools
#[cfg(test)]
const MAX_COMPARE_SIZE: usize = 256;
pub fn compare256_slice(src0: &[u8], src1: &[u8]) -> usize {
let src0 = first_chunk::<_, 256>(src0).unwrap();
let src1 = first_chunk::<_, 256>(src1).unwrap();
compare256(src0, src1)
}
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) };
}
#[cfg(target_arch = "aarch64")]
if crate::cpu_features::is_enabled_neon() {
return unsafe { neon::compare256(src0, src1) };
}
#[cfg(target_arch = "wasm32")]
if crate::cpu_features::is_enabled_simd128() {
return wasm32::compare256(src0, src1);
}
rust::compare256(src0, src1)
}
pub fn compare256_rle_slice(byte: u8, src: &[u8]) -> usize {
rust::compare256_rle(byte, src)
}
#[inline]
pub const fn first_chunk<T, const N: usize>(slice: &[T]) -> Option<&[T; N]> {
if slice.len() < N {
None
} else {
// SAFETY: We explicitly check for the correct number of elements,
// and do not let the reference outlive the slice.
Some(unsafe { &*(slice.as_ptr() as *const [T; N]) })
}
}
mod rust {
pub fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
// only unrolls 4 iterations; zlib-ng unrolls 8
src0.iter().zip(src1).take_while(|(x, y)| x == y).count()
}
// run-length encoding
pub fn compare256_rle(byte: u8, src: &[u8]) -> usize {
assert!(src.len() >= 256, "too short {}", src.len());
let sv = u64::from_ne_bytes([byte; 8]);
let mut len = 0;
// this optimizes well because we statically limit the slice to 256 bytes.
// the loop gets unrolled 4 times automatically.
for chunk in src[..256].chunks_exact(8) {
let mv = u64::from_le_bytes(chunk.try_into().unwrap());
let diff = sv ^ mv;
if diff > 0 {
let match_byte = diff.trailing_zeros() / 8;
return len + match_byte as usize;
}
len += 8
}
256
}
#[test]
fn test_compare256() {
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 = compare256(&str1, &str2);
assert_eq!(match_len, i);
str2[i] = b'a';
}
}
#[test]
fn test_compare256_rle() {
let mut string = [b'a'; super::MAX_COMPARE_SIZE];
for i in 0..string.len() {
string[i] = 0;
let match_len = compare256_rle(b'a', &string);
assert_eq!(match_len, i);
string[i] = b'a';
}
}
}
#[cfg(target_arch = "aarch64")]
mod neon {
use core::arch::aarch64::{
uint8x16x4_t, vceqq_u8, vget_lane_u64, vld4q_u8, vreinterpret_u64_u8, vreinterpretq_u16_u8,
vshrn_n_u16, vsriq_n_u8,
};
/// # Safety
///
/// Behavior is undefined if the `neon` target feature is not enabled
#[target_feature(enable = "neon")]
pub unsafe fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
type Chunk = uint8x16x4_t;
let src0 = src0.chunks_exact(core::mem::size_of::<Chunk>());
let src1 = src1.chunks_exact(core::mem::size_of::<Chunk>());
let mut len = 0;
for (a, b) in src0.zip(src1) {
unsafe {
// Load 4 vectors *deinterleaved* from the two slices
// e.g. the first vector contains the 0, 4, 8, ... bytes of the input, the
// second vector contains the 1, 5, 9, ... bytes of the input, etc.
let a: Chunk = vld4q_u8(a.as_ptr());
let b: Chunk = vld4q_u8(b.as_ptr());
// Compare each vector element-wise, each resulting vector will contain
// 0xFF for equal bytes, and 0x00 for unequal bytes.
let cmp0 = vceqq_u8(a.0, b.0);
let cmp1 = vceqq_u8(a.1, b.1);
let cmp2 = vceqq_u8(a.2, b.2);
let cmp3 = vceqq_u8(a.3, b.3);
// Pack bits from the 4 vectors into a single vector to convert to a 64-bit integer.
// shift the second vector right by one, insert the top bit from the first vector
// The top two bits each element of the result are from the first and second vector
let first_two_bits = vsriq_n_u8::<1>(cmp1, cmp0);
// shift the fourth vector right by one, insert the top bit from the third vector
// The top two bits each element of the result are from the third and fourth vector
let last_two_bits = vsriq_n_u8::<1>(cmp3, cmp2);
// shift last_two_bits (the top two bits of which are from the third and fourth
// vector) right by 2, insert the top two bits from first_two_bits (the top two
// bits of which are from the first and second vector).
// The top four bits of each element of the result are from the
// first, second, third, and fourth vector
let first_four_bits = vsriq_n_u8::<2>(last_two_bits, first_two_bits);
// duplicate the top 4 bits into the bottom 4 bits of each element.
let bitmask_vector = vsriq_n_u8::<4>(first_four_bits, first_four_bits);
// Reinterpret as 16-bit integers, and shift right by 4 bits narrowing:
// shifting right by 4 bits means the top 4 bits of each 16 bit element contains the
// low 4 bits of the 0th 8-bit element and the high 4 bits of the 1nth 8-bit
// element. Narrowing takes the top 8 bits of each (16-bit) element.
let result_vector = vshrn_n_u16::<4>(vreinterpretq_u16_u8(bitmask_vector));
// Convert the vector to a 64-bit integer, where each bit represents whether
// the corresponding byte in the original vectors was equal.
let bitmask = vget_lane_u64::<0>(vreinterpret_u64_u8(result_vector));
// We reinterpreted the vector as a 64-bit integer, so endianness matters.
// We want things to be in little-endian (where the least significant bit is in the
// first byte), but in big-endian, the first vector element will be the most
// significant byte, so we need to convert to little-endian.
let bitmask = bitmask.to_le();
if bitmask != u64::MAX {
// Find the first byte that is not equal, which is the first bit that is not set
let match_byte = bitmask.trailing_ones();
return len + match_byte as usize;
}
len += core::mem::size_of::<Chunk>();
}
}
256
}
#[test]
fn test_compare256() {
if crate::cpu_features::is_enabled_neon() {
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 = "x86_64")]
mod avx2 {
use core::arch::x86_64::{
__m256i, _mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8,
};
/// # Safety
///
/// Behavior is undefined if the `avx` target feature is not enabled
#[target_feature(enable = "avx2")]
#[target_feature(enable = "bmi2")]
#[target_feature(enable = "bmi1")]
pub unsafe fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
let src0 = src0.chunks_exact(32);
let src1 = src1.chunks_exact(32);
let mut len = 0;
unsafe {
for (chunk0, chunk1) in src0.zip(src1) {
let ymm_src0 = _mm256_loadu_si256(chunk0.as_ptr() as *const __m256i);
let ymm_src1 = _mm256_loadu_si256(chunk1.as_ptr() as *const __m256i);
// element-wise compare of the 8-bit elements
let ymm_cmp = _mm256_cmpeq_epi8(ymm_src0, ymm_src1);
// turn an 32 * 8-bit vector into a 32-bit integer.
// a bit in the output is one if the corresponding element is non-zero.
let mask = _mm256_movemask_epi8(ymm_cmp) as u32;
if mask != 0xFFFFFFFF {
let match_byte = mask.trailing_ones();
return len + match_byte as usize;
}
len += 32;
}
}
256
}
#[test]
fn test_compare256() {
if crate::cpu_features::is_enabled_avx2_and_bmi2() {
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(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};
#[target_feature(enable = "simd128")]
pub fn compare256(src0: &[u8; 256], src1: &[u8; 256]) -> usize {
let src0 = src0.chunks_exact(16);
let src1 = src1.chunks_exact(16);
let mut len = 0;
for (chunk0, chunk1) in src0.zip(src1) {
// SAFETY: these are valid pointers to slice data.
let v128_src0 = unsafe { v128_load(chunk0.as_ptr() as *const v128) };
let v128_src1 = unsafe { v128_load(chunk1.as_ptr() as *const v128) };
let v128_cmp = u8x16_eq(v128_src0, v128_src1);
let mask = u8x16_bitmask(v128_cmp);
if mask != 0xFFFF {
let match_byte = mask.trailing_ones();
return len + match_byte as usize;
}
len += 16;
}
256
}
#[test]
fn test_compare256() {
if crate::cpu_features::is_enabled_simd128() {
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';
}
}
}
}