#![allow(path_statements)]
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::__m256i;
#[cfg(target_arch = "x86")]
use core::arch::x86::__m256i;
use core::any::TypeId;
use core::fmt::Debug;
use core::marker::PhantomData;
use core::mem::{transmute, transmute_copy};
use core::ops::{Add, AddAssign, Not};
use crate::sealed::Sealed;
pub trait Simd256Scalar: Sealed
+ Sized
+ Copy
+ Add<Self, Output = Self>
+ AddAssign
+ Default
+ Not<Output = Self>
+ PartialEq
+ 'static
{
const ZERO: Self;
}
pub trait Simd256SaturatingAdd: Simd256Scalar {}
pub trait Simd256IntegerAbs: Simd256Scalar {}
macro_rules! impl_scalars {
( $($t:ty $(| $extra:ident )*)* ) => {$(
impl Simd256Scalar for $t {
const ZERO: $t = 0;
}
$(
impl $extra for $t {}
)*
)*};
}
impl_scalars! {
u8
| Simd256SaturatingAdd
i8
| Simd256SaturatingAdd
| Simd256IntegerAbs
u16
| Simd256SaturatingAdd
i16
| Simd256SaturatingAdd
| Simd256IntegerAbs
u32
i32
| Simd256IntegerAbs
u64
i64
}
#[allow(non_camel_case_types)]
pub type u8x32 = Simd256Integer<u8, 32>;
#[allow(non_camel_case_types)]
pub type i8x32 = Simd256Integer<i8, 32>;
#[allow(non_camel_case_types)]
pub type u16x16 = Simd256Integer<u16, 16>;
#[allow(non_camel_case_types)]
pub type i16x16 = Simd256Integer<i16, 16>;
#[allow(non_camel_case_types)]
pub type u32x8 = Simd256Integer<u32, 8>;
#[allow(non_camel_case_types)]
pub type i32x8 = Simd256Integer<i32, 8>;
#[allow(non_camel_case_types)]
pub type u64x4 = Simd256Integer<u64, 4>;
#[allow(non_camel_case_types)]
pub type i64x4 = Simd256Integer<i64, 4>;
#[derive(Clone, Copy, Debug)]
pub struct Simd256Integer<S: Simd256Scalar, const LANES: usize> {
phantom: PhantomData<[S; LANES]>,
pub inner: Simd256IntegerInner,
}
#[derive(Clone, Copy)]
pub union Simd256IntegerInner {
#[cfg(any(feature = "std", target_feature = "avx"))]
pub avx: __m256i,
pub fallback: [u8; size_of::<__m256i>() / size_of::<u8>()],
}
impl Debug for Simd256IntegerInner {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
#[cfg(all(not(feature = "std"), target_feature = "avx2"))]
{
return f
.debug_struct(stringify!(Simd256IntegerInner))
.field("avx", &unsafe { self.avx })
.finish();
}
#[cfg(feature = "std")]
if std::is_x86_feature_detected!("avx2") {
return f
.debug_struct(stringify!(Simd256IntegerInner))
.field("avx", &unsafe { self.avx })
.finish();
}
f.debug_struct(stringify!(Simd256IntegerInner))
.field("fallback", &unsafe { self.fallback.as_ref() })
.finish()
}
}
impl<S: Simd256Scalar, const LANES: usize> Simd256Integer<S, LANES> {
const _MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE: () = assert!(
LANES == size_of::<__m256i>() / size_of::<S>(),
"The number of lanes needs to be consistent with the size of the SIMD vector for the scalar type."
);
pub fn from_array(array: [S; LANES]) -> Self {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
unsafe { transmute_copy(&array) }
}
pub fn splat(s: S) -> Self {
Simd256Integer::from_array([s; LANES])
}
pub fn try_from_iter(iter: &mut impl Iterator<Item = S>) -> Option<Self> {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
let mut array: [S; LANES] = [S::ZERO; LANES];
#[allow(clippy::needless_range_loop)]
for i in 0..LANES {
array[i] = iter.next()?;
}
Some(Simd256Integer::from_array(array))
}
pub const fn to_array(self) -> [S; LANES] {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
unsafe { transmute_copy(&self) }
}
pub fn as_array_ref(&self) -> &[S; LANES] {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
unsafe { transmute(self) }
}
#[inline(always)]
pub const fn from_intrinsic(intrinsic: __m256i) -> Self {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
unsafe { transmute(intrinsic) }
}
#[inline(always)]
pub fn element_type_is<T: Simd256Scalar>() -> bool {
TypeId::of::<S>() == TypeId::of::<T>()
}
#[cfg(any(feature = "std", target_feature = "avx2"))]
#[target_feature(enable = "avx2")]
pub unsafe fn avx2_vertical_add(a: Self, b: Self) -> Self {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
let result = match size_of::<S>() {
1 => _mm256_add_epi8(a.inner.avx, b.inner.avx),
2 => _mm256_add_epi16(a.inner.avx, b.inner.avx),
4 => _mm256_add_epi32(a.inner.avx, b.inner.avx),
8 => _mm256_add_epi64(a.inner.avx, b.inner.avx),
_ => crate::unreachable_uncheched_on_release(),
};
Self::from_intrinsic(result)
}
#[cfg(any(feature = "std", target_feature = "avx2"))]
#[target_feature(enable = "avx2")]
pub unsafe fn avx2_vertical_saturating_add(a: Self, b: Self) -> Self
where
S: Simd256SaturatingAdd,
{
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
let result = match size_of::<S>() {
1 if Self::element_type_is::<u8>() => {
_mm256_adds_epu8(a.inner.avx, b.inner.avx)
}
1 if Self::element_type_is::<i8>() => {
_mm256_adds_epi8(a.inner.avx, b.inner.avx)
}
2 if Self::element_type_is::<u16>() => {
_mm256_adds_epu16(a.inner.avx, b.inner.avx)
}
2 if Self::element_type_is::<i16>() => {
_mm256_adds_epi16(a.inner.avx, b.inner.avx)
}
_ => crate::unreachable_uncheched_on_release(),
};
Self::from_intrinsic(result)
}
pub fn saturating_add(a: Self, b: Self) -> Self
where S: Simd256SaturatingAdd
{
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_feature = "avx2")]
return unsafe { Self::avx2_vertical_saturating_add(a, b) };
#[cfg(feature = "std")]
if std::is_x86_feature_detected!("avx2") {
return unsafe { Self::avx2_vertical_saturating_add(a, b) };
}
match size_of::<S>() {
1 if Self::element_type_is::<u8>() => {
let a = unsafe { transmute::<_, u8x32>(a) }.to_array();
let b = unsafe { transmute::<_, u8x32>(b) }.to_array();
let mut result: [u8; 32] = [0; 32];
for i in 0..LANES {
result[i] = u8::saturating_add(a[i], b[i]);
}
unsafe { transmute(result) }
}
1 if Self::element_type_is::<i8>() => {
let a = unsafe { transmute::<_, i8x32>(a) }.to_array();
let b = unsafe { transmute::<_, i8x32>(b) }.to_array();
let mut result: [i8; 32] = [0; 32];
for i in 0..LANES {
result[i] = i8::saturating_add(a[i], b[i]);
}
unsafe { transmute(result) }
}
2 if Self::element_type_is::<u16>() => {
let a = unsafe { transmute::<_, u16x16>(a) }.to_array();
let b = unsafe { transmute::<_, u16x16>(b) }.to_array();
let mut result: [u16; 16] = [0; 16];
for i in 0..LANES {
result[i] = u16::saturating_add(a[i], b[i]);
}
unsafe { transmute(result) }
}
2 if Self::element_type_is::<i16>() => {
let a = unsafe { transmute::<_, i16x16>(a) }.to_array();
let b = unsafe { transmute::<_, i16x16>(b) }.to_array();
let mut result: [i16; 16] = [0; 16];
for i in 0..LANES {
result[i] = i16::saturating_add(a[i], b[i]);
}
unsafe { transmute(result) }
}
_ => unsafe { crate::unreachable_uncheched_on_release() }
}
}
#[cfg(any(feature = "std", target_feature = "avx2"))]
#[target_feature(enable = "avx2")]
pub unsafe fn avx2_vertical_cmp_eq(a: Self, b: Self) -> Self {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
let result = match size_of::<S>() {
1 => _mm256_cmpeq_epi8(a.inner.avx, b.inner.avx),
2 => _mm256_cmpeq_epi16(a.inner.avx, b.inner.avx),
4 => _mm256_cmpeq_epi32(a.inner.avx, b.inner.avx),
8 => _mm256_cmpeq_epi64(a.inner.avx, b.inner.avx),
_ => crate::unreachable_uncheched_on_release(),
};
Self::from_intrinsic(result)
}
pub fn vertical_cmp_eq(a: Self, b: Self) -> Self {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_feature = "avx2")]
return unsafe { Self::avx2_vertical_cmp_eq(a, b) };
#[cfg(feature = "std")]
if std::is_x86_feature_detected!("avx2") {
return unsafe { Self::avx2_vertical_cmp_eq(a, b) };
}
let mut result = [S::ZERO; LANES];
#[allow(clippy::needless_range_loop)]
for i in 0..LANES {
if a.as_array_ref()[i] == b.as_array_ref()[i] {
result[i] = !S::ZERO;
}
}
Self::from_array(result)
}
#[cfg(any(feature = "std", target_feature = "avx2"))]
#[target_feature(enable = "avx2")]
pub unsafe fn avx2_vertical_abs(self) -> Self
where S: Simd256IntegerAbs
{
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
let result = match size_of::<S>() {
1 => _mm256_abs_epi8(self.inner.avx),
2 => _mm256_abs_epi16(self.inner.avx),
4 => _mm256_abs_epi32(self.inner.avx),
_ => crate::unreachable_uncheched_on_release(),
};
Self::from_intrinsic(result)
}
pub fn abs(self) -> Self
where S: Simd256IntegerAbs {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_feature = "avx2")]
return unsafe { Self::avx2_vertical_abs(self) };
#[cfg(feature = "std")]
if std::is_x86_feature_detected!("avx2") {
return unsafe { Self::avx2_vertical_abs(self) };
}
match size_of::<S>() {
1 => unsafe {
let mut array = transmute::<_, i8x32>(self).to_array();
for element in &mut array {
*element = i8::abs(*element);
}
transmute(array)
}
2 => unsafe {
let mut array = transmute::<_, i16x16>(self).to_array();
for element in &mut array {
*element = i16::abs(*element);
}
transmute(array)
}
4 => unsafe {
let mut array = transmute::<_, i32x8>(self).to_array();
for element in &mut array {
*element = i32::abs(*element);
}
transmute(array)
}
_ => unsafe { crate::unreachable_uncheched_on_release() }
}
}
#[cfg(any(feature = "std", target_feature = "avx"))]
#[target_feature(enable = "avx")]
pub unsafe fn avx_load(ptr: *const __m256i) -> Self {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
transmute(_mm256_loadu_si256(ptr))
}
pub fn load_from_slice(slice: &[S]) -> Self {
Self::try_load_from_slice(slice)
.expect("slice must contain enough elements to load into a SIMD vector")
}
pub fn try_load_from_slice(slice: &[S]) -> Option<Self> {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
if slice.len() < LANES {
return None;
}
#[cfg(any(feature = "std", target_feature = "avx"))]
let cptr = slice as *const [S] as *const () as *const __m256i;
#[cfg(target_feature = "avx")]
return Some(unsafe { Self::avx_load(cptr) });
#[cfg(feature = "std")]
if std::is_x86_feature_detected!("avx") {
return Some(unsafe { Self::avx_load(cptr) });
}
let mut result = [S::ZERO; LANES];
result[..LANES].copy_from_slice(&slice[..LANES]);
Some(Self::from_array(result))
}
}
impl<S: Simd256Scalar, const LANES: usize> core::ops::Add for Simd256Integer<S, LANES> {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
Self::_MENTION_ME_TO_ASSERT_LANES_MATCH_SIZE;
#[cfg(target_feature = "avx2")]
return unsafe { Simd256Integer::avx2_vertical_add(self, rhs) };
#[cfg(feature = "std")]
if std::is_x86_feature_detected!("avx2") {
return unsafe { Simd256Integer::avx2_vertical_add(self, rhs) };
}
let mut result: [S; LANES] = [S::ZERO; LANES];
let a = self.to_array();
let b = rhs.to_array();
for i in 0..result.len() {
result[i] = a[i] + b[i];
}
Simd256Integer::from_array(result)
}
}
impl<S: Simd256Scalar, const LANES: usize> PartialEq for Simd256Integer<S, LANES> {
fn eq(&self, other: &Self) -> bool {
self.as_array_ref() == other.as_array_ref()
}
}
impl<S: Simd256Scalar, const LANES: usize> Eq for Simd256Integer<S, LANES> {}
#[cfg(test)]
mod tests {
use core::u16;
use crate::integers::int256::i64x4;
use super::i32x8;
use super::i8x32;
use super::u64x4;
use super::u8x32;
use super::u16x16;
use super::i16x16;
#[test]
#[cfg(feature = "std")]
fn test_debug() {
let simd_value = u8x32::try_from_iter(&mut (0..32).into_iter()).unwrap();
println!("{simd_value:#x?}");
}
#[test]
fn test_add() {
let simd_10x16 = u16x16::splat(10);
let simd_12x16 = u16x16::splat(12);
let added = simd_10x16 + simd_12x16;
assert_eq!(added.to_array(), [22; 16]);
}
#[test]
fn test_saturating_add() {
let simd_max = u16x16::splat(u16::MAX);
assert_eq!(u16x16::saturating_add(simd_max, simd_max).as_array_ref(), simd_max.as_array_ref());
}
#[test]
fn test_abs() {
let simd_neg1 = i16x16::splat(-1);
assert_eq!(simd_neg1.abs().to_array(), [1; 16]);
}
#[test]
fn test_vertical_cmp_eq() {
let simd_a = u64x4::from_array([10, 20, 30, 40]);
let simd_b = u64x4::from_array([0, 20, 40, 60]);
assert_eq!(u64x4::vertical_cmp_eq(simd_a, simd_b).to_array(), [0, 0xFFFF_FFFF_FFFF_FFFF, 0, 0]);
}
#[test]
fn test_try_load_fails() {
assert!(i8x32::try_load_from_slice(&[0;10]).is_none());
}
#[test]
#[should_panic]
fn test_load_panics() {
i32x8::load_from_slice(&[0; 1]);
}
#[test]
fn test_load() {
assert_eq!(i64x4::load_from_slice(&[100; 4]).to_array(), [100; 4]);
}
}