use std::arch::x86_64::*;
use std::sync::LazyLock;
use crate::scalar;
const PREFETCH_DIST: usize = 16384;
#[target_feature(enable = "avx512f")]
unsafe fn avx512_impl<const SHOULD_HASH: bool>(
bytes: &[u8],
multiplier: u32,
addend: u32,
) -> usize {
if bytes.len() < 64 + 3 {
return unsafe { avx2_impl::<SHOULD_HASH>(bytes, multiplier, addend) };
}
let vmul = _mm512_set1_epi32(multiplier as i32);
let vadd = _mm512_set1_epi32(addend as i32);
let mut accum_lo = _mm512_set1_epi64(u64::MAX as i64);
let mut accum_hi = _mm512_set1_epi64(u64::MAX as i64);
let mut offset = 0;
unsafe {
while offset + 4 <= bytes.len() {
if offset + 64 + 3 > bytes.len() {
offset = bytes.len() - 64 - 3;
}
_mm_prefetch::<_MM_HINT_T0>(bytes.as_ptr().add(offset + PREFETCH_DIST).cast());
let mut v0 = _mm512_loadu_si512(bytes.as_ptr().add(offset).cast());
let mut v1 = _mm512_loadu_si512(bytes.as_ptr().add(offset + 1).cast());
let mut v2 = _mm512_loadu_si512(bytes.as_ptr().add(offset + 2).cast());
let mut v3 = _mm512_loadu_si512(bytes.as_ptr().add(offset + 3).cast());
if SHOULD_HASH {
v0 = _mm512_mullo_epi32(v0, vmul);
v1 = _mm512_mullo_epi32(v1, vmul);
v2 = _mm512_mullo_epi32(v2, vmul);
v3 = _mm512_mullo_epi32(v3, vmul);
v0 = _mm512_add_epi32(v0, vadd);
v1 = _mm512_add_epi32(v1, vadd);
v2 = _mm512_add_epi32(v2, vadd);
v3 = _mm512_add_epi32(v3, vadd);
}
let m01 = _mm512_min_epu32(v0, v1);
let m23 = _mm512_min_epu32(v2, v3);
let m0123 = _mm512_min_epu32(m01, m23);
let voffset = _mm512_set1_epi32(offset as i32);
let pairs_lo = _mm512_unpacklo_epi32(voffset, m0123);
let pairs_hi = _mm512_unpackhi_epi32(voffset, m0123);
accum_lo = _mm512_min_epu64(accum_lo, pairs_lo);
accum_hi = _mm512_min_epu64(accum_hi, pairs_hi);
offset += 64;
}
static SUBREG_OFFSET_LO: [u64; 8] = [0, 4, 16, 20, 32, 36, 48, 52];
static SUBREG_OFFSET_HI: [u64; 8] = [8, 12, 24, 28, 40, 44, 56, 60];
accum_lo = _mm512_add_epi64(
accum_lo,
_mm512_loadu_epi64(SUBREG_OFFSET_LO.as_ptr().cast()),
);
accum_hi = _mm512_add_epi64(
accum_hi,
_mm512_loadu_epi64(SUBREG_OFFSET_HI.as_ptr().cast()),
);
accum_lo = _mm512_min_epu64(accum_lo, accum_hi);
let mut out = [0u64; 8];
_mm512_storeu_si512(out.as_mut_ptr().cast(), accum_lo);
out.iter().copied().min().unwrap() as u32 as usize
}
}
#[target_feature(enable = "avx2")]
unsafe fn avx2_impl<const SHOULD_HASH: bool>(bytes: &[u8], multiplier: u32, addend: u32) -> usize {
if bytes.len() < 32 + 3 {
return sse41_impl::<SHOULD_HASH>(bytes, multiplier, addend);
}
let vmul = _mm256_set1_epi32(multiplier as i32);
let vadd = _mm256_set1_epi32((addend as i32).wrapping_add(i32::MIN));
let mut offset = 0;
let mut min_val = _mm256_set1_epi32(i32::MAX);
let mut min_offset = _mm256_set1_epi32(0);
let mut body = |offset: &mut usize| unsafe {
let mut v0 = _mm256_loadu_si256(bytes.as_ptr().add(*offset).cast());
let mut v1 = _mm256_loadu_si256(bytes.as_ptr().add(*offset + 1).cast());
let mut v2 = _mm256_loadu_si256(bytes.as_ptr().add(*offset + 2).cast());
let mut v3 = _mm256_loadu_si256(bytes.as_ptr().add(*offset + 3).cast());
let m0123s;
if SHOULD_HASH {
v0 = _mm256_mullo_epi32(v0, vmul);
v1 = _mm256_mullo_epi32(v1, vmul);
v2 = _mm256_mullo_epi32(v2, vmul);
v3 = _mm256_mullo_epi32(v3, vmul);
v0 = _mm256_add_epi32(v0, vadd);
v1 = _mm256_add_epi32(v1, vadd);
v2 = _mm256_add_epi32(v2, vadd);
v3 = _mm256_add_epi32(v3, vadd);
let m01s = _mm256_min_epi32(v0, v1);
let m23s = _mm256_min_epi32(v2, v3);
m0123s = _mm256_min_epi32(m01s, m23s);
} else {
let m01u = _mm256_min_epu32(v0, v1);
let m23u = _mm256_min_epu32(v2, v3);
let m0123u = _mm256_min_epu32(m01u, m23u);
m0123s = _mm256_add_epi32(m0123u, _mm256_set1_epi32(i32::MIN));
}
let voffset = _mm256_set1_epi32(*offset as i32);
let better = _mm256_cmpgt_epi32(min_val, m0123s);
min_val = _mm256_min_epi32(min_val, m0123s);
min_offset = _mm256_blendv_epi8(min_offset, voffset, better);
*offset += 32;
};
unsafe {
while offset + 64 + 4 <= bytes.len() {
_mm_prefetch::<_MM_HINT_T0>(bytes.as_ptr().add(offset + PREFETCH_DIST).cast());
body(&mut offset);
body(&mut offset);
}
while offset + 4 <= bytes.len() {
if offset + 32 + 3 > bytes.len() {
offset = bytes.len() - 32 - 3;
}
body(&mut offset);
}
static SUBREG_OFFSET: [u32; 8] = [0, 4, 8, 12, 16, 20, 24, 28];
min_offset = _mm256_add_epi32(
min_offset,
_mm256_loadu_si256(SUBREG_OFFSET.as_ptr().cast()),
);
let pairs_lo = _mm256_unpacklo_epi32(min_offset, min_val);
let pairs_hi = _mm256_unpackhi_epi32(min_offset, min_val);
let mut out = [0i64; 8];
_mm256_storeu_si256(out.as_mut_ptr().cast(), pairs_lo);
_mm256_storeu_si256(out.as_mut_ptr().add(4).cast(), pairs_hi);
out.iter().copied().min().unwrap() as u32 as usize
}
}
#[target_feature(enable = "sse4.1")]
fn sse41_impl<const SHOULD_HASH: bool>(bytes: &[u8], multiplier: u32, addend: u32) -> usize {
if bytes.len() < 16 + 3 {
return scalar::argmin_u32_overlapping_hashed::<SHOULD_HASH>(bytes, multiplier, addend);
}
let vmul = _mm_set1_epi32(multiplier as i32);
let vadd = _mm_set1_epi32((addend as i32).wrapping_add(i32::MIN));
let mut offset = 0;
let mut min_val = _mm_set1_epi32(i32::MAX);
let mut min_offset = _mm_set1_epi32(0);
let mut body = |offset: &mut usize| unsafe {
let mut v0 = _mm_loadu_si128(bytes.as_ptr().add(*offset).cast());
let mut v1 = _mm_loadu_si128(bytes.as_ptr().add(*offset + 1).cast());
let mut v2 = _mm_loadu_si128(bytes.as_ptr().add(*offset + 2).cast());
let mut v3 = _mm_loadu_si128(bytes.as_ptr().add(*offset + 3).cast());
let m0123s;
if SHOULD_HASH {
v0 = _mm_mullo_epi32(v0, vmul);
v1 = _mm_mullo_epi32(v1, vmul);
v2 = _mm_mullo_epi32(v2, vmul);
v3 = _mm_mullo_epi32(v3, vmul);
v0 = _mm_add_epi32(v0, vadd);
v1 = _mm_add_epi32(v1, vadd);
v2 = _mm_add_epi32(v2, vadd);
v3 = _mm_add_epi32(v3, vadd);
let m01s = _mm_min_epi32(v0, v1);
let m23s = _mm_min_epi32(v2, v3);
m0123s = _mm_min_epi32(m01s, m23s);
} else {
let m01u = _mm_min_epu32(v0, v1);
let m23u = _mm_min_epu32(v2, v3);
let m0123u = _mm_min_epu32(m01u, m23u);
m0123s = _mm_add_epi32(m0123u, _mm_set1_epi32(i32::MIN));
}
let voffset = _mm_set1_epi32(*offset as i32);
let better = _mm_cmpgt_epi32(min_val, m0123s);
min_val = _mm_min_epi32(min_val, m0123s);
min_offset = _mm_blendv_epi8(min_offset, voffset, better);
*offset += 16;
};
unsafe {
while offset + 32 + 4 <= bytes.len() {
_mm_prefetch::<_MM_HINT_T0>(bytes.as_ptr().add(offset + PREFETCH_DIST).cast());
body(&mut offset);
body(&mut offset);
}
while offset + 4 <= bytes.len() {
if offset + 16 + 3 > bytes.len() {
offset = bytes.len() - 16 - 3;
}
body(&mut offset);
}
static SUBREG_OFFSET: [u32; 4] = [0, 4, 8, 12];
min_offset = _mm_add_epi32(min_offset, _mm_loadu_si128(SUBREG_OFFSET.as_ptr().cast()));
let pairs_lo = _mm_unpacklo_epi32(min_offset, min_val);
let pairs_hi = _mm_unpackhi_epi32(min_offset, min_val);
let mut out = [0i64; 4];
_mm_storeu_si128(out.as_mut_ptr().cast(), pairs_lo);
_mm_storeu_si128(out.as_mut_ptr().add(2).cast(), pairs_hi);
out.iter().copied().min().unwrap() as u32 as usize
}
}
type ArgMinFn = unsafe fn(&[u8], u32, u32) -> usize;
static ARGMIN_IMPL: LazyLock<ArgMinFn> = LazyLock::new(|| {
if is_x86_feature_detected!("avx512f") {
avx512_impl::<false>
} else if is_x86_feature_detected!("avx2") {
avx2_impl::<false>
} else if is_x86_feature_detected!("sse4.1") {
sse41_impl::<false>
} else {
scalar::argmin_u32_overlapping_hashed::<false>
}
});
static ARGMIN_HASH_IMPL: LazyLock<ArgMinFn> = LazyLock::new(|| {
if is_x86_feature_detected!("avx512f") {
avx512_impl::<true>
} else if is_x86_feature_detected!("avx2") {
avx2_impl::<true>
} else if is_x86_feature_detected!("sse4.1") {
sse41_impl::<true>
} else {
scalar::argmin_u32_overlapping_hashed::<true>
}
});
pub fn argmin_u32_overlapping_hashed<const SHOULD_HASH: bool>(
bytes: &[u8],
multiplier: u32,
addend: u32,
) -> usize {
const MIN_SIMD_LEN: usize = 16 + 3;
assert!(bytes.len() <= u32::MAX as usize);
let within_four_offset = if SHOULD_HASH {
unsafe { ARGMIN_HASH_IMPL(bytes, multiplier, addend) }
} else {
unsafe { ARGMIN_IMPL(bytes, multiplier, addend) }
};
if bytes.len() >= MIN_SIMD_LEN {
let final_bump = scalar::argmin_u32_overlapping_hashed_four::<SHOULD_HASH>(
&bytes[within_four_offset..],
multiplier,
addend,
);
within_four_offset + final_bump
} else {
within_four_offset
}
}