1use core::marker::PhantomData;
12use core::ops::{Add, Div, Mul, Rem, Sub};
13
14use num_traits::Bounded;
15
16use thermite::Swizzle;
17use thermite::element::{Element, FloatElement, SignedElement};
18use thermite::generic_array::{GenericArray, IntoArrayLength, typenum::Const};
19use thermite::mask::{GenericMask, GenericSelectable};
20use thermite::math::RealMathWithPolicy;
21use thermite::math::algorithms::reduce_in_place;
22use thermite::math::policy::DefaultPolicy;
23use thermite::register::SwizzleIndices;
24use thermite::vector::ops::{AddSubExt, AddSubExtMasked, NegMasked, Square, SquareMasked};
25use thermite::vector::{NewConst, NewVector, SplatConst, SplatVector, VectorValue, const_new, const_splat};
26use thermite::{LargeInt, features, prelude::*};
27
28use crate::{Complex, RealValue};
29
30pub trait RealFloatVector:
37 RealValue + FloatVector<Element: RealValue> + CastVector<Self> + RealMathWithPolicy + SwizzleVector
38{
39}
40
41impl<V> RealFloatVector for V where
42 V: RealValue + FloatVector<Element: RealValue> + CastVector<V> + RealMathWithPolicy + SwizzleVector
43{
44}
45
46impl<V: RealFloatVector> Swizzle<V::Lanes> for Complex<V> {
50 #[inline(always)]
51 fn swizzle(self, other: Self, indices: GenericArray<u32, V::Lanes>) -> Self {
52 Self {
53 re: self.re.swizzle(other.re, indices.clone()),
54 im: self.im.swizzle(other.im, indices),
55 }
56 }
57
58 #[inline(always)]
59 fn permute(self, indices: GenericArray<u32, V::Lanes>) -> Self {
60 Self {
61 re: self.re.permute(indices.clone()),
62 im: self.im.permute(indices),
63 }
64 }
65
66 #[inline(always)]
70 fn swizzle_const<I: SwizzleIndices<V::Lanes>>(self, other: Self) -> Self {
71 Self {
72 re: self.re.swizzle_const::<I>(other.re),
73 im: self.im.swizzle_const::<I>(other.im),
74 }
75 }
76
77 #[inline(always)]
78 fn permute_const<I: SwizzleIndices<V::Lanes>>(self) -> Self {
79 Self {
80 re: self.re.permute_const::<I>(),
81 im: self.im.permute_const::<I>(),
82 }
83 }
84}
85
86#[rustfmt::skip]
89impl<E: RealValue + Element> Element for Complex<E> {
90 type Signed = <E as Element>::Signed;
91 type Unsigned = <E as Element>::Unsigned;
92
93 const ZERO: Self = Self::ZERO;
94 const ONE: Self = Self::ONE;
95
96 const ORDER_MAX: Self = Self { re: E::ORDER_MAX, im: E::ORDER_MAX };
101 const ORDER_MIN: Self = Self { re: E::ORDER_MIN, im: E::ORDER_MIN };
102 const HAS_UNORDERED: bool = E::HAS_UNORDERED;
103 const IS_FLOAT: bool = E::IS_FLOAT;
104
105 #[inline(always)] fn from_i8(value: i8) -> Self { Self::real(E::from_i8(value)) }
106 #[inline(always)] fn from_u8(value: u8) -> Self { Self::real(E::from_u8(value)) }
107 #[inline(always)] fn from_u16(value: u16) -> Self { Self::real(E::from_u16(value)) }
108}
109
110impl<E: RealValue + FloatElement> Complex<E> {
111 #[inline(always)]
114 fn elem_modulus(self) -> E {
115 E::sqrt(self.re.mul_add(self.re, self.im * self.im))
116 }
117}
118
119impl<E: RealValue + FloatElement> SignedElement for Complex<E> {
120 #[inline(always)]
122 fn abs(self) -> Self {
123 Self::real(self.elem_modulus())
124 }
125
126 #[inline(always)]
128 fn signum(self) -> Self {
129 let m = self.elem_modulus();
130
131 if m == E::ZERO {
132 return Self::ZERO;
133 }
134
135 Self::new(self.re / m, self.im / m)
136 }
137}
138
139pub struct ComplexIntConst<E, const VAL: LargeInt>(PhantomData<E>);
141
142pub struct ComplexRatioConst<E, const NUM: LargeInt, const DEN: LargeInt>(PhantomData<E>);
144
145impl<E: RealValue + FloatElement, const VAL: LargeInt> SplatConst<Complex<E>> for ComplexIntConst<E, VAL> {
146 const VALUE: Complex<E> = Complex::real(<E::ConstInt<VAL> as SplatConst<E>>::VALUE);
147}
148
149impl<E: RealValue + FloatElement, const NUM: LargeInt, const DEN: LargeInt> SplatConst<Complex<E>>
150 for ComplexRatioConst<E, NUM, DEN>
151{
152 const VALUE: Complex<E> = Complex::real(<E::ConstRatio<NUM, DEN> as SplatConst<E>>::VALUE);
153}
154
155#[rustfmt::skip]
156impl<E: RealValue + FloatElement> FloatElement for Complex<E> {
157 #[inline(always)]
160 fn sqrt(this: Self) -> Self {
161 let half = E::from_ratio(1, 2);
162
163 let t = E::sqrt((SignedElement::abs(this.re) + this.elem_modulus()) * half);
164
165 if t == E::ZERO {
166 return Self::ZERO;
167 }
168
169 let half_im = this.im * half;
170
171 if this.re >= E::ZERO {
172 Self::new(t, half_im / t)
173 } else {
174 let i = if this.im < E::ZERO { -t } else { t };
175
176 Self::new(SignedElement::abs(half_im) / t, i)
177 }
178 }
179
180 #[inline(always)] fn floor(this: Self) -> Self { Self::new(E::floor(this.re), E::floor(this.im)) }
182 #[inline(always)] fn ceil(this: Self) -> Self { Self::new(E::ceil(this.re), E::ceil(this.im)) }
183 #[inline(always)] fn round(this: Self) -> Self { Self::new(E::round(this.re), E::round(this.im)) }
184 #[inline(always)] fn trunc(this: Self) -> Self { Self::new(E::trunc(this.re), E::trunc(this.im)) }
185
186 #[inline(always)] fn next_up(this: Self) -> Self { Self::new(E::next_up(this.re), E::next_up(this.im)) }
187 #[inline(always)] fn next_down(this: Self) -> Self { Self::new(E::next_down(this.re), E::next_down(this.im)) }
188
189 #[inline(always)]
190 fn try_from_int(value: LargeInt) -> Option<Self> {
191 E::try_from_int(value).map(Self::real)
192 }
193
194 #[inline(always)]
195 fn try_from_ratio(n: LargeInt, d: LargeInt) -> Option<Self> {
196 E::try_from_ratio(n, d).map(Self::real)
197 }
198
199 const HAS_INFINITY: bool = E::HAS_INFINITY;
200 const HAS_SIGNED_ZERO: bool = E::HAS_SIGNED_ZERO;
201 const HAS_SUBNORMALS: bool = E::HAS_SUBNORMALS;
202
203 type ConstInt<const VAL: LargeInt> = ComplexIntConst<E, VAL>;
204 type ConstRatio<const NUM: LargeInt, const DEN: LargeInt> = ComplexRatioConst<E, NUM, DEN>;
205}
206
207impl<V: thermite::simd::HasIsa> thermite::simd::HasIsa for Complex<V> {
210 type Native = V::Native;
211
212 const ISA: thermite::isa::InstructionSet = V::ISA;
213}
214
215impl<V: RealFloatVector> GenericSelectable for Complex<V> {
216 type SelectableMask = <V as GenericSelectable>::SelectableMask;
217
218 #[inline(always)]
219 fn select<M>(mask: M, t: Self, f: Self) -> Self
220 where
221 Self::SelectableMask: CastMask<M>,
222 {
223 let mask = <Self::SelectableMask as CastMask<M>>::mask_from(mask);
224
225 Self::new(mask.select(t.re, f.re), mask.select(t.im, f.im))
226 }
227}
228
229#[rustfmt::skip]
230impl<V: RealFloatVector> Interleave for Complex<V> {
231 #[inline(always)]
232 fn interleave(self, other: Self) -> (Self, Self) {
233 let (re_lo, re_hi) = self.re.interleave(other.re);
234 let (im_lo, im_hi) = self.im.interleave(other.im);
235
236 (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
237 }
238
239 #[inline(always)]
240 fn deinterleave(self, other: Self) -> (Self, Self) {
241 let (re_lo, re_hi) = self.re.deinterleave(other.re);
242 let (im_lo, im_hi) = self.im.deinterleave(other.im);
243
244 (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
245 }
246}
247
248struct ComplexReSplat<E, V>(PhantomData<(E, V)>);
252struct ComplexImSplat<E, V>(PhantomData<(E, V)>);
253
254impl<E, V: RealFloatVector> SplatConst<V::Element> for ComplexReSplat<E, V>
255where
256 E: SplatConst<Complex<V::Element>>,
257{
258 const VALUE: V::Element = <E as SplatConst<Complex<V::Element>>>::VALUE.re;
259}
260
261impl<E, V: RealFloatVector> SplatConst<V::Element> for ComplexImSplat<E, V>
262where
263 E: SplatConst<Complex<V::Element>>,
264{
265 const VALUE: V::Element = <E as SplatConst<Complex<V::Element>>>::VALUE.im;
266}
267
268impl<V: RealFloatVector> SplatVector<Complex<V::Element>> for Complex<V> {
269 type Splat<T: SplatConst<Complex<V::Element>>> = Self;
270}
271
272impl<V: RealFloatVector, E: SplatConst<Complex<V::Element>>> VectorValue<E, Complex<V>> for Complex<V> {
273 const VALUE: Complex<V> = Complex {
274 re: const_splat::<V, ComplexReSplat<E, V>>(),
275 im: const_splat::<V, ComplexImSplat<E, V>>(),
276 };
277}
278
279struct ComplexReNew<C, V>(PhantomData<(C, V)>);
281struct ComplexImNew<C, V>(PhantomData<(C, V)>);
282
283macro_rules! impl_new_const {
284 ($($carrier:ident => $field:ident),* $(,)?) => {$(
285 impl<C, V: RealFloatVector> NewConst<V::Element, V::Lanes> for $carrier<C, V>
286 where
287 C: NewConst<Complex<V::Element>, V::Lanes>,
288 {
289 const VALUES: GenericArray<V::Element, V::Lanes> = const {
290 let c_vals = C::VALUES;
291 let src = c_vals.as_slice();
292 let mut out: GenericArray<V::Element, V::Lanes> = unsafe { core::mem::zeroed() };
293 let dst = out.as_mut_slice();
294 let mut i = 0;
295 while i < V::LANES {
296 dst[i] = src[i].$field;
297 i += 1;
298 }
299 core::mem::forget(c_vals);
300 out
301 };
302 }
303 )*};
304}
305
306impl_new_const!(ComplexReNew => re, ComplexImNew => im);
307
308pub struct ComplexNewImpl;
310
311impl<T, V: RealFloatVector> VectorValue<T, Complex<V>> for ComplexNewImpl
312where
313 T: NewConst<Complex<V::Element>, V::Lanes>,
314{
315 const VALUE: Complex<V> = Complex {
316 re: const_new::<V, V::Lanes, ComplexReNew<T, V>>(),
317 im: const_new::<V, V::Lanes, ComplexImNew<T, V>>(),
318 };
319}
320
321impl<V: RealFloatVector> NewVector<Complex<V::Element>, V::Lanes> for Complex<V> {
322 type New<T: NewConst<Complex<V::Element>, V::Lanes>> = ComplexNewImpl;
323}
324
325impl<FROM, TO> CastVector<Complex<FROM>> for Complex<TO>
328where
329 FROM: RealFloatVector + CastVector<TO>,
330 TO: RealFloatVector + CastVector<FROM>,
331{
332 #[inline(always)]
333 fn cast_into(self) -> Complex<FROM> {
334 Complex::<FROM>::cast_from(self)
335 }
336
337 #[inline(always)]
338 fn cast_from(from: Complex<FROM>) -> Self {
339 Complex::new(TO::cast_from(from.re), TO::cast_from(from.im))
340 }
341}
342
343#[rustfmt::skip]
346impl<V: RealFloatVector> crate::math::specialized::ComplexVector for Complex<V> {
347 type Real = V;
348
349 #[inline(always)] fn re(self) -> V { self.re }
350 #[inline(always)] fn im(self) -> V { self.im }
351 #[inline(always)] fn from_parts(re: V, im: V) -> Self { Complex::new(re, im) }
352
353 #[inline(always)]
354 unsafe fn store_streaming_block(self, ptr: *mut Self) {
355 unsafe {
361 self.re.store_streaming((&raw mut (*ptr).re).cast());
362 self.im.store_streaming((&raw mut (*ptr).im).cast());
363 }
364 }
365
366 #[inline(always)] fn conj(self) -> Self { Complex::conj(self) }
369 #[inline(always)] fn norm_sqr(self) -> V { Complex::norm_sqr(self) }
370 #[inline(always)] fn inv(self) -> Self { Complex::inv(self) }
371
372 #[inline(always)] fn norm_l1(self) -> V { self.re.abs() + self.im.abs() }
373}
374
375#[inline(always)]
384fn radix_per_component<V: RealFloatVector, const N: usize>(
385 inputs: [Complex<V>; N],
386 f: impl Fn([V; N]) -> [V; N],
387) -> [Complex<V>; N] {
388 let (mut re, mut im) = ([V::EMPTY; N], [V::EMPTY; N]);
389
390 let mut i = 0;
391 while i < N {
392 re[i] = inputs[i].re;
393 im[i] = inputs[i].im;
394 i += 1;
395 }
396
397 let re = f(re);
398 let im = f(im);
399
400 let mut out = [Complex::<V>::EMPTY; N];
401
402 let mut i = 0;
403 while i < N {
404 out[i] = Complex::new(re[i], im[i]);
405 i += 1;
406 }
407
408 out
409}
410
411impl<V: RealFloatVector> Complex<V> {
412 #[inline(always)]
414 pub fn splat_parts(re: V::Element, im: V::Element) -> Self {
415 Self::new(V::splat(re), V::splat(im))
416 }
417}
418
419impl<V: RealFloatVector> GenericVector for Complex<V> {
420 type Element = Complex<V::Element>;
421
422 const EMPTY: Self = Self::ZERO;
423 const LANES: usize = V::LANES;
424
425 type Lanes = V::Lanes;
426
427 type Unsigned = V::Unsigned;
428 type Signed = V::Signed;
429 type Mask = V::Mask;
430
431 #[inline(always)]
432 fn new<const N: usize>(value: [Self::Element; N]) -> Self
433 where
434 Const<N>: IntoArrayLength<ArrayLength = Self::Lanes>,
435 {
436 let mut re = [<V::Element as Element>::ZERO; N];
437 let mut im = [<V::Element as Element>::ZERO; N];
438
439 let mut i = 0;
440 while i < N {
441 re[i] = value[i].re;
442 im[i] = value[i].im;
443 i += 1;
444 }
445
446 Complex::new(V::new(re), V::new(im))
447 }
448
449 #[inline(always)]
450 fn into_array(self) -> GenericArray<Self::Element, Self::Lanes> {
451 let mut arr = GenericArray::default();
452
453 for i in 0..Self::LANES {
454 arr[i] = Complex::new(self.re.extractv(i), self.im.extractv(i));
455 }
456
457 arr
458 }
459
460 #[inline(always)]
461 fn splat(value: Self::Element) -> Self {
462 Self::splat_parts(value.re, value.im)
463 }
464
465 #[inline(always)]
466 fn single(value: Self::Element) -> Self {
467 Complex::new(V::single(value.re), V::single(value.im))
468 }
469
470 #[inline(always)]
473 unsafe fn load(ptr: *const Self::Element) -> Self {
474 unsafe { Self::load_unaligned(ptr) }
475 }
476
477 #[inline(always)]
481 unsafe fn load_unaligned(ptr: *const Self::Element) -> Self {
482 let [out] = unsafe { Self::load_deinterleaved::<1>(ptr) };
483 out
484 }
485
486 #[inline(always)]
487 unsafe fn load_streaming(ptr: *const Self::Element) -> Self {
488 unsafe { Self::load(ptr) }
489 }
490
491 #[inline(always)]
492 fn interleave_by<const GROUP: usize>(self, other: Self) -> (Self, Self) {
493 let (re_lo, re_hi) = self.re.interleave_by::<GROUP>(other.re);
494 let (im_lo, im_hi) = self.im.interleave_by::<GROUP>(other.im);
495 (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
496 }
497
498 #[inline(always)]
499 fn deinterleave_by<const GROUP: usize>(self, other: Self) -> (Self, Self) {
500 let (re_lo, re_hi) = self.re.deinterleave_by::<GROUP>(other.re);
501 let (im_lo, im_hi) = self.im.deinterleave_by::<GROUP>(other.im);
502 (Self::new(re_lo, im_lo), Self::new(re_hi, im_hi))
503 }
504
505 #[inline(always)]
506 fn interleave_radix<const N: usize>(inputs: [Self; N]) -> [Self; N] {
507 radix_per_component::<V, N>(inputs, V::interleave_radix::<N>)
508 }
509
510 #[inline(always)]
511 fn deinterleave_radix<const N: usize>(inputs: [Self; N]) -> [Self; N] {
512 radix_per_component::<V, N>(inputs, V::deinterleave_radix::<N>)
513 }
514
515 #[inline(always)]
516 fn deinterleave_radix_by<const N: usize, const GROUP: usize>(inputs: [Self; N]) -> [Self; N] {
517 radix_per_component::<V, N>(inputs, V::deinterleave_radix_by::<N, GROUP>)
518 }
519
520 #[inline(always)]
521 fn interleave_radix_by<const N: usize, const GROUP: usize>(inputs: [Self; N]) -> [Self; N] {
522 radix_per_component::<V, N>(inputs, V::interleave_radix_by::<N, GROUP>)
523 }
524
525 #[inline(always)]
530 unsafe fn load_deinterleaved<const M: usize>(ptr: *const Self::Element) -> [Self; M] {
531 let groups = unsafe { V::load_deinterleaved_grouped::<M, 1>(ptr as *const V::Element) };
532
533 let mut out = [Self::EMPTY; M];
534 let mut j = 0;
535 while j < M {
536 out[j] = Complex::new(groups[j].head, groups[j].tail[0]);
537 j += 1;
538 }
539 out
540 }
541
542 #[inline(always)]
544 unsafe fn store_interleaved<const M: usize>(ptr: *mut Self::Element, values: [Self; M]) {
545 let mut groups = [StreamGroup {
546 head: V::ZERO,
547 tail: [V::ZERO; 1],
548 }; M];
549
550 let mut j = 0;
551 while j < M {
552 groups[j] = StreamGroup {
553 head: values[j].re,
554 tail: [values[j].im],
555 };
556 j += 1;
557 }
558
559 unsafe { V::store_interleaved_grouped::<M, 1>(ptr as *mut V::Element, groups) }
560 }
561
562 #[inline(always)]
563 unsafe fn load_m(src: Self, mask: Self::Mask, ptr: *const Self::Element) -> Self {
564 let ptr = ptr as *const V::Element;
565
566 let (a_mask, b_mask) = mask.interleave(mask);
569 let (src_a, src_b) = src.re.interleave(src.im);
570
571 let a = unsafe { V::load_m(src_a, a_mask, ptr) };
572 let b = unsafe { V::load_m(src_b, b_mask, ptr.add(V::LANES)) };
573
574 let (re, im) = a.deinterleave(b);
575 Complex::new(re, im)
576 }
577
578 #[inline(always)]
579 unsafe fn load_z(mask: Self::Mask, ptr: *const Self::Element) -> Self {
580 unsafe { Self::load_m(Self::EMPTY, mask, ptr) }
581 }
582
583 #[inline(always)]
585 unsafe fn store(self, ptr: *mut Self::Element) {
586 unsafe { self.store_unaligned(ptr) }
587 }
588
589 #[inline(always)]
590 unsafe fn store_unaligned(self, ptr: *mut Self::Element) {
591 unsafe { Self::store_interleaved::<1>(ptr, [self]) }
592 }
593
594 #[inline(always)]
595 unsafe fn store_streaming(self, ptr: *mut Self::Element) {
596 unsafe { self.store(ptr) }
597 }
598
599 #[inline(always)]
600 unsafe fn store_masked(self, mask: Self::Mask, ptr: *mut Self::Element) {
601 let ptr = ptr as *mut V::Element;
602
603 let (a_mask, b_mask) = mask.interleave(mask);
604 let (a, b) = self.re.interleave(self.im);
605
606 unsafe {
607 a.store_masked(a_mask, ptr);
608 b.store_masked(b_mask, ptr.add(V::LANES));
609 }
610 }
611
612 #[inline(always)]
613 unsafe fn lookup_unchecked(values: &[Self::Element], indices: Self::Unsigned) -> Self {
614 let floats = unsafe { core::slice::from_raw_parts(values.as_ptr() as *const V::Element, values.len() * 2) };
618
619 let re_idx = indices << 1;
620 let im_idx = re_idx + Self::Unsigned::ONE;
621
622 let re = unsafe { V::lookup_unchecked(floats, re_idx) };
623 let im = unsafe { V::lookup_unchecked(floats, im_idx) };
624
625 Complex::new(re, im)
626 }
627
628 #[inline(always)]
629 fn broadcast<const I: usize>(self) -> Self {
630 Complex::new(V::broadcast::<I>(self.re), V::broadcast::<I>(self.im))
631 }
632
633 #[inline(always)]
634 fn broadcastv(self, idx: usize) -> Self {
635 Complex::new(self.re.broadcastv(idx), self.im.broadcastv(idx))
636 }
637
638 #[inline(always)]
639 fn extract<const I: usize>(self) -> Self::Element {
640 Complex::new(V::extract::<I>(self.re), V::extract::<I>(self.im))
641 }
642
643 #[inline(always)]
644 fn extractv(self, idx: usize) -> Self::Element {
645 Complex::new(self.re.extractv(idx), self.im.extractv(idx))
646 }
647
648 #[inline(always)]
649 fn insert<const I: usize>(self, value: Self::Element) -> Self {
650 Complex::new(V::insert::<I>(self.re, value.re), V::insert::<I>(self.im, value.im))
651 }
652
653 #[inline(always)]
654 fn insertv(self, idx: usize, value: Self::Element) -> Self {
655 Complex::new(self.re.insertv(idx, value.re), self.im.insertv(idx, value.im))
656 }
657
658 #[inline(always)]
659 fn reverse(self) -> Self {
660 Complex::new(self.re.reverse(), self.im.reverse())
661 }
662
663 #[inline(always)]
664 fn swap_bytes(self) -> Self {
665 Complex::new(self.re.swap_bytes(), self.im.swap_bytes())
666 }
667
668 #[inline(always)]
669 fn zz(self, mask: Self::Mask) -> Self {
670 Complex::new(self.re.zz(mask), self.im.zz(mask))
671 }
672
673 #[inline(always)]
674 fn nz(self, mask: Self::Mask) -> Self {
675 Complex::new(self.re.nz(mask), self.im.nz(mask))
676 }
677
678 #[inline(always)]
679 fn compress(self, mask: Self::Mask) -> Self {
680 Complex::new(self.re.compress(mask), self.im.compress(mask))
681 }
682
683 #[inline(always)]
684 fn compress_z(self, mask: Self::Mask) -> Self {
685 Complex::new(self.re.compress_z(mask), self.im.compress_z(mask))
686 }
687
688 #[inline(always)]
693 fn compress_m(self, src: Self, mask: Self::Mask) -> Self {
694 Complex::new(self.re.compress_m(src.re, mask), self.im.compress_m(src.im, mask))
695 }
696
697 #[inline(always)]
698 fn expand(self, mask: Self::Mask) -> Self {
699 Complex::new(self.re.expand(mask), self.im.expand(mask))
700 }
701
702 #[inline(always)]
703 fn expand_z(self, mask: Self::Mask) -> Self {
704 Complex::new(self.re.expand_z(mask), self.im.expand_z(mask))
705 }
706
707 #[inline(always)]
708 fn expand_m(self, src: Self, mask: Self::Mask) -> Self {
709 Complex::new(self.re.expand_m(src.re, mask), self.im.expand_m(src.im, mask))
710 }
711
712 #[inline(always)]
713 fn align<const OFFSET: usize>(self, other: Self) -> Self {
714 Complex::new(self.re.align::<OFFSET>(other.re), self.im.align::<OFFSET>(other.im))
715 }
716
717 const HAS_NATIVE_ALIGN: bool = V::HAS_NATIVE_ALIGN;
719
720 #[inline(always)]
721 fn map<F>(mut self, f: F) -> Self
722 where
723 F: Fn(Self::Element) -> Self::Element,
724 {
725 for i in 0..Self::LANES {
726 self = self.insertv(i, f(self.extractv(i)));
727 }
728 self
729 }
730
731 #[inline(always)]
732 fn fold<F>(self, mut init: Self::Element, f: F) -> Self::Element
733 where
734 F: Fn(Self::Element, Self::Element) -> Self::Element,
735 {
736 for i in 0..Self::LANES {
737 init = f(init, self.extractv(i));
738 }
739 init
740 }
741
742 #[inline(always)]
743 fn reduce<F>(self, f: F) -> Self::Element
744 where
745 F: Fn(Self::Element, Self::Element) -> Self::Element,
746 {
747 let mut result = self.extractv(0);
748 for i in 1..Self::LANES {
749 result = f(result, self.extractv(i));
750 }
751 result
752 }
753
754 #[rustfmt::skip]
755 #[inline(always)] fn splat_m(src: Self, mask: Self::Mask, value: Self::Element) -> Self { mask.select(Self::splat(value), src) }
756 #[rustfmt::skip]
757 #[inline(always)] fn splat_z(mask: Self::Mask, value: Self::Element) -> Self { mask.select(Self::splat(value), Self::EMPTY) }
758 #[rustfmt::skip]
759 #[inline(always)] fn broadcast_c<const I: usize>(self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), self) }
760 #[rustfmt::skip]
761 #[inline(always)] fn broadcast_m<const I: usize>(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), src) }
762 #[rustfmt::skip]
763 #[inline(always)] fn broadcast_z<const I: usize>(self, mask: Self::Mask) -> Self { mask.select(self.broadcast::<I>(), Self::EMPTY) }
764 #[rustfmt::skip]
765 #[inline(always)] fn broadcastv_c(self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), self) }
766 #[rustfmt::skip]
767 #[inline(always)] fn broadcastv_m(self, src: Self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), src) }
768 #[rustfmt::skip]
769 #[inline(always)] fn broadcastv_z(self, mask: Self::Mask, idx: usize) -> Self { mask.select(self.broadcastv(idx), Self::EMPTY) }
770 #[rustfmt::skip]
771 #[inline(always)] fn reverse_c(self, mask: Self::Mask) -> Self { mask.select(self.reverse(), self) }
772 #[rustfmt::skip]
773 #[inline(always)] fn reverse_m(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.reverse(), src) }
774 #[rustfmt::skip]
775 #[inline(always)] fn reverse_z(self, mask: Self::Mask) -> Self { mask.select(self.reverse(), Self::EMPTY) }
776 #[rustfmt::skip]
777 #[inline(always)] fn swap_bytes_c(self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), self) }
778 #[rustfmt::skip]
779 #[inline(always)] fn swap_bytes_m(self, src: Self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), src) }
780 #[rustfmt::skip]
781 #[inline(always)] fn swap_bytes_z(self, mask: Self::Mask) -> Self { mask.select(self.swap_bytes(), Self::EMPTY) }
782}
783
784#[rustfmt::skip]
787impl<V: RealFloatVector> PartialOrdVector for Complex<V> {
788 #[inline(always)]
789 fn cmp_eq(self, other: Self) -> Self::Mask {
790 self.re.cmp_eq(other.re) & self.im.cmp_eq(other.im)
791 }
792
793 #[inline(always)]
794 fn cmp_ne(self, other: Self) -> Self::Mask {
795 self.re.cmp_ne(other.re) | self.im.cmp_ne(other.im)
796 }
797
798 #[inline(always)]
800 fn cmp_lt(self, other: Self) -> Self::Mask {
801 let re_lt = self.re.cmp_lt(other.re);
802 let re_eq = self.re.cmp_eq(other.re);
803 let im_lt = self.im.cmp_lt(other.im);
804
805 GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_lt, re_eq, im_lt)
806 }
807
808 #[inline(always)]
809 fn cmp_gt(self, other: Self) -> Self::Mask {
810 let re_gt = self.re.cmp_gt(other.re);
811 let re_eq = self.re.cmp_eq(other.re);
812 let im_gt = self.im.cmp_gt(other.im);
813
814 GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_gt, re_eq, im_gt)
815 }
816
817 #[inline(always)]
818 fn cmp_le(self, other: Self) -> Self::Mask {
819 let re_lt = self.re.cmp_lt(other.re);
820 let re_eq = self.re.cmp_eq(other.re);
821 let im_le = self.im.cmp_le(other.im);
822
823 GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_lt, re_eq, im_le)
824 }
825
826 #[inline(always)]
827 fn cmp_ge(self, other: Self) -> Self::Mask {
828 let re_gt = self.re.cmp_gt(other.re);
829 let re_eq = self.re.cmp_eq(other.re);
830 let im_ge = self.im.cmp_ge(other.im);
831
832 GenericMask::ternlog::<{ thermite::ternlog_imm!(A | (B & C)) }>(re_gt, re_eq, im_ge)
833 }
834}
835
836macro_rules! impl_masked {
839 (MUL_ADD: $($method:ident),*) => {paste::paste! {
840 impl<V: RealFloatVector, A, B> thermite::vector::ops::MulAddExtMasked<V::Mask, A, B> for Complex<V>
841 where
842 Complex<V>: thermite::vector::ops::MulAddExt<A, B, Output = Self>,
843 {
844 $(
845 #[inline(always)]
846 fn [<$method _c>](self, mask: V::Mask, a: A, b: B) -> Self {
847 mask.select(self.$method(a, b), self)
848 }
849 #[inline(always)]
850 fn [<$method _m>](self, src: Self, mask: V::Mask, a: A, b: B) -> Self {
851 mask.select(self.$method(a, b), src)
852 }
853 #[inline(always)]
854 fn [<$method _z>](self, mask: V::Mask, a: A, b: B) -> Self {
855 mask.select(self.$method(a, b), Self::EMPTY)
856 }
857 )*
858 }
859
860 impl<V: RealFloatVector, A, B> thermite::vector::ops::MulAddAssignExtMasked<V::Mask, A, B> for Complex<V>
861 where
862 Complex<V>: thermite::vector::ops::MulAddExt<A, B, Output = Self>,
863 {
864 $(
865 #[inline(always)]
866 fn [<$method _assign_c>](&mut self, mask: V::Mask, a: A, b: B) {
867 *self = mask.select(self.$method(a, b), *self);
868 }
869 #[inline(always)]
870 fn [<$method _assign_m>](&mut self, src: Self, mask: V::Mask, a: A, b: B) {
871 *self = mask.select(self.$method(a, b), src);
872 }
873 #[inline(always)]
874 fn [<$method _assign_z>](&mut self, mask: V::Mask, a: A, b: B) {
875 *self = mask.select(self.$method(a, b), Self::EMPTY);
876 }
877 )*
878 }
879 }};
880
881 ($trait:ident::$method:ident) => {paste::paste! {
882 impl<V: RealFloatVector, Rhs> thermite::vector::ops::[<$trait Masked>]<V::Mask, Rhs> for Complex<V>
883 where
884 Complex<V>: core::ops::$trait<Rhs, Output = Self>,
885 {
886 #[inline(always)]
887 fn [<$method _c>](self, mask: V::Mask, rhs: Rhs) -> Self {
888 mask.select(self.$method(rhs), self)
889 }
890 #[inline(always)]
891 fn [<$method _m>](self, src: Self, mask: V::Mask, rhs: Rhs) -> Self {
892 mask.select(self.$method(rhs), src)
893 }
894 #[inline(always)]
895 fn [<$method _z>](self, mask: V::Mask, rhs: Rhs) -> Self {
896 mask.select(self.$method(rhs), Self::EMPTY)
897 }
898 }
899
900 impl<V: RealFloatVector, Rhs> thermite::vector::ops::[<$trait AssignMasked>]<V::Mask, Rhs> for Complex<V>
901 where
902 Complex<V>: core::ops::$trait<Rhs, Output = Self>,
903 {
904 #[inline(always)]
905 fn [<$method _assign_c>](&mut self, mask: V::Mask, rhs: Rhs) {
906 *self = mask.select(self.$method(rhs), *self);
907 }
908 #[inline(always)]
909 fn [<$method _assign_m>](&mut self, src: Self, mask: V::Mask, rhs: Rhs) {
910 *self = mask.select(self.$method(rhs), src);
911 }
912 #[inline(always)]
913 fn [<$method _assign_z>](&mut self, mask: V::Mask, rhs: Rhs) {
914 *self = mask.select(self.$method(rhs), Self::EMPTY);
915 }
916 }
917 }};
918}
919
920impl_masked!(MUL_ADD: mul_add, mul_sub, nmul_add, nmul_sub, mul_adde, mul_sube, nmul_adde, nmul_sube);
921impl_masked!(Add::add);
922impl_masked!(Sub::sub);
923impl_masked!(Mul::mul);
924impl_masked!(Div::div);
925impl_masked!(Rem::rem);
926
927#[inline(always)]
938fn neg_even_complex<V: RealFloatVector>(x: Complex<V>) -> Complex<V> {
939 Complex::new(V::ZERO.addsub(x.re), V::ZERO.addsub(x.im))
941}
942
943impl<V: RealFloatVector> AddSubExt for Complex<V> {
944 type Output = Self;
945
946 #[inline(always)]
947 fn addsub(self, b: Self) -> Self {
948 self + neg_even_complex(b)
949 }
950 #[inline(always)]
951 fn fmaddsub(self, b: Self, c: Self) -> Self {
952 self.mul_adde(b, neg_even_complex(c))
953 }
954 #[inline(always)]
955 fn fmsubadd(self, b: Self, c: Self) -> Self {
956 self.mul_sube(b, neg_even_complex(c))
957 }
958}
959
960impl<V: RealFloatVector> AddSubExtMasked<V::Mask> for Complex<V> {
961 #[inline(always)]
962 fn addsub_c(self, mask: V::Mask, b: Self) -> Self {
963 mask.select(self.addsub(b), self)
964 }
965 #[inline(always)]
966 fn addsub_m(self, src: Self, mask: V::Mask, b: Self) -> Self {
967 mask.select(self.addsub(b), src)
968 }
969 #[inline(always)]
970 fn addsub_z(self, mask: V::Mask, b: Self) -> Self {
971 mask.select(self.addsub(b), Self::EMPTY)
972 }
973
974 #[inline(always)]
975 fn fmaddsub_c(self, mask: V::Mask, b: Self, c: Self) -> Self {
976 mask.select(self.fmaddsub(b, c), self)
977 }
978 #[inline(always)]
979 fn fmaddsub_m(self, src: Self, mask: V::Mask, b: Self, c: Self) -> Self {
980 mask.select(self.fmaddsub(b, c), src)
981 }
982 #[inline(always)]
983 fn fmaddsub_z(self, mask: V::Mask, b: Self, c: Self) -> Self {
984 mask.select(self.fmaddsub(b, c), Self::EMPTY)
985 }
986
987 #[inline(always)]
988 fn fmsubadd_c(self, mask: V::Mask, b: Self, c: Self) -> Self {
989 mask.select(self.fmsubadd(b, c), self)
990 }
991 #[inline(always)]
992 fn fmsubadd_m(self, src: Self, mask: V::Mask, b: Self, c: Self) -> Self {
993 mask.select(self.fmsubadd(b, c), src)
994 }
995 #[inline(always)]
996 fn fmsubadd_z(self, mask: V::Mask, b: Self, c: Self) -> Self {
997 mask.select(self.fmsubadd(b, c), Self::EMPTY)
998 }
999}
1000
1001impl<V: RealFloatVector> SquareMasked<V::Mask> for Complex<V> {
1002 #[inline(always)]
1003 fn square_c(self, mask: V::Mask) -> Self::Output {
1004 mask.select(self.square(), self)
1005 }
1006
1007 #[inline(always)]
1008 fn square_m(self, src: Self, mask: V::Mask) -> Self::Output {
1009 mask.select(self.square(), src)
1010 }
1011
1012 #[inline(always)]
1013 fn square_z(self, mask: V::Mask) -> Self::Output {
1014 mask.select(self.square(), Self::ZERO)
1015 }
1016}
1017
1018macro_rules! complex_masked {
1021 (unary: $($m:ident),* $(,)?) => { paste::paste! {
1022 $(
1023 #[inline(always)] fn [<$m _c>](self, mask: Self::Mask) -> Self { mask.select(self.$m(), self) }
1024 #[inline(always)] fn [<$m _m>](self, src: Self, mask: Self::Mask) -> Self { mask.select(self.$m(), src) }
1025 #[inline(always)] fn [<$m _z>](self, mask: Self::Mask) -> Self { mask.select(self.$m(), Self::ZERO) }
1026 )*
1027 }};
1028 (binary: $($m:ident),* $(,)?) => { paste::paste! {
1029 $(
1030 #[inline(always)] fn [<$m _c>](self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), self) }
1031 #[inline(always)] fn [<$m _m>](self, src: Self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), src) }
1032 #[inline(always)] fn [<$m _z>](self, mask: Self::Mask, rhs: Self) -> Self { mask.select(self.$m(rhs), Self::ZERO) }
1033 )*
1034 }};
1035}
1036
1037#[rustfmt::skip]
1040impl<V: RealFloatVector> Bounded for Complex<V> {
1041 #[inline(always)] fn min_value() -> Self { Complex::new(V::MIN, V::MIN) }
1042 #[inline(always)] fn max_value() -> Self { Complex::new(V::MAX, V::MAX) }
1043}
1044
1045#[rustfmt::skip]
1046impl<V: RealFloatVector> thermite::sort::SortKey<Self> for Complex<V> {
1050 #[inline(always)]
1051 fn key_lt(a: Self, b: Self) -> V::Mask {
1052 a.cmp_lt(b)
1053 }
1054}
1055
1056#[inline(always)]
1060fn sort_lanes_scalar<V: NumericVector, O: thermite::sort::SortOrder>(v: V) -> V
1061where
1062 V::Element: PartialOrd,
1063{
1064 let mut out = v;
1065 let mut i = 1;
1066 while i < V::LANES {
1067 let key = out.extractv(i);
1068 let mut j = i;
1069 while j > 0 {
1070 let prev = out.extractv(j - 1);
1071 let misplaced = if O::IS_ASCENDING { prev > key } else { prev < key };
1072 if !misplaced {
1073 break;
1074 }
1075 out = out.insertv(j, prev);
1076 j -= 1;
1077 }
1078 out = out.insertv(j, key);
1079 i += 1;
1080 }
1081 out
1082}
1083
1084impl<V: RealFloatVector> NumericVector for Complex<V> {
1085 #[inline(always)]
1089 fn to_signed_integer(self) -> Self::Signed {
1090 self.re.to_signed_integer()
1091 }
1092
1093 #[inline(always)]
1094 fn from_signed_integer(v: Self::Signed) -> Self {
1095 Self::real(V::from_signed_integer(v))
1096 }
1097
1098 #[inline(always)]
1099 fn to_unsigned_integer(self) -> Self::Unsigned {
1100 self.re.to_unsigned_integer()
1101 }
1102
1103 #[inline(always)]
1104 fn from_unsigned_integer(v: Self::Unsigned) -> Self {
1105 Self::real(V::from_unsigned_integer(v))
1106 }
1107
1108 const ZERO: Self = Complex::new(V::ZERO, V::ZERO);
1109 const ONE: Self = Complex::new(V::ONE, V::ZERO);
1110 const TWO: Self = Self::real(V::TWO);
1111
1112 const MIN: Self = Complex::new(V::MIN, V::MIN);
1114 const MAX: Self = Complex::new(V::MAX, V::MAX);
1115
1116 #[inline(always)]
1117 fn is_zero(self) -> Self::Mask {
1118 self.re.is_zero() & self.im.is_zero()
1119 }
1120 #[inline(always)]
1121 fn is_all_zero(self) -> bool {
1122 self.re.is_all_zero() && self.im.is_all_zero()
1123 }
1124
1125 #[inline(always)]
1130 fn sort_by<O: thermite::sort::SortOrder>(self) -> Self {
1131 if const { Self::LANES <= 16 && Self::LANES.is_power_of_two() } {
1132 thermite::sort::sort_lanes_by_key::<Self, O, Self>(self)
1133 } else {
1134 sort_lanes_scalar::<Self, O>(self)
1135 }
1136 }
1137
1138 #[inline(always)]
1139 fn bitonic_clean_by<O: thermite::sort::SortOrder>(self) -> Self {
1140 if const { Self::LANES <= 16 && Self::LANES.is_power_of_two() } {
1141 thermite::sort::bitonic_clean_lanes_by_key::<Self, O, Self>(self)
1142 } else {
1143 sort_lanes_scalar::<Self, O>(self)
1145 }
1146 }
1147
1148 #[inline(always)]
1149 fn min(self, other: Self) -> Self {
1150 self.cmp_lt(other).select(self, other)
1151 }
1152 #[inline(always)]
1153 fn max(self, other: Self) -> Self {
1154 self.cmp_gt(other).select(self, other)
1155 }
1156
1157 #[inline(always)]
1159 fn clamp(self, min: Self, max: Self) -> Self {
1160 let is_lt = self.cmp_lt(min);
1161 let is_gt = self.cmp_gt(max);
1162
1163 Complex::new(
1164 is_lt.select(min.re, is_gt.select(max.re, self.re)),
1165 is_lt.select(min.im, is_gt.select(max.im, self.im)),
1166 )
1167 }
1168
1169 #[inline(always)]
1173 fn min_element(self) -> Self::Element {
1174 let mut arr = self.into_array();
1175 reduce_in_place(&mut arr, |a, b| if b < a { b } else { a });
1176 arr[0]
1177 }
1178
1179 #[inline(always)]
1180 fn max_element(self) -> Self::Element {
1181 let mut arr = self.into_array();
1182 reduce_in_place(&mut arr, |a, b| if b > a { b } else { a });
1183 arr[0]
1184 }
1185
1186 #[inline(always)]
1187 fn min_max_element(self) -> (Self::Element, Self::Element) {
1188 (self.min_element(), self.max_element())
1189 }
1190
1191 #[inline(always)]
1192 fn arg_minmax(self) -> (usize, usize) {
1193 let arr = self.into_array();
1194
1195 let (mut lo, mut hi) = (0, 0);
1196
1197 for i in 1..Self::LANES {
1198 if arr[i] < arr[lo] {
1199 lo = i;
1200 }
1201 if arr[i] > arr[hi] {
1202 hi = i;
1203 }
1204 }
1205
1206 (lo, hi)
1207 }
1208
1209 #[inline(always)]
1213 fn sum_elements(self) -> Self::Element {
1214 Complex::new(self.re.sum_elements(), self.im.sum_elements())
1215 }
1216
1217 #[inline(always)]
1220 fn prefix_sum(self) -> Self {
1221 Complex::new(self.re.prefix_sum(), self.im.prefix_sum())
1222 }
1223
1224 #[inline(always)]
1225 fn reverse_prefix_sum(self) -> Self {
1226 Complex::new(self.re.reverse_prefix_sum(), self.im.reverse_prefix_sum())
1227 }
1228
1229 #[inline(always)]
1234 fn prefix_min(self) -> Self {
1235 thermite::scan_ladder!(forward, self, self.broadcast::<0>(), Self::min)
1236 }
1237
1238 #[inline(always)]
1239 fn prefix_max(self) -> Self {
1240 thermite::scan_ladder!(forward, self, self.broadcast::<0>(), Self::max)
1241 }
1242
1243 #[inline(always)]
1244 fn reverse_prefix_min(self) -> Self {
1245 thermite::scan_ladder!(reverse, self, self.reverse().broadcast::<0>(), Self::min)
1246 }
1247
1248 #[inline(always)]
1249 fn reverse_prefix_max(self) -> Self {
1250 thermite::scan_ladder!(reverse, self, self.reverse().broadcast::<0>(), Self::max)
1251 }
1252
1253 #[inline(always)]
1257 fn prod_elements(self) -> Self::Element {
1258 let mut arr = self.into_array();
1259 reduce_in_place(&mut arr, |a, b| a * b);
1260 arr[0]
1261 }
1262
1263 #[inline(always)]
1264 fn offset() -> Self {
1265 Self::real(V::offset())
1266 }
1267 #[inline(always)]
1268 fn indexed() -> Self {
1269 Self::real(V::indexed())
1270 }
1271
1272 #[inline(always)]
1273 fn scale(self, factor: Self::Element) -> Self {
1274 self * Self::splat(factor)
1275 }
1276
1277 #[inline(always)]
1278 fn scale_c(self, mask: Self::Mask, factor: Self::Element) -> Self {
1279 mask.select(<Self as NumericVector>::scale(self, factor), self)
1280 }
1281 #[inline(always)]
1282 fn scale_m(self, src: Self, mask: Self::Mask, factor: Self::Element) -> Self {
1283 mask.select(<Self as NumericVector>::scale(self, factor), src)
1284 }
1285 #[inline(always)]
1286 fn scale_z(self, mask: Self::Mask, factor: Self::Element) -> Self {
1287 mask.select(<Self as NumericVector>::scale(self, factor), Self::ZERO)
1288 }
1289 complex_masked!(binary: min, max);
1290
1291 #[inline(always)]
1294 fn pairwise_sum(lo: Self, hi: Self) -> Self {
1295 Complex::new(V::pairwise_sum(lo.re, hi.re), V::pairwise_sum(lo.im, hi.im))
1296 }
1297
1298 #[inline(always)]
1299 fn relaxed_pairwise_sum(lo: Self, hi: Self) -> Self {
1300 Complex::new(
1301 V::relaxed_pairwise_sum(lo.re, hi.re),
1302 V::relaxed_pairwise_sum(lo.im, hi.im),
1303 )
1304 }
1305}
1306
1307impl<V: RealFloatVector> NegMasked<V::Mask> for Complex<V> {
1310 #[inline(always)]
1311 fn neg_c(self, mask: V::Mask) -> Self {
1312 Complex::new(self.re.neg_c(mask), self.im.neg_c(mask))
1313 }
1314
1315 #[inline(always)]
1316 fn neg_m(self, src: Self, mask: V::Mask) -> Self {
1317 Complex::new(self.re.neg_m(src.re, mask), self.im.neg_m(src.im, mask))
1318 }
1319
1320 #[inline(always)]
1321 fn neg_z(self, mask: V::Mask) -> Self {
1322 Complex::new(self.re.neg_z(mask), self.im.neg_z(mask))
1323 }
1324}
1325
1326impl<V: RealFloatVector> Complex<V> {
1327 #[inline(always)]
1331 pub(crate) fn modulus(self) -> V {
1332 self.re.hypot_p::<DefaultPolicy>(self.im)
1333 }
1334}
1335
1336#[rustfmt::skip]
1337impl<V: RealFloatVector> SignedVector for Complex<V> {
1338 const NEG_ONE: Self = Self::real(V::NEG_ONE);
1339 const MIN_POSITIVE: Self = Self::real(V::MIN_POSITIVE);
1340
1341 #[inline(always)]
1343 fn abs(self) -> Self {
1344 Self::real(self.modulus())
1345 }
1346
1347 #[inline(always)]
1351 fn signum(self) -> Self {
1352 let m = self.modulus();
1353 let is_zero = m.is_zero();
1354 let inv = m.reciprocal_p::<DefaultPolicy>();
1355
1356 Complex::new((self.re * inv).nz(is_zero), (self.im * inv).nz(is_zero))
1359 }
1360
1361 #[inline(always)] fn is_positive(self) -> Self::Mask { self.re.is_positive() }
1364 #[inline(always)] fn is_negative(self) -> Self::Mask { self.re.is_negative() }
1365 #[inline(always)] fn select_negative(self, if_neg: Self, if_pos: Self) -> Self { self.is_negative().select(if_neg, if_pos) }
1366
1367 #[inline(always)]
1368 fn copysign(self, sign: Self) -> Self {
1369 Complex::new(self.re.copysign(sign.re), self.im.copysign(sign.im))
1370 }
1371
1372 complex_masked!(unary: abs);
1373 complex_masked!(binary: copysign);
1374}
1375
1376#[inline(always)]
1384fn sqrt_kahan<V: RealFloatVector>(z: Complex<V>) -> (Complex<V>, V::Mask) {
1385 let half = <V as FloatVector>::HALF;
1386
1387 let m = z.modulus();
1388
1389 let t = z.re.abs().mul_adde(half, m * half).sqrt(); let half_im = z.im * half;
1396
1397 let q = half_im / t;
1399
1400 let re_pos = z.re.is_positive();
1402
1403 let degenerate = if const { features::PRESERVE_DENORMALS || features::IGNORE_DENORMALS } {
1416 t.is_zero()
1417 } else {
1418 m.is_zero()
1419 };
1420
1421 (
1422 Complex::new(re_pos.select(t, q.abs()), re_pos.select(q, t.mul_sign(z.im))),
1423 degenerate,
1424 )
1425}
1426
1427#[inline(always)]
1441fn sqrt_rescaled<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
1442 let up = V::ONE / <V as SignedVector>::MIN_POSITIVE;
1443 let down = <V as SignedVector>::MIN_POSITIVE.sqrt();
1444
1445 let (r, _) = sqrt_kahan(Complex::new(z.re * up, z.im * up));
1446
1447 let is_zero = z.is_zero();
1448
1449 Complex::new((r.re * down).nz(is_zero), (r.im * down).nz(is_zero))
1450}
1451
1452#[rustfmt::skip]
1453impl<V: RealFloatVector> FloatVector for Complex<V> {
1454 const HALF: Self = Self::real(<V as FloatVector>::HALF);
1455 const NEG_ZERO: Self = Self::real(<V as FloatVector>::NEG_ZERO);
1456 const EPSILON: Self = Self::real(<V as FloatVector>::EPSILON);
1457
1458 const INFINITY: Self = Self::real(<V as FloatVector>::INFINITY);
1461 const NEG_INFINITY: Self = Self::real(<V as FloatVector>::NEG_INFINITY);
1462 const NAN: Self = Complex::new(<V as FloatVector>::NAN, <V as FloatVector>::NAN);
1463
1464 type ExtendedPrecision = Self;
1465
1466 const HAS_APPROX_RCP: bool = V::HAS_APPROX_RCP;
1469 const HAS_APPROX_RSQRT: bool = V::HAS_APPROX_RCP;
1470
1471 #[inline(always)] fn is_nan(self) -> Self::Mask { self.re.is_nan() | self.im.is_nan() }
1472 #[inline(always)] fn is_infinite(self) -> Self::Mask { self.re.is_infinite() | self.im.is_infinite() }
1473 #[inline(always)] fn is_finite(self) -> Self::Mask { self.re.is_finite() & self.im.is_finite() }
1474
1475 #[inline(always)] fn is_zero_or_subnormal(self) -> Self::Mask { self.re.is_zero_or_subnormal() & self.im.is_zero_or_subnormal() }
1478 #[inline(always)] fn is_normal(self) -> Self::Mask { (self.re.is_normal() | self.im.is_normal()) & self.is_finite() }
1479 #[inline(always)] fn is_subnormal(self) -> Self::Mask { (self.re.is_subnormal() | self.im.is_subnormal()) & self.is_zero_or_subnormal() }
1480
1481 #[inline(always)]
1483 fn sqrt(self) -> Self {
1484 let (res, degenerate) = sqrt_kahan(self);
1495
1496 if thermite::unlikely(degenerate.any()) {
1500 return degenerate.select(sqrt_rescaled(self), res);
1501 }
1502
1503 res
1504 }
1505
1506 #[inline(always)]
1508 fn rcp(self) -> Self {
1509 let inv = self.norm_sqr().rcp();
1510
1511 Complex::new(self.re * inv, -(self.im * inv))
1512 }
1513
1514 #[inline(always)]
1516 fn rsqrt(self) -> Self {
1517 let s = self.sqrt();
1518 let inv = self.modulus().rcp();
1519
1520 Complex::new(s.re * inv, -(s.im * inv))
1521 }
1522
1523 #[inline(always)] fn floor(self) -> Self { Complex::new(self.re.floor(), self.im.floor()) }
1525 #[inline(always)] fn ceil(self) -> Self { Complex::new(self.re.ceil(), self.im.ceil()) }
1526 #[inline(always)] fn round(self) -> Self { Complex::new(self.re.round(), self.im.round()) }
1527 #[inline(always)] fn trunc(self) -> Self { Complex::new(self.re.trunc(), self.im.trunc()) }
1528 #[inline(always)] fn fract(self) -> Self { Complex::new(self.re.fract(), self.im.fract()) }
1529
1530 #[inline(always)] fn mul_sign(self, sign: Self) -> Self { Complex::new(self.re.mul_sign(sign.re), self.im.mul_sign(sign.im)) }
1531 #[inline(always)] fn signed_zero(self) -> Self { Complex::new(self.re.signed_zero(), self.im.signed_zero()) }
1532
1533 #[inline(always)] fn next_up(self) -> Self { Complex::new(self.re.next_up(), self.im.next_up()) }
1534 #[inline(always)] fn next_down(self) -> Self { Complex::new(self.re.next_down(), self.im.next_down()) }
1535
1536 #[inline(always)]
1537 unsafe fn block_autovectorization(&mut self) {
1538 unsafe {
1539 self.re.block_autovectorization();
1540 self.im.block_autovectorization();
1541 }
1542 }
1543
1544 #[inline(always)] fn mix(self, a: Self, b: Self) -> Self { (b - a).mul_adde(self, a) }
1547
1548 complex_masked!(unary: sqrt, rsqrt, rcp, floor, ceil, round, trunc, fract, signed_zero, next_up, next_down);
1549 complex_masked!(binary: mul_sign);
1550}