use super::traits::VecTrait;
use crate::{dtype::TypeCommon, type_promote::NormalOut};
#[macro_export]
macro_rules! impl_std_simd_bit_logic {
($ty: ident) => {
impl std::ops::BitAnd for $ty {
type Output = Self;
fn bitand(self, rhs: Self) -> Self::Output {
$ty(self.0 & rhs.0)
}
}
impl std::ops::BitOr for $ty {
type Output = Self;
fn bitor(self, rhs: Self) -> Self::Output {
$ty(self.0 | rhs.0)
}
}
impl std::ops::BitXor for $ty {
type Output = Self;
fn bitxor(self, rhs: Self) -> Self::Output {
$ty(self.0 ^ rhs.0)
}
}
impl std::ops::Not for $ty {
type Output = Self;
fn not(self) -> Self::Output {
$ty(!self.0)
}
}
impl std::ops::Shl for $ty {
type Output = Self;
fn shl(self, rhs: Self) -> Self {
$ty(self.0 << rhs.0)
}
}
impl std::ops::Shr for $ty {
type Output = Self;
fn shr(self, rhs: Self) -> Self {
$ty(self.0 >> rhs.0)
}
}
};
}
#[inline(always)]
pub fn vec_sum<T: TypeCommon + NormalOut<T, Output = T>>(vec: T::Vec) -> T {
let mut sum = T::ZERO;
for i in 0..T::Vec::SIZE {
sum = sum._add(vec.extract(i));
}
sum
}
#[inline(always)]
pub fn array_vec_sum<T: TypeCommon + NormalOut<T, Output = T> + Copy>(array: &[T]) -> T {
let remain = array.len() % T::Vec::SIZE;
let vecs = unsafe {
std::slice::from_raw_parts(
array as *const _ as *const T::Vec,
array.len() / T::Vec::SIZE,
)
};
let mut sum_vec = T::Vec::splat(T::ZERO);
for vec in vecs.iter() {
sum_vec = sum_vec._add(vec.read_unaligned());
}
let mut sum = T::ZERO;
for i in array.iter().skip(array.len() - remain) {
sum = sum._add(*i);
}
for i in 0..T::Vec::SIZE {
sum = sum._add(sum_vec.extract(i));
}
sum
}
#[inline(always)]
pub fn array_vec_reduce<T: TypeCommon + NormalOut<T, Output = T> + Copy>(
array: &[T],
init: T,
scalar_preop: impl Fn(T) -> T,
scalar_cumulate: impl Fn(T, T) -> T,
vec_preop: impl Fn(T::Vec) -> T::Vec,
vec_cumulate: impl Fn(T::Vec, T::Vec) -> T::Vec,
) -> T {
let remain = array.len() % T::Vec::SIZE;
let vecs = array as *const _ as *const T::Vec;
let mut red_vec = T::Vec::splat(init);
for i in 0..array.len() / T::Vec::SIZE {
red_vec = vec_cumulate(red_vec, vec_preop(unsafe { vecs.add(i).read_unaligned() }));
}
let mut red = init;
for i in array.iter().skip(array.len() - remain) {
red = scalar_cumulate(red, scalar_preop(*i));
}
for i in 0..T::Vec::SIZE {
red = scalar_cumulate(red, red_vec.extract(i));
}
red
}