use core::marker::PhantomData;
use core::ops::{Add, Div, Mul, Rem, Sub};
use num_traits::Bounded;
use thermite::Swizzle;
use thermite::element::{Element, FloatElement, SignedElement};
use thermite::generic_array::{GenericArray, IntoArrayLength, typenum::Const};
use thermite::mask::{GenericMask, GenericSelectable};
use thermite::math::RealMathWithPolicy;
use thermite::math::algorithms::reduce_in_place;
use thermite::math::policy::DefaultPolicy;
use thermite::register::SwizzleIndices;
use thermite::vector::ops::{AddSubExt, AddSubExtMasked, NegMasked, Square, SquareMasked};
use thermite::vector::{NewConst, NewVector, SplatConst, SplatVector, VectorValue, const_new, const_splat};
use thermite::{LargeInt, features, prelude::*};
use crate::{Complex, RealValue};
pub trait RealFloatVector:
RealValue + FloatVector<Element: RealValue> + CastVector<Self> + RealMathWithPolicy + SwizzleVector
{
}
impl<V> RealFloatVector for V where
V: RealValue + FloatVector<Element: RealValue> + CastVector<V> + RealMathWithPolicy + SwizzleVector
{
}
impl<V: RealFloatVector> Swizzle<V::Lanes> for Complex<V> {
#[inline(always)]
fn swizzle(self, other: Self, indices: GenericArray<u32, V::Lanes>) -> Self {
Self {
re: self.re.swizzle(other.re, indices.clone()),
im: self.im.swizzle(other.im, indices),
}
}
#[inline(always)]
fn permute(self, indices: GenericArray<u32, V::Lanes>) -> Self {
Self {
re: self.re.permute(indices.clone()),
im: self.im.permute(indices),
}
}
#[inline(always)]
fn swizzle_const<I: SwizzleIndices<V::Lanes>>(self, other: Self) -> Self {
Self {
re: self.re.swizzle_const::<I>(other.re),
im: self.im.swizzle_const::<I>(other.im),
}
}
#[inline(always)]
fn permute_const<I: SwizzleIndices<V::Lanes>>(self) -> Self {
Self {
re: self.re.permute_const::<I>(),
im: self.im.permute_const::<I>(),
}
}
}
#[rustfmt::skip]
impl<E: RealValue + Element> Element for Complex<E> {
type Signed = <E as Element>::Signed;
type Unsigned = <E as Element>::Unsigned;
const ZERO: Self = Self::ZERO;
const ONE: Self = Self::ONE;
const ORDER_MAX: Self = Self { re: E::ORDER_MAX, im: E::ORDER_MAX };
const ORDER_MIN: Self = Self { re: E::ORDER_MIN, im: E::ORDER_MIN };
const HAS_UNORDERED: bool = E::HAS_UNORDERED;
const IS_FLOAT: bool = E::IS_FLOAT;
#[inline(always)] fn from_i8(value: i8) -> Self { Self::real(E::from_i8(value)) }
#[inline(always)] fn from_u8(value: u8) -> Self { Self::real(E::from_u8(value)) }
#[inline(always)] fn from_u16(value: u16) -> Self { Self::real(E::from_u16(value)) }
}
impl<E: RealValue + FloatElement> Complex<E> {
#[inline(always)]
fn elem_modulus(self) -> E {
E::sqrt(self.re.mul_add(self.re, self.im * self.im))
}
}
impl<E: RealValue + FloatElement> SignedElement for Complex<E> {
#[inline(always)]
fn abs(self) -> Self {
Self::real(self.elem_modulus())
}
#[inline(always)]
fn signum(self) -> Self {
let m = self.elem_modulus();
if m == E::ZERO {
return Self::ZERO;
}
Self::new(self.re / m, self.im / m)
}
}
pub struct ComplexIntConst<E, const VAL: LargeInt>(PhantomData<E>);
pub struct ComplexRatioConst<E, const NUM: LargeInt, const DEN: LargeInt>(PhantomData<E>);
impl<E: RealValue + FloatElement, const VAL: LargeInt> SplatConst<Complex<E>> for ComplexIntConst<E, VAL> {
const VALUE: Complex<E> = Complex::real(<E::ConstInt<VAL> as SplatConst<E>>::VALUE);
}
impl<E: RealValue + FloatElement, const NUM: LargeInt, const DEN: LargeInt> SplatConst<Complex<E>>
for ComplexRatioConst<E, NUM, DEN>
{
const VALUE: Complex<E> = Complex::real(<E::ConstRatio<NUM, DEN> as SplatConst<E>>::VALUE);
}
#[rustfmt::skip]
impl<E: RealValue + FloatElement> FloatElement for Complex<E> {
#[inline(always)]
fn sqrt(this: Self) -> Self {
let half = E::from_ratio(1, 2);
let t = E::sqrt((SignedElement::abs(this.re) + this.elem_modulus()) * half);
if t == E::ZERO {
return Self::ZERO;
}
let half_im = this.im * half;
if this.re >= E::ZERO {
Self::new(t, half_im / t)
} else {
let i = if this.im < E::ZERO { -t } else { t };
Self::new(SignedElement::abs(half_im) / t, i)
}
}
#[inline(always)] fn floor(this: Self) -> Self { Self::new(E::floor(this.re), E::floor(this.im)) }
#[inline(always)] fn ceil(this: Self) -> Self { Self::new(E::ceil(this.re), E::ceil(this.im)) }
#[inline(always)] fn round(this: Self) -> Self { Self::new(E::round(this.re), E::round(this.im)) }
#[inline(always)] fn trunc(this: Self) -> Self { Self::new(E::trunc(this.re), E::trunc(this.im)) }
#[inline(always)] fn next_up(this: Self) -> Self { Self::new(E::next_up(this.re), E::next_up(this.im)) }
#[inline(always)] fn next_down(this: Self) -> Self { Self::new(E::next_down(this.re), E::next_down(this.im)) }
#[inline(always)]
fn try_from_int(value: LargeInt) -> Option<Self> {
E::try_from_int(value).map(Self::real)
}
#[inline(always)]
fn try_from_ratio(n: LargeInt, d: LargeInt) -> Option<Self> {
E::try_from_ratio(n, d).map(Self::real)
}
const HAS_INFINITY: bool = E::HAS_INFINITY;
const HAS_SIGNED_ZERO: bool = E::HAS_SIGNED_ZERO;
const HAS_SUBNORMALS: bool = E::HAS_SUBNORMALS;
type ConstInt<const VAL: LargeInt> = ComplexIntConst<E, VAL>;
type ConstRatio<const NUM: LargeInt, const DEN: LargeInt> = ComplexRatioConst<E, NUM, DEN>;
}
impl<V: thermite::simd::HasIsa> thermite::simd::HasIsa for Complex<V> {
type Native = V::Native;
const ISA: thermite::isa::InstructionSet = V::ISA;
}
impl<V: RealFloatVector> GenericSelectable for Complex<V> {
type SelectableMask = <V as GenericSelectable>::SelectableMask;
#[inline(always)]
fn select<M>(mask: M, t: Self, f: Self) -> Self
where
Self::SelectableMask: CastMask<M>,
{
let mask = <Self::SelectableMask as CastMask<M>>::mask_from(mask);
Self::new(mask.select(t.re, f.re), mask.select(t.im, f.im))
}
}
#[rustfmt::skip]
impl<V: RealFloatVector> Interleave for Complex<V> {
#[inline(always)]
fn interleave(self, other: Self) -> (Self, Self) {
let (re_lo, re_hi) = self.re.interleave(other.re);
let (im_lo, im_hi) = self.im.interleave(other.im);
(Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
}
#[inline(always)]
fn deinterleave(self, other: Self) -> (Self, Self) {
let (re_lo, re_hi) = self.re.deinterleave(other.re);
let (im_lo, im_hi) = self.im.deinterleave(other.im);
(Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
}
}
struct ComplexReSplat<E, V>(PhantomData<(E, V)>);
struct ComplexImSplat<E, V>(PhantomData<(E, V)>);
impl<E, V: RealFloatVector> SplatConst<V::Element> for ComplexReSplat<E, V>
where
E: SplatConst<Complex<V::Element>>,
{
const VALUE: V::Element = <E as SplatConst<Complex<V::Element>>>::VALUE.re;
}
impl<E, V: RealFloatVector> SplatConst<V::Element> for ComplexImSplat<E, V>
where
E: SplatConst<Complex<V::Element>>,
{
const VALUE: V::Element = <E as SplatConst<Complex<V::Element>>>::VALUE.im;
}
impl<V: RealFloatVector> SplatVector<Complex<V::Element>> for Complex<V> {
type Splat<T: SplatConst<Complex<V::Element>>> = Self;
}
impl<V: RealFloatVector, E: SplatConst<Complex<V::Element>>> VectorValue<E, Complex<V>> for Complex<V> {
const VALUE: Complex<V> = Complex {
re: const_splat::<V, ComplexReSplat<E, V>>(),
im: const_splat::<V, ComplexImSplat<E, V>>(),
};
}
struct ComplexReNew<C, V>(PhantomData<(C, V)>);
struct ComplexImNew<C, V>(PhantomData<(C, V)>);
macro_rules! impl_new_const {
($($carrier:ident => $field:ident),* $(,)?) => {$(
impl<C, V: RealFloatVector> NewConst<V::Element, V::Lanes> for $carrier<C, V>
where
C: NewConst<Complex<V::Element>, V::Lanes>,
{
const VALUES: GenericArray<V::Element, V::Lanes> = const {
let c_vals = C::VALUES;
let src = c_vals.as_slice();
let mut out: GenericArray<V::Element, V::Lanes> = unsafe { core::mem::zeroed() };
let dst = out.as_mut_slice();
let mut i = 0;
while i < V::LANES {
dst[i] = src[i].$field;
i += 1;
}
core::mem::forget(c_vals);
out
};
}
)*};
}
impl_new_const!(ComplexReNew => re, ComplexImNew => im);
pub struct ComplexNewImpl;
impl<T, V: RealFloatVector> VectorValue<T, Complex<V>> for ComplexNewImpl
where
T: NewConst<Complex<V::Element>, V::Lanes>,
{
const VALUE: Complex<V> = Complex {
re: const_new::<V, V::Lanes, ComplexReNew<T, V>>(),
im: const_new::<V, V::Lanes, ComplexImNew<T, V>>(),
};
}
impl<V: RealFloatVector> NewVector<Complex<V::Element>, V::Lanes> for Complex<V> {
type New<T: NewConst<Complex<V::Element>, V::Lanes>> = ComplexNewImpl;
}
impl<FROM, TO> CastVector<Complex<FROM>> for Complex<TO>
where
FROM: RealFloatVector + CastVector<TO>,
TO: RealFloatVector + CastVector<FROM>,
{
#[inline(always)]
fn cast_into(self) -> Complex<FROM> {
Complex::<FROM>::cast_from(self)
}
#[inline(always)]
fn cast_from(from: Complex<FROM>) -> Self {
Complex::new(TO::cast_from(from.re), TO::cast_from(from.im))
}
}
#[rustfmt::skip]
impl<V: RealFloatVector> crate::math::specialized::ComplexVector for Complex<V> {
type Real = V;
#[inline(always)] fn re(self) -> V { self.re }
#[inline(always)] fn im(self) -> V { self.im }
#[inline(always)] fn from_parts(re: V, im: V) -> Self { Complex::new(re, im) }
#[inline(always)]
unsafe fn store_streaming_block(self, ptr: *mut Self) {
unsafe {
self.re.store_streaming((&raw mut (*ptr).re).cast());
self.im.store_streaming((&raw mut (*ptr).im).cast());
}
}
#[inline(always)] fn conj(self) -> Self { Complex::conj(self) }
#[inline(always)] fn norm_sqr(self) -> V { Complex::norm_sqr(self) }
#[inline(always)] fn inv(self) -> Self { Complex::inv(self) }
#[inline(always)] fn norm_l1(self) -> V { self.re.abs() + self.im.abs() }
}
#[inline(always)]
fn radix_per_component<V: RealFloatVector, const N: usize>(
inputs: [Complex<V>; N],
f: impl Fn([V; N]) -> [V; N],
) -> [Complex<V>; N] {
let (mut re, mut im) = ([V::EMPTY; N], [V::EMPTY; N]);
let mut i = 0;
while i < N {
re[i] = inputs[i].re;
im[i] = inputs[i].im;
i += 1;
}
let re = f(re);
let im = f(im);
let mut out = [Complex::<V>::EMPTY; N];
let mut i = 0;
while i < N {
out[i] = Complex::new(re[i], im[i]);
i += 1;
}
out
}
impl<V: RealFloatVector> Complex<V> {
#[inline(always)]
pub fn splat_parts(re: V::Element, im: V::Element) -> Self {
Self::new(V::splat(re), V::splat(im))
}
}
impl<V: RealFloatVector> GenericVector for Complex<V> {
type Element = Complex<V::Element>;
const EMPTY: Self = Self::ZERO;
const LANES: usize = V::LANES;
type Lanes = V::Lanes;
type Unsigned = V::Unsigned;
type Signed = V::Signed;
type Mask = V::Mask;
#[inline(always)]
fn new<const N: usize>(value: [Self::Element; N]) -> Self
where
Const<N>: IntoArrayLength<ArrayLength = Self::Lanes>,
{
let mut re = [<V::Element as Element>::ZERO; N];
let mut im = [<V::Element as Element>::ZERO; N];
let mut i = 0;
while i < N {
re[i] = value[i].re;
im[i] = value[i].im;
i += 1;
}
Complex::new(V::new(re), V::new(im))
}
#[inline(always)]
fn into_array(self) -> GenericArray<Self::Element, Self::Lanes> {
let mut arr = GenericArray::default();
for i in 0..Self::LANES {
arr[i] = Complex::new(self.re.extractv(i), self.im.extractv(i));
}
arr
}
#[inline(always)]
fn splat(value: Self::Element) -> Self {
Self::splat_parts(value.re, value.im)
}
#[inline(always)]
fn single(value: Self::Element) -> Self {
Complex::new(V::single(value.re), V::single(value.im))
}
#[inline(always)]
unsafe fn load(ptr: *const Self::Element) -> Self {
unsafe { Self::load_unaligned(ptr) }
}
#[inline(always)]
unsafe fn load_unaligned(ptr: *const Self::Element) -> Self {
let [out] = unsafe { Self::load_deinterleaved::<1>(ptr) };
out
}
#[inline(always)]
unsafe fn load_streaming(ptr: *const Self::Element) -> Self {
unsafe { Self::load(ptr) }
}
#[inline(always)]
fn interleave_by<const GROUP: usize>(self, other: Self) -> (Self, Self) {
let (re_lo, re_hi) = self.re.interleave_by::<GROUP>(other.re);
let (im_lo, im_hi) = self.im.interleave_by::<GROUP>(other.im);
(Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
}
#[inline(always)]
fn deinterleave_by<const GROUP: usize>(self, other: Self) -> (Self, Self) {
let (re_lo, re_hi) = self.re.deinterleave_by::<GROUP>(other.re);
let (im_lo, im_hi) = self.im.deinterleave_by::<GROUP>(other.im);
(Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
}
#[inline(always)]
fn interleave_radix<const N: usize>(inputs: [Self; N]) -> [Self; N] {
radix_per_component::<V, N>(inputs, V::interleave_radix::<N>)
}
#[inline(always)]
fn deinterleave_radix<const N: usize>(inputs: [Self; N]) -> [Self; N] {
radix_per_component::<V, N>(inputs, V::deinterleave_radix::<N>)
}
#[inline(always)]
fn deinterleave_radix_by<const N: usize, const GROUP: usize>(inputs: [Self; N]) -> [Self; N] {
radix_per_component::<V, N>(inputs, V::deinterleave_radix_by::<N, GROUP>)
}
#[inline(always)]
fn interleave_radix_by<const N: usize, const GROUP: usize>(inputs: [Self; N]) -> [Self; N] {
radix_per_component::<V, N>(inputs, V::interleave_radix_by::<N, GROUP>)
}
#[inline(always)]
unsafe fn load_deinterleaved<const M: usize>(ptr: *const Self::Element) -> [Self; M] {
let groups = unsafe { V::load_deinterleaved_grouped::<M, 1>(ptr as *const V::Element) };
let mut out = [Self::EMPTY; M];
let mut j = 0;
while j < M {
out[j] = Complex::new(groups[j].head, groups[j].tail[0]);
j += 1;
}
out
}
#[inline(always)]
unsafe fn store_interleaved<const M: usize>(ptr: *mut Self::Element, values: [Self; M]) {
let mut groups = [StreamGroup {
head: V::ZERO,
tail: [V::ZERO; 1],
}; M];
let mut j = 0;
while j < M {
groups[j] = StreamGroup {
head: values[j].re,
tail: [values[j].im],
};
j += 1;
}
unsafe { V::store_interleaved_grouped::<M, 1>(ptr as *mut V::Element, groups) }
}
#[inline(always)]
unsafe fn load_m(src: Self, mask: Self::Mask, ptr: *const Self::Element) -> Self {
let ptr = ptr as *const V::Element;
let (a_mask, b_mask) = mask.interleave(mask);
let (src_a, src_b) = src.re.interleave(src.im);
let a = unsafe { V::load_m(src_a, a_mask, ptr) };
let b = unsafe { V::load_m(src_b, b_mask, ptr.add(V::LANES)) };
let (re, im) = a.deinterleave(b);
Complex::new(re, im)
}
#[inline(always)]
unsafe fn load_z(mask: Self::Mask, ptr: *const Self::Element) -> Self {
unsafe { Self::load_m(Self::EMPTY, mask, ptr) }
}
#[inline(always)]
unsafe fn store(self, ptr: *mut Self::Element) {
unsafe { self.store_unaligned(ptr) }
}
#[inline(always)]
unsafe fn store_unaligned(self, ptr: *mut Self::Element) {
unsafe { Self::store_interleaved::<1>(ptr, [self]) }
}
#[inline(always)]
unsafe fn store_streaming(self, ptr: *mut Self::Element) {
unsafe { self.store(ptr) }
}
#[inline(always)]
unsafe fn store_masked(self, mask: Self::Mask, ptr: *mut Self::Element) {
let ptr = ptr as *mut V::Element;
let (a_mask, b_mask) = mask.interleave(mask);
let (a, b) = self.re.interleave(self.im);
unsafe {
a.store_masked(a_mask, ptr);
b.store_masked(b_mask, ptr.add(V::LANES));
}
}
#[inline(always)]
unsafe fn lookup_unchecked(values: &[Self::Element], indices: Self::Unsigned) -> Self {
let floats = unsafe { core::slice::from_raw_parts(values.as_ptr() as *const V::Element, values.len() * 2) };
let re_idx = indices << 1;
let im_idx = re_idx + Self::Unsigned::ONE;
let re = unsafe { V::lookup_unchecked(floats, re_idx) };
let im = unsafe { V::lookup_unchecked(floats, im_idx) };
Complex::new(re, im)
}
#[inline(always)]
fn broadcast<const I: usize>(self) -> Self {
Complex::new(V::broadcast::<I>(self.re), V::broadcast::<I>(self.im))
}
#[inline(always)]
fn broadcastv(self, idx: usize) -> Self {
Complex::new(self.re.broadcastv(idx), self.im.broadcastv(idx))
}
#[inline(always)]
fn extract<const I: usize>(self) -> Self::Element {
Complex::new(V::extract::<I>(self.re), V::extract::<I>(self.im))
}
#[inline(always)]
fn extractv(self, idx: usize) -> Self::Element {
Complex::new(self.re.extractv(idx), self.im.extractv(idx))
}
#[inline(always)]
fn insert<const I: usize>(self, value: Self::Element) -> Self {
Complex::new(V::insert::<I>(self.re, value.re), V::insert::<I>(self.im, value.im))
}
#[inline(always)]
fn insertv(self, idx: usize, value: Self::Element) -> Self {
Complex::new(self.re.insertv(idx, value.re), self.im.insertv(idx, value.im))
}
#[inline(always)]
fn reverse(self) -> Self {
Complex::new(self.re.reverse(), self.im.reverse())
}
#[inline(always)]
fn swap_bytes(self) -> Self {
Complex::new(self.re.swap_bytes(), self.im.swap_bytes())
}
#[inline(always)]
fn zz(self, mask: Self::Mask) -> Self {
Complex::new(self.re.zz(mask), self.im.zz(mask))
}
#[inline(always)]
fn nz(self, mask: Self::Mask) -> Self {
Complex::new(self.re.nz(mask), self.im.nz(mask))
}
#[inline(always)]
fn compress(self, mask: Self::Mask) -> Self {
Complex::new(self.re.compress(mask), self.im.compress(mask))
}
#[inline(always)]
fn compress_z(self, mask: Self::Mask) -> Self {
Complex::new(self.re.compress_z(mask), self.im.compress_z(mask))
}
#[inline(always)]
fn compress_m(self, src: Self, mask: Self::Mask) -> Self {
Complex::new(self.re.compress_m(src.re, mask), self.im.compress_m(src.im, mask))
}
#[inline(always)]
fn expand(self, mask: Self::Mask) -> Self {
Complex::new(self.re.expand(mask), self.im.expand(mask))
}
#[inline(always)]
fn expand_z(self, mask: Self::Mask) -> Self {
Complex::new(self.re.expand_z(mask), self.im.expand_z(mask))
}
#[inline(always)]
fn expand_m(self, src: Self, mask: Self::Mask) -> Self {
Complex::new(self.re.expand_m(src.re, mask), self.im.expand_m(src.im, mask))
}
#[inline(always)]
fn align<const OFFSET: usize>(self, other: Self) -> Self {
Complex::new(self.re.align::<OFFSET>(other.re), self.im.align::<OFFSET>(other.im))
}
const HAS_NATIVE_ALIGN: bool = V::HAS_NATIVE_ALIGN;
#[inline(always)]
fn map<F>(mut self, f: F) -> Self
where
F: Fn(Self::Element) -> Self::Element,
{
for i in 0..Self::LANES {
self = self.insertv(i, f(self.extractv(i)));
}
self
}
#[inline(always)]
fn fold<F>(self, mut init: Self::Element, f: F) -> Self::Element
where
F: Fn(Self::Element, Self::Element) -> Self::Element,
{
for i in 0..Self::LANES {
init = f(init, self.extractv(i));
}
init
}
#[inline(always)]
fn reduce<F>(self, f: F) -> Self::Element
where
F: Fn(Self::Element, Self::Element) -> Self::Element,
{
let mut result = self.extractv(0);
for i in 1..Self::LANES {
result = f(result, self.extractv(i));
}
result
}
#[rustfmt::skip]
#[inline(always)] fn splat_m(src: Self, mask: Self::Mask, value: Self::Element) -> Self { mask.select(Self::splat(value), src) }
#[rustfmt::skip]
#[inline(always)] fn splat_z(mask: Self::Mask, value: Self::Element) -> Self { mask.select(Self::splat(value), Self::EMPTY) }
#[rustfmt::skip]
#[inline(always)] fn broadcast_c<const I: usize>(self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), self) }
#[rustfmt::skip]
#[inline(always)] fn broadcast_m<const I: usize>(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), src) }
#[rustfmt::skip]
#[inline(always)] fn broadcast_z<const I: usize>(self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), Self::EMPTY) }
#[rustfmt::skip]
#[inline(always)] fn broadcastv_c(self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), self) }
#[rustfmt::skip]
#[inline(always)] fn broadcastv_m(self, src: Self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), src) }
#[rustfmt::skip]
#[inline(always)] fn broadcastv_z(self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), Self::EMPTY) }
#[rustfmt::skip]
#[inline(always)] fn reverse_c(self, mask: Self::Mask) -> Self { mask.select(self.reverse(), self) }
#[rustfmt::skip]
#[inline(always)] fn reverse_m(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.reverse(), src) }
#[rustfmt::skip]
#[inline(always)] fn reverse_z(self, mask: Self::Mask) -> Self { mask.select(self.reverse(), Self::EMPTY) }
#[rustfmt::skip]
#[inline(always)] fn swap_bytes_c(self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), self) }
#[rustfmt::skip]
#[inline(always)] fn swap_bytes_m(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), src) }
#[rustfmt::skip]
#[inline(always)] fn swap_bytes_z(self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), Self::EMPTY) }
}
#[rustfmt::skip]
impl<V: RealFloatVector> PartialOrdVector for Complex<V> {
#[inline(always)]
fn cmp_eq(self, other: Self) -> Self::Mask {
self.re.cmp_eq(other.re) & self.im.cmp_eq(other.im)
}
#[inline(always)]
fn cmp_ne(self, other: Self) -> Self::Mask {
self.re.cmp_ne(other.re) | self.im.cmp_ne(other.im)
}
#[inline(always)]
fn cmp_lt(self, other: Self) -> Self::Mask {
let re_lt = self.re.cmp_lt(other.re);
let re_eq = self.re.cmp_eq(other.re);
let im_lt = self.im.cmp_lt(other.im);
GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_lt, re_eq, im_lt)
}
#[inline(always)]
fn cmp_gt(self, other: Self) -> Self::Mask {
let re_gt = self.re.cmp_gt(other.re);
let re_eq = self.re.cmp_eq(other.re);
let im_gt = self.im.cmp_gt(other.im);
GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_gt, re_eq, im_gt)
}
#[inline(always)]
fn cmp_le(self, other: Self) -> Self::Mask {
let re_lt = self.re.cmp_lt(other.re);
let re_eq = self.re.cmp_eq(other.re);
let im_le = self.im.cmp_le(other.im);
GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_lt, re_eq, im_le)
}
#[inline(always)]
fn cmp_ge(self, other: Self) -> Self::Mask {
let re_gt = self.re.cmp_gt(other.re);
let re_eq = self.re.cmp_eq(other.re);
let im_ge = self.im.cmp_ge(other.im);
GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_gt, re_eq, im_ge)
}
}
macro_rules! impl_masked {
(MUL_ADD: $($method:ident),*) => {paste::paste! {
impl<V: RealFloatVector, A, B> thermite::vector::ops::MulAddExtMasked<V::Mask, A, B> for Complex<V>
where
Complex<V>: thermite::vector::ops::MulAddExt<A, B, Output = Self>,
{
$(
#[inline(always)]
fn [<$method _c>](self, mask: V::Mask, a: A, b: B) -> Self {
mask.select(self.$method(a, b), self)
}
#[inline(always)]
fn [<$method _m>](self, src: Self, mask: V::Mask, a: A, b: B) -> Self {
mask.select(self.$method(a, b), src)
}
#[inline(always)]
fn [<$method _z>](self, mask: V::Mask, a: A, b: B) -> Self {
mask.select(self.$method(a, b), Self::EMPTY)
}
)*
}
impl<V: RealFloatVector, A, B> thermite::vector::ops::MulAddAssignExtMasked<V::Mask, A, B> for Complex<V>
where
Complex<V>: thermite::vector::ops::MulAddExt<A, B, Output = Self>,
{
$(
#[inline(always)]
fn [<$method _assign_c>](&mut self, mask: V::Mask, a: A, b: B) {
*self = mask.select(self.$method(a, b), *self);
}
#[inline(always)]
fn [<$method _assign_m>](&mut self, src: Self, mask: V::Mask, a: A, b: B) {
*self = mask.select(self.$method(a, b), src);
}
#[inline(always)]
fn [<$method _assign_z>](&mut self, mask: V::Mask, a: A, b: B) {
*self = mask.select(self.$method(a, b), Self::EMPTY);
}
)*
}
}};
($trait:ident::$method:ident) => {paste::paste! {
impl<V: RealFloatVector, Rhs> thermite::vector::ops::[<$trait Masked>]<V::Mask, Rhs> for Complex<V>
where
Complex<V>: core::ops::$trait<Rhs, Output = Self>,
{
#[inline(always)]
fn [<$method _c>](self, mask: V::Mask, rhs: Rhs) -> Self {
mask.select(self.$method(rhs), self)
}
#[inline(always)]
fn [<$method _m>](self, src: Self, mask: V::Mask, rhs: Rhs) -> Self {
mask.select(self.$method(rhs), src)
}
#[inline(always)]
fn [<$method _z>](self, mask: V::Mask, rhs: Rhs) -> Self {
mask.select(self.$method(rhs), Self::EMPTY)
}
}
impl<V: RealFloatVector, Rhs> thermite::vector::ops::[<$trait AssignMasked>]<V::Mask, Rhs> for Complex<V>
where
Complex<V>: core::ops::$trait<Rhs, Output = Self>,
{
#[inline(always)]
fn [<$method _assign_c>](&mut self, mask: V::Mask, rhs: Rhs) {
*self = mask.select(self.$method(rhs), *self);
}
#[inline(always)]
fn [<$method _assign_m>](&mut self, src: Self, mask: V::Mask, rhs: Rhs) {
*self = mask.select(self.$method(rhs), src);
}
#[inline(always)]
fn [<$method _assign_z>](&mut self, mask: V::Mask, rhs: Rhs) {
*self = mask.select(self.$method(rhs), Self::EMPTY);
}
}
}};
}
impl_masked!(MUL_ADD: mul_add, mul_sub, nmul_add, nmul_sub, mul_adde, mul_sube, nmul_adde, nmul_sube);
impl_masked!(Add::add);
impl_masked!(Sub::sub);
impl_masked!(Mul::mul);
impl_masked!(Div::div);
impl_masked!(Rem::rem);
#[inline(always)]
fn neg_even_complex<V: RealFloatVector>(x: Complex<V>) -> Complex<V> {
Complex::new(V::ZERO.addsub(x.re), V::ZERO.addsub(x.im))
}
impl<V: RealFloatVector> AddSubExt for Complex<V> {
type Output = Self;
#[inline(always)]
fn addsub(self, b: Self) -> Self {
self + neg_even_complex(b)
}
#[inline(always)]
fn fmaddsub(self, b: Self, c: Self) -> Self {
self.mul_adde(b, neg_even_complex(c))
}
#[inline(always)]
fn fmsubadd(self, b: Self, c: Self) -> Self {
self.mul_sube(b, neg_even_complex(c))
}
}
impl<V: RealFloatVector> AddSubExtMasked<V::Mask> for Complex<V> {
#[inline(always)]
fn addsub_c(self, mask: V::Mask, b: Self) -> Self {
mask.select(self.addsub(b), self)
}
#[inline(always)]
fn addsub_m(self, src: Self, mask: V::Mask, b: Self) -> Self {
mask.select(self.addsub(b), src)
}
#[inline(always)]
fn addsub_z(self, mask: V::Mask, b: Self) -> Self {
mask.select(self.addsub(b), Self::EMPTY)
}
#[inline(always)]
fn fmaddsub_c(self, mask: V::Mask, b: Self, c: Self) -> Self {
mask.select(self.fmaddsub(b, c), self)
}
#[inline(always)]
fn fmaddsub_m(self, src: Self, mask: V::Mask, b: Self, c: Self) -> Self {
mask.select(self.fmaddsub(b, c), src)
}
#[inline(always)]
fn fmaddsub_z(self, mask: V::Mask, b: Self, c: Self) -> Self {
mask.select(self.fmaddsub(b, c), Self::EMPTY)
}
#[inline(always)]
fn fmsubadd_c(self, mask: V::Mask, b: Self, c: Self) -> Self {
mask.select(self.fmsubadd(b, c), self)
}
#[inline(always)]
fn fmsubadd_m(self, src: Self, mask: V::Mask, b: Self, c: Self) -> Self {
mask.select(self.fmsubadd(b, c), src)
}
#[inline(always)]
fn fmsubadd_z(self, mask: V::Mask, b: Self, c: Self) -> Self {
mask.select(self.fmsubadd(b, c), Self::EMPTY)
}
}
impl<V: RealFloatVector> SquareMasked<V::Mask> for Complex<V> {
#[inline(always)]
fn square_c(self, mask: V::Mask) -> Self::Output {
mask.select(self.square(), self)
}
#[inline(always)]
fn square_m(self, src: Self, mask: V::Mask) -> Self::Output {
mask.select(self.square(), src)
}
#[inline(always)]
fn square_z(self, mask: V::Mask) -> Self::Output {
mask.select(self.square(), Self::ZERO)
}
}
macro_rules! complex_masked {
(unary: $($m:ident),* $(,)?) => { paste::paste! {
$(
#[inline(always)] fn [<$m _c>](self, mask: Self::Mask) -> Self { mask.select(self.$m(), self) }
#[inline(always)] fn [<$m _m>](self, src: Self, mask: Self::Mask) -> Self { mask.select(self.$m(), src) }
#[inline(always)] fn [<$m _z>](self, mask: Self::Mask) -> Self { mask.select(self.$m(), Self::ZERO) }
)*
}};
(binary: $($m:ident),* $(,)?) => { paste::paste! {
$(
#[inline(always)] fn [<$m _c>](self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), self) }
#[inline(always)] fn [<$m _m>](self, src: Self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), src) }
#[inline(always)] fn [<$m _z>](self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), Self::ZERO) }
)*
}};
}
#[rustfmt::skip]
impl<V: RealFloatVector> Bounded for Complex<V> {
#[inline(always)] fn min_value() -> Self { Complex::new(V::MIN, V::MIN) }
#[inline(always)] fn max_value() -> Self { Complex::new(V::MAX, V::MAX) }
}
#[rustfmt::skip]
impl<V: RealFloatVector> thermite::sort::SortKey<Self> for Complex<V> {
#[inline(always)]
fn key_lt(a: Self, b: Self) -> V::Mask {
a.cmp_lt(b)
}
}
#[inline(always)]
fn sort_lanes_scalar<V: NumericVector, O: thermite::sort::SortOrder>(v: V) -> V
where
V::Element: PartialOrd,
{
let mut out = v;
let mut i = 1;
while i < V::LANES {
let key = out.extractv(i);
let mut j = i;
while j > 0 {
let prev = out.extractv(j - 1);
let misplaced = if O::IS_ASCENDING { prev > key } else { prev < key };
if !misplaced {
break;
}
out = out.insertv(j, prev);
j -= 1;
}
out = out.insertv(j, key);
i += 1;
}
out
}
impl<V: RealFloatVector> NumericVector for Complex<V> {
#[inline(always)]
fn to_signed_integer(self) -> Self::Signed {
self.re.to_signed_integer()
}
#[inline(always)]
fn from_signed_integer(v: Self::Signed) -> Self {
Self::real(V::from_signed_integer(v))
}
#[inline(always)]
fn to_unsigned_integer(self) -> Self::Unsigned {
self.re.to_unsigned_integer()
}
#[inline(always)]
fn from_unsigned_integer(v: Self::Unsigned) -> Self {
Self::real(V::from_unsigned_integer(v))
}
const ZERO: Self = Complex::new(V::ZERO, V::ZERO);
const ONE: Self = Complex::new(V::ONE, V::ZERO);
const TWO: Self = Self::real(V::TWO);
const MIN: Self = Complex::new(V::MIN, V::MIN);
const MAX: Self = Complex::new(V::MAX, V::MAX);
#[inline(always)]
fn is_zero(self) -> Self::Mask {
self.re.is_zero() & self.im.is_zero()
}
#[inline(always)]
fn is_all_zero(self) -> bool {
self.re.is_all_zero() && self.im.is_all_zero()
}
#[inline(always)]
fn sort_by<O: thermite::sort::SortOrder>(self) -> Self {
if const { Self::LANES <= 16 && Self::LANES.is_power_of_two() } {
thermite::sort::sort_lanes_by_key::<Self, O, Self>(self)
} else {
sort_lanes_scalar::<Self, O>(self)
}
}
#[inline(always)]
fn bitonic_clean_by<O: thermite::sort::SortOrder>(self) -> Self {
if const { Self::LANES <= 16 && Self::LANES.is_power_of_two() } {
thermite::sort::bitonic_clean_lanes_by_key::<Self, O, Self>(self)
} else {
sort_lanes_scalar::<Self, O>(self)
}
}
#[inline(always)]
fn min(self, other: Self) -> Self {
self.cmp_lt(other).select(self, other)
}
#[inline(always)]
fn max(self, other: Self) -> Self {
self.cmp_gt(other).select(self, other)
}
#[inline(always)]
fn clamp(self, min: Self, max: Self) -> Self {
let is_lt = self.cmp_lt(min);
let is_gt = self.cmp_gt(max);
Complex::new(
is_lt.select(min.re, is_gt.select(max.re, self.re)),
is_lt.select(min.im, is_gt.select(max.im, self.im)),
)
}
#[inline(always)]
fn min_element(self) -> Self::Element {
let mut arr = self.into_array();
reduce_in_place(&mut arr, |a, b| if b < a { b } else { a });
arr[0]
}
#[inline(always)]
fn max_element(self) -> Self::Element {
let mut arr = self.into_array();
reduce_in_place(&mut arr, |a, b| if b > a { b } else { a });
arr[0]
}
#[inline(always)]
fn min_max_element(self) -> (Self::Element, Self::Element) {
(self.min_element(), self.max_element())
}
#[inline(always)]
fn arg_minmax(self) -> (usize, usize) {
let arr = self.into_array();
let (mut lo, mut hi) = (0, 0);
for i in 1..Self::LANES {
if arr[i] < arr[lo] {
lo = i;
}
if arr[i] > arr[hi] {
hi = i;
}
}
(lo, hi)
}
#[inline(always)]
fn sum_elements(self) -> Self::Element {
Complex::new(self.re.sum_elements(), self.im.sum_elements())
}
#[inline(always)]
fn prefix_sum(self) -> Self {
Complex::new(self.re.prefix_sum(), self.im.prefix_sum())
}
#[inline(always)]
fn reverse_prefix_sum(self) -> Self {
Complex::new(self.re.reverse_prefix_sum(), self.im.reverse_prefix_sum())
}
#[inline(always)]
fn prefix_min(self) -> Self {
thermite::scan_ladder!(forward, self, self.broadcast::<0>(), Self::min)
}
#[inline(always)]
fn prefix_max(self) -> Self {
thermite::scan_ladder!(forward, self, self.broadcast::<0>(), Self::max)
}
#[inline(always)]
fn reverse_prefix_min(self) -> Self {
thermite::scan_ladder!(reverse, self, self.reverse().broadcast::<0>(), Self::min)
}
#[inline(always)]
fn reverse_prefix_max(self) -> Self {
thermite::scan_ladder!(reverse, self, self.reverse().broadcast::<0>(), Self::max)
}
#[inline(always)]
fn prod_elements(self) -> Self::Element {
let mut arr = self.into_array();
reduce_in_place(&mut arr, |a, b| a * b);
arr[0]
}
#[inline(always)]
fn offset() -> Self {
Self::real(V::offset())
}
#[inline(always)]
fn indexed() -> Self {
Self::real(V::indexed())
}
#[inline(always)]
fn scale(self, factor: Self::Element) -> Self {
self * Self::splat(factor)
}
#[inline(always)]
fn scale_c(self, mask: Self::Mask, factor: Self::Element) -> Self {
mask.select(<Self as NumericVector>::scale(self, factor), self)
}
#[inline(always)]
fn scale_m(self, src: Self, mask: Self::Mask, factor: Self::Element) -> Self {
mask.select(<Self as NumericVector>::scale(self, factor), src)
}
#[inline(always)]
fn scale_z(self, mask: Self::Mask, factor: Self::Element) -> Self {
mask.select(<Self as NumericVector>::scale(self, factor), Self::ZERO)
}
complex_masked!(binary: min, max);
#[inline(always)]
fn pairwise_sum(lo: Self, hi: Self) -> Self {
Complex::new(V::pairwise_sum(lo.re, hi.re), V::pairwise_sum(lo.im, hi.im))
}
#[inline(always)]
fn relaxed_pairwise_sum(lo: Self, hi: Self) -> Self {
Complex::new(
V::relaxed_pairwise_sum(lo.re, hi.re),
V::relaxed_pairwise_sum(lo.im, hi.im),
)
}
}
impl<V: RealFloatVector> NegMasked<V::Mask> for Complex<V> {
#[inline(always)]
fn neg_c(self, mask: V::Mask) -> Self {
Complex::new(self.re.neg_c(mask), self.im.neg_c(mask))
}
#[inline(always)]
fn neg_m(self, src: Self, mask: V::Mask) -> Self {
Complex::new(self.re.neg_m(src.re, mask), self.im.neg_m(src.im, mask))
}
#[inline(always)]
fn neg_z(self, mask: V::Mask) -> Self {
Complex::new(self.re.neg_z(mask), self.im.neg_z(mask))
}
}
impl<V: RealFloatVector> Complex<V> {
#[inline(always)]
pub(crate) fn modulus(self) -> V {
self.re.hypot_p::<DefaultPolicy>(self.im)
}
}
#[rustfmt::skip]
impl<V: RealFloatVector> SignedVector for Complex<V> {
const NEG_ONE: Self = Self::real(V::NEG_ONE);
const MIN_POSITIVE: Self = Self::real(V::MIN_POSITIVE);
#[inline(always)]
fn abs(self) -> Self {
Self::real(self.modulus())
}
#[inline(always)]
fn signum(self) -> Self {
let m = self.modulus();
let is_zero = m.is_zero();
let inv = m.reciprocal_p::<DefaultPolicy>();
Complex::new((self.re * inv).nz(is_zero), (self.im * inv).nz(is_zero))
}
#[inline(always)] fn is_positive(self) -> Self::Mask { self.re.is_positive() }
#[inline(always)] fn is_negative(self) -> Self::Mask { self.re.is_negative() }
#[inline(always)] fn select_negative(self, if_neg: Self, if_pos: Self) -> Self { self.is_negative().select(if_neg, if_pos) }
#[inline(always)]
fn copysign(self, sign: Self) -> Self {
Complex::new(self.re.copysign(sign.re), self.im.copysign(sign.im))
}
complex_masked!(unary: abs);
complex_masked!(binary: copysign);
}
#[inline(always)]
fn sqrt_kahan<V: RealFloatVector>(z: Complex<V>) -> (Complex<V>, V::Mask) {
let half = <V as FloatVector>::HALF;
let m = z.modulus();
let t = z.re.abs().mul_adde(half, m * half).sqrt(); let half_im = z.im * half;
let q = half_im / t;
let re_pos = z.re.is_positive();
let degenerate = if const { features::PRESERVE_DENORMALS || features::IGNORE_DENORMALS } {
t.is_zero()
} else {
m.is_zero()
};
(
Complex::new(re_pos.select(t, q.abs()), re_pos.select(q, t.mul_sign(z.im))),
degenerate,
)
}
#[inline(always)]
fn sqrt_rescaled<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
let up = V::ONE / <V as SignedVector>::MIN_POSITIVE;
let down = <V as SignedVector>::MIN_POSITIVE.sqrt();
let (r, _) = sqrt_kahan(Complex::new(z.re * up, z.im * up));
let is_zero = z.is_zero();
Complex::new((r.re * down).nz(is_zero), (r.im * down).nz(is_zero))
}
#[rustfmt::skip]
impl<V: RealFloatVector> FloatVector for Complex<V> {
const HALF: Self = Self::real(<V as FloatVector>::HALF);
const NEG_ZERO: Self = Self::real(<V as FloatVector>::NEG_ZERO);
const EPSILON: Self = Self::real(<V as FloatVector>::EPSILON);
const INFINITY: Self = Self::real(<V as FloatVector>::INFINITY);
const NEG_INFINITY: Self = Self::real(<V as FloatVector>::NEG_INFINITY);
const NAN: Self = Complex::new(<V as FloatVector>::NAN, <V as FloatVector>::NAN);
type ExtendedPrecision = Self;
const HAS_APPROX_RCP: bool = V::HAS_APPROX_RCP;
const HAS_APPROX_RSQRT: bool = V::HAS_APPROX_RCP;
#[inline(always)] fn is_nan(self) -> Self::Mask { self.re.is_nan() | self.im.is_nan() }
#[inline(always)] fn is_infinite(self) -> Self::Mask { self.re.is_infinite() | self.im.is_infinite() }
#[inline(always)] fn is_finite(self) -> Self::Mask { self.re.is_finite() & self.im.is_finite() }
#[inline(always)] fn is_zero_or_subnormal(self) -> Self::Mask { self.re.is_zero_or_subnormal() & self.im.is_zero_or_subnormal() }
#[inline(always)] fn is_normal(self) -> Self::Mask { (self.re.is_normal() | self.im.is_normal()) & self.is_finite() }
#[inline(always)] fn is_subnormal(self) -> Self::Mask { (self.re.is_subnormal() | self.im.is_subnormal()) & self.is_zero_or_subnormal() }
#[inline(always)]
fn sqrt(self) -> Self {
let (res, degenerate) = sqrt_kahan(self);
if thermite::unlikely(degenerate.any()) {
return degenerate.select(sqrt_rescaled(self), res);
}
res
}
#[inline(always)]
fn rcp(self) -> Self {
let inv = self.norm_sqr().rcp();
Complex::new(self.re * inv, -(self.im * inv))
}
#[inline(always)]
fn rsqrt(self) -> Self {
let s = self.sqrt();
let inv = self.modulus().rcp();
Complex::new(s.re * inv, -(s.im * inv))
}
#[inline(always)] fn floor(self) -> Self { Complex::new(self.re.floor(), self.im.floor()) }
#[inline(always)] fn ceil(self) -> Self { Complex::new(self.re.ceil(), self.im.ceil()) }
#[inline(always)] fn round(self) -> Self { Complex::new(self.re.round(), self.im.round()) }
#[inline(always)] fn trunc(self) -> Self { Complex::new(self.re.trunc(), self.im.trunc()) }
#[inline(always)] fn fract(self) -> Self { Complex::new(self.re.fract(), self.im.fract()) }
#[inline(always)] fn mul_sign(self, sign: Self) -> Self { Complex::new(self.re.mul_sign(sign.re), self.im.mul_sign(sign.im)) }
#[inline(always)] fn signed_zero(self) -> Self { Complex::new(self.re.signed_zero(), self.im.signed_zero()) }
#[inline(always)] fn next_up(self) -> Self { Complex::new(self.re.next_up(), self.im.next_up()) }
#[inline(always)] fn next_down(self) -> Self { Complex::new(self.re.next_down(), self.im.next_down()) }
#[inline(always)]
unsafe fn block_autovectorization(&mut self) {
unsafe {
self.re.block_autovectorization();
self.im.block_autovectorization();
}
}
#[inline(always)] fn mix(self, a: Self, b: Self) -> Self { (b - a).mul_adde(self, a) }
complex_masked!(unary: sqrt, rsqrt, rcp, floor, ceil, round, trunc, fract, signed_zero, next_up, next_down);
complex_masked!(binary: mul_sign);
}