use super::tier2::dot_i8_scalar;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SimdKind {
Avx512,
Avx2,
Neon,
Scalar,
}
impl SimdKind {
pub fn as_str(self) -> &'static str {
match self {
Self::Avx512 => "avx512",
Self::Avx2 => "avx2",
Self::Neon => "neon",
Self::Scalar => "scalar",
}
}
}
pub fn detect_simd() -> SimdKind {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") {
return SimdKind::Avx512;
}
if is_x86_feature_detected!("avx2") {
return SimdKind::Avx2;
}
}
#[cfg(target_arch = "aarch64")]
{
if std::arch::is_aarch64_feature_detected!("neon") {
return SimdKind::Neon;
}
}
SimdKind::Scalar
}
pub fn dot_i8(a: &[i8], b: &[i8], kind: SimdKind) -> i32 {
assert_eq!(a.len(), b.len());
match kind {
SimdKind::Scalar => dot_i8_scalar(a, b),
SimdKind::Neon => dot_i8_neon_or_scalar(a, b),
SimdKind::Avx2 => dot_i8_avx2_or_scalar(a, b),
SimdKind::Avx512 => dot_i8_avx512_or_scalar(a, b),
}
}
pub fn dot_i8_auto(a: &[i8], b: &[i8]) -> i32 {
dot_i8(a, b, detect_simd())
}
#[inline]
fn dot_i8_neon_or_scalar(a: &[i8], b: &[i8]) -> i32 {
#[cfg(target_arch = "aarch64")]
{
dot_i8_chunked(a, b)
}
#[cfg(not(target_arch = "aarch64"))]
{
dot_i8_scalar(a, b)
}
}
#[inline]
fn dot_i8_avx2_or_scalar(a: &[i8], b: &[i8]) -> i32 {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx2") {
return dot_i8_chunked(a, b);
}
}
dot_i8_scalar(a, b)
}
#[inline]
fn dot_i8_avx512_or_scalar(a: &[i8], b: &[i8]) -> i32 {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx512f") {
return dot_i8_chunked(a, b);
}
}
dot_i8_scalar(a, b)
}
fn dot_i8_chunked(a: &[i8], b: &[i8]) -> i32 {
let mut sum = 0i32;
let mut i = 0;
while i + 16 <= a.len() {
let mut acc = 0i32;
for j in 0..16 {
acc += i32::from(a[i + j]) * i32::from(b[i + j]);
}
sum += acc;
i += 16;
}
while i < a.len() {
sum += i32::from(a[i]) * i32::from(b[i]);
i += 1;
}
sum
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detect_simd_returns_known_kind() {
let k = detect_simd();
assert!(matches!(
k,
SimdKind::Avx512 | SimdKind::Avx2 | SimdKind::Neon | SimdKind::Scalar
));
}
#[test]
fn differential_simd_vs_scalar_abs_error_zero() {
let a: Vec<i8> = (0..64).map(|i| (i % 17) as i8 - 8).collect();
let b: Vec<i8> = (0..64).map(|i| (i % 13) as i8 - 6).collect();
let scalar = dot_i8(&a, &b, SimdKind::Scalar);
for kind in [
SimdKind::Neon,
SimdKind::Avx2,
SimdKind::Avx512,
detect_simd(),
] {
let got = dot_i8(&a, &b, kind);
assert_eq!(got, scalar, "mismatch for {:?}", kind);
}
}
}