Skip to main content

strided_basic/
simd.rs

1#[inline(always)]
2pub(crate) fn dispatch<R>(f: impl FnOnce() -> R) -> R {
3    #[cfg(feature = "simd")]
4    {
5        pulp::Arch::new().dispatch(f)
6    }
7    #[cfg(not(feature = "simd"))]
8    {
9        f()
10    }
11}
12
13#[inline(always)]
14pub(crate) fn dispatch_if_large<R>(len: usize, f: impl FnOnce() -> R) -> R {
15    // Avoid runtime-dispatch overhead for tiny loops (especially common for small-array cases).
16    // This is a heuristic; correctness does not depend on it.
17    if len >= 64 {
18        dispatch(f)
19    } else {
20        f()
21    }
22}
23
24#[cfg(feature = "simd")]
25#[inline(always)]
26unsafe fn cast_slice<T, U>(src: &[T]) -> &[U] {
27    debug_assert_eq!(std::mem::size_of::<T>(), std::mem::size_of::<U>());
28    unsafe { std::slice::from_raw_parts(src.as_ptr().cast::<U>(), src.len()) }
29}
30
31#[cfg(feature = "simd")]
32macro_rules! impl_simd_mul_ptr {
33    (
34        $mul_into:ident,
35        $ty:ty,
36        $split:ident,
37        $split_uninit:ident,
38        $load:ident,
39        $mask:ident,
40        $store_ptr:ident,
41        $mul:ident,
42        $mask_scale:expr
43    ) => {
44        unsafe fn $mul_into(dst: *mut $ty, len: usize, a: &[$ty], b: &[$ty]) {
45            struct Mul<'a> {
46                dst: *mut $ty,
47                len: usize,
48                a: &'a [$ty],
49                b: &'a [$ty],
50            }
51
52            impl<'a> pulp::WithSimd for Mul<'a> {
53                type Output = ();
54
55                #[inline(always)]
56                fn with_simd<S: pulp::Simd>(self, simd: S) -> Self::Output {
57                    debug_assert_eq!(self.len, self.a.len());
58                    debug_assert_eq!(self.len, self.b.len());
59
60                    // SAFETY: the caller supplies len writable, nonaliasing
61                    // elements. MaybeUninit does not assert initialized output.
62                    let output = unsafe {
63                        core::slice::from_raw_parts_mut(
64                            self.dst.cast::<core::mem::MaybeUninit<$ty>>(),
65                            self.len,
66                        )
67                    };
68                    let (dest, tail) = S::$split_uninit(output);
69                    let (a, a_tail) = S::$split(self.a);
70                    let (b, b_tail) = S::$split(self.b);
71                    // Full vectors use ordinary stores. Pulp's partial-access
72                    // machinery is needed only for the final incomplete vector.
73                    for ((dst, &a), &b) in dest.iter_mut().zip(a).zip(b) {
74                        dst.write(simd.$mul(a, b));
75                    }
76                    if !tail.is_empty() {
77                        let va = simd.$load(a_tail);
78                        let vb = simd.$load(b_tail);
79                        // SAFETY: the mask writes exactly the remaining lanes.
80                        unsafe {
81                            simd.$store_ptr(
82                                simd.$mask(0, (tail.len() * $mask_scale) as _),
83                                tail.as_mut_ptr().cast::<$ty>(),
84                                simd.$mul(va, vb),
85                            );
86                        }
87                    }
88                }
89            }
90
91            pulp::Arch::new().dispatch(Mul { dst, len, a, b });
92        }
93    };
94}
95
96#[cfg(feature = "simd")]
97impl_simd_mul_ptr!(
98    simd_mul_f32_into,
99    f32,
100    as_simd_f32s,
101    as_uninit_mut_simd_f32s,
102    partial_load_f32s,
103    mask_between_m32s,
104    mask_store_ptr_f32s,
105    mul_f32s,
106    1
107);
108
109#[cfg(feature = "simd")]
110impl_simd_mul_ptr!(
111    simd_mul_f64_into,
112    f64,
113    as_simd_f64s,
114    as_uninit_mut_simd_f64s,
115    partial_load_f64s,
116    mask_between_m64s,
117    mask_store_ptr_f64s,
118    mul_f64s,
119    1
120);
121
122#[cfg(feature = "simd")]
123impl_simd_mul_ptr!(
124    simd_mul_c32_into,
125    num_complex::Complex32,
126    as_simd_c32s,
127    as_uninit_mut_simd_c32s,
128    partial_load_c32s,
129    mask_between_m32s,
130    mask_store_ptr_c32s,
131    mul_e_c32s,
132    2
133);
134
135#[cfg(feature = "simd")]
136impl_simd_mul_ptr!(
137    simd_mul_c64_into,
138    num_complex::Complex64,
139    as_simd_c64s,
140    as_uninit_mut_simd_c64s,
141    partial_load_c64s,
142    mask_between_m64s,
143    mask_store_ptr_c64s,
144    mul_e_c64s,
145    2
146);
147
148#[cfg(feature = "simd")]
149#[inline]
150#[cfg_attr(not(test), allow(dead_code))]
151pub(crate) fn try_mul_contiguous<D: 'static, A: 'static, B: 'static>(
152    dst: &mut [D],
153    a: &[A],
154    b: &[B],
155) -> bool {
156    unsafe { try_mul_contiguous_ptr(dst.as_mut_ptr(), dst.len(), a, b) }
157}
158
159#[cfg(feature = "simd")]
160#[inline]
161pub(crate) unsafe fn try_mul_contiguous_ptr<D: 'static, A: 'static, B: 'static>(
162    dst: *mut D,
163    len: usize,
164    a: &[A],
165    b: &[B],
166) -> bool {
167    use std::any::TypeId;
168
169    macro_rules! try_same_type {
170        ($ty:ty, $mul_into:ident) => {
171            if TypeId::of::<D>() == TypeId::of::<$ty>()
172                && TypeId::of::<A>() == TypeId::of::<$ty>()
173                && TypeId::of::<B>() == TypeId::of::<$ty>()
174            {
175                unsafe { $mul_into(dst.cast::<$ty>(), len, cast_slice(a), cast_slice(b)) };
176                return true;
177            }
178        };
179    }
180
181    try_same_type!(f64, simd_mul_f64_into);
182    try_same_type!(f32, simd_mul_f32_into);
183    try_same_type!(num_complex::Complex64, simd_mul_c64_into);
184    try_same_type!(num_complex::Complex32, simd_mul_c32_into);
185
186    false
187}
188
189#[cfg(not(feature = "simd"))]
190#[inline]
191#[cfg_attr(not(test), allow(dead_code))]
192pub(crate) fn try_mul_contiguous<D: 'static, A: 'static, B: 'static>(
193    _dst: &mut [D],
194    _a: &[A],
195    _b: &[B],
196) -> bool {
197    false
198}
199
200#[cfg(not(feature = "simd"))]
201#[inline]
202pub(crate) unsafe fn try_mul_contiguous_ptr<D: 'static, A: 'static, B: 'static>(
203    _dst: *mut D,
204    _len: usize,
205    _a: &[A],
206    _b: &[B],
207) -> bool {
208    false
209}
210
211const TRANSPOSE_TILE: usize = 8;
212
213unsafe fn mul_transposed_scalar_rhs_source_contiguous<T>(
214    dst: *mut T,
215    src: *const T,
216    scalar: T,
217    inner_len: usize,
218    row_len: usize,
219    src_fast_stride: isize,
220    src_row_stride: isize,
221) where
222    T: Copy + std::ops::Mul<Output = T>,
223{
224    let mut inner0 = 0usize;
225
226    while inner0 < inner_len {
227        let inner_count = TRANSPOSE_TILE.min(inner_len - inner0);
228        let mut row0 = 0usize;
229
230        while row0 < row_len {
231            let row_count = TRANSPOSE_TILE.min(row_len - row0);
232
233            for inner in 0..inner_count {
234                let src_base = src.offset((inner0 + inner) as isize * src_fast_stride);
235                for row in 0..row_count {
236                    let row_index = row0 + row;
237                    let src_offset = row_index as isize * src_row_stride;
238                    *dst.add(row_index * inner_len + inner0 + inner) =
239                        *src_base.offset(src_offset) * scalar;
240                }
241            }
242
243            row0 += TRANSPOSE_TILE;
244        }
245
246        inner0 += TRANSPOSE_TILE;
247    }
248}
249
250unsafe fn mul_transposed_scalar_lhs_source_contiguous<T>(
251    dst: *mut T,
252    scalar: T,
253    src: *const T,
254    inner_len: usize,
255    row_len: usize,
256    src_fast_stride: isize,
257    src_row_stride: isize,
258) where
259    T: Copy + std::ops::Mul<Output = T>,
260{
261    let mut inner0 = 0usize;
262
263    while inner0 < inner_len {
264        let inner_count = TRANSPOSE_TILE.min(inner_len - inner0);
265        let mut row0 = 0usize;
266
267        while row0 < row_len {
268            let row_count = TRANSPOSE_TILE.min(row_len - row0);
269
270            for inner in 0..inner_count {
271                let src_base = src.offset((inner0 + inner) as isize * src_fast_stride);
272                for row in 0..row_count {
273                    let row_index = row0 + row;
274                    let src_offset = row_index as isize * src_row_stride;
275                    *dst.add(row_index * inner_len + inner0 + inner) =
276                        scalar * *src_base.offset(src_offset);
277                }
278            }
279
280            row0 += TRANSPOSE_TILE;
281        }
282
283        inner0 += TRANSPOSE_TILE;
284    }
285}
286
287#[inline]
288pub(crate) unsafe fn try_mul_transposed_scalar_rhs_2d<D: 'static, A: 'static, B: 'static>(
289    dst: *mut D,
290    src: *const A,
291    scalar: *const B,
292    inner_len: usize,
293    row_len: usize,
294    src_fast_stride: isize,
295    src_row_stride: isize,
296) -> bool {
297    use std::any::TypeId;
298
299    macro_rules! try_same_type {
300        ($ty:ty) => {
301            if TypeId::of::<D>() == TypeId::of::<$ty>()
302                && TypeId::of::<A>() == TypeId::of::<$ty>()
303                && TypeId::of::<B>() == TypeId::of::<$ty>()
304            {
305                unsafe {
306                    mul_transposed_scalar_rhs_source_contiguous(
307                        dst.cast::<$ty>(),
308                        src.cast::<$ty>(),
309                        *scalar.cast::<$ty>(),
310                        inner_len,
311                        row_len,
312                        src_fast_stride,
313                        src_row_stride,
314                    );
315                }
316                return true;
317            }
318        };
319    }
320
321    try_same_type!(f64);
322    try_same_type!(f32);
323    try_same_type!(num_complex::Complex64);
324    try_same_type!(num_complex::Complex32);
325
326    false
327}
328
329#[inline]
330pub(crate) unsafe fn try_mul_transposed_scalar_lhs_2d<D: 'static, A: 'static, B: 'static>(
331    dst: *mut D,
332    scalar: *const A,
333    src: *const B,
334    inner_len: usize,
335    row_len: usize,
336    src_fast_stride: isize,
337    src_row_stride: isize,
338) -> bool {
339    use std::any::TypeId;
340
341    macro_rules! try_same_type {
342        ($ty:ty) => {
343            if TypeId::of::<D>() == TypeId::of::<$ty>()
344                && TypeId::of::<A>() == TypeId::of::<$ty>()
345                && TypeId::of::<B>() == TypeId::of::<$ty>()
346            {
347                unsafe {
348                    mul_transposed_scalar_lhs_source_contiguous(
349                        dst.cast::<$ty>(),
350                        *scalar.cast::<$ty>(),
351                        src.cast::<$ty>(),
352                        inner_len,
353                        row_len,
354                        src_fast_stride,
355                        src_row_stride,
356                    );
357                }
358                return true;
359            }
360        };
361    }
362
363    try_same_type!(f64);
364    try_same_type!(f32);
365    try_same_type!(num_complex::Complex64);
366    try_same_type!(num_complex::Complex32);
367
368    false
369}
370
371/// Trait for types that may have SIMD-accelerated sum/dot operations.
372///
373/// Default implementations return `None` (no SIMD available).
374/// f32/f64 override these with SIMD kernels when the `simd` feature is enabled.
375pub trait MaybeSimdOps: Copy + Sized {
376    fn try_simd_sum(_src: &[Self]) -> Option<Self> {
377        None
378    }
379    fn try_simd_dot(_a: &[Self], _b: &[Self]) -> Option<Self> {
380        None
381    }
382}
383
384pub(crate) trait MaybeSimdProduct: Copy + Sized {
385    fn try_simd_product(_src: &[Self]) -> Option<Self> {
386        None
387    }
388}
389
390pub(crate) trait MaybeSimdSumSquares: Copy + Sized {
391    fn try_simd_sum_squares(_src: &[Self]) -> Option<Self> {
392        None
393    }
394}
395
396/// Minimum accumulator-block length worth a SIMD dispatch.
397#[cfg(feature = "simd")]
398const BLOCK_FOLD_MIN_LEN: usize = 16;
399
400/// SIMD fold of one unit-stride value run into adjacent accumulators.
401///
402/// This is a reduction-kernel implementation hook, not part of the public
403/// surface. Every method requires `acc.len() == values.len()`; the reduction
404/// callers always pass the same block length for both, so the vector loop and
405/// the scalar tail cover exactly the same indices. A method returns false when
406/// no SIMD path applies and leaves `acc` untouched.
407pub(crate) trait MaybeSimdBlockFold: Copy + Sized {
408    fn try_simd_add_block(_acc: &mut [Self], _values: &[Self]) -> bool {
409        false
410    }
411    fn try_simd_mul_block(_acc: &mut [Self], _values: &[Self]) -> bool {
412        false
413    }
414    fn try_simd_sum_squares_block(_acc: &mut [Self], _values: &[Self]) -> bool {
415        false
416    }
417}
418
419// Default (no-op) impls for integer types and Complex
420macro_rules! impl_no_simd {
421    ($($t:ty),*) => {
422        $(
423            impl MaybeSimdOps for $t {}
424            impl MaybeSimdProduct for $t {}
425            impl MaybeSimdSumSquares for $t {}
426            impl MaybeSimdBlockFold for $t {}
427        )*
428    };
429}
430
431impl_no_simd!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);
432
433impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdOps
434    for num_complex::Complex<T>
435{
436}
437impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdProduct
438    for num_complex::Complex<T>
439{
440}
441impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdSumSquares
442    for num_complex::Complex<T>
443{
444}
445impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdBlockFold
446    for num_complex::Complex<T>
447{
448}
449
450// f32/f64: SIMD-accelerated when feature enabled, no-op otherwise
451#[cfg(not(feature = "simd"))]
452impl MaybeSimdOps for f32 {}
453#[cfg(not(feature = "simd"))]
454impl MaybeSimdProduct for f32 {}
455#[cfg(not(feature = "simd"))]
456impl MaybeSimdSumSquares for f32 {}
457
458#[cfg(not(feature = "simd"))]
459impl MaybeSimdOps for f64 {}
460#[cfg(not(feature = "simd"))]
461impl MaybeSimdProduct for f64 {}
462#[cfg(not(feature = "simd"))]
463impl MaybeSimdSumSquares for f64 {}
464#[cfg(not(feature = "simd"))]
465impl MaybeSimdBlockFold for f32 {}
466#[cfg(not(feature = "simd"))]
467impl MaybeSimdBlockFold for f64 {}
468
469#[cfg(feature = "simd")]
470mod simd_impls {
471    use super::{
472        MaybeSimdBlockFold, MaybeSimdOps, MaybeSimdProduct, MaybeSimdSumSquares, BLOCK_FOLD_MIN_LEN,
473    };
474    use pulp::{Simd, WithSimd};
475
476    impl MaybeSimdOps for f32 {
477        fn try_simd_sum(src: &[f32]) -> Option<f32> {
478            struct Sum<'a>(&'a [f32]);
479            impl<'a> WithSimd for Sum<'a> {
480                type Output = f32;
481
482                #[inline(always)]
483                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
484                    let (head, tail) = S::as_simd_f32s(self.0);
485
486                    let mut acc0 = simd.splat_f32s(0.0);
487                    let mut acc1 = simd.splat_f32s(0.0);
488                    let mut acc2 = simd.splat_f32s(0.0);
489                    let mut acc3 = simd.splat_f32s(0.0);
490
491                    let mut i = 0usize;
492                    while i + 4 <= head.len() {
493                        acc0 = simd.add_f32s(acc0, head[i]);
494                        acc1 = simd.add_f32s(acc1, head[i + 1]);
495                        acc2 = simd.add_f32s(acc2, head[i + 2]);
496                        acc3 = simd.add_f32s(acc3, head[i + 3]);
497                        i += 4;
498                    }
499                    for &v in &head[i..] {
500                        acc0 = simd.add_f32s(acc0, v);
501                    }
502
503                    let acc = simd.add_f32s(simd.add_f32s(acc0, acc1), simd.add_f32s(acc2, acc3));
504                    let mut sum = simd.reduce_sum_f32s(acc);
505                    for &x in tail {
506                        sum += x;
507                    }
508                    sum
509                }
510            }
511
512            Some(pulp::Arch::new().dispatch(Sum(src)))
513        }
514
515        fn try_simd_dot(a: &[f32], b: &[f32]) -> Option<f32> {
516            try_simd_dot_f32(a, b)
517        }
518    }
519
520    impl MaybeSimdProduct for f32 {
521        fn try_simd_product(src: &[f32]) -> Option<f32> {
522            struct Product<'a>(&'a [f32]);
523            impl<'a> WithSimd for Product<'a> {
524                type Output = f32;
525
526                #[inline(always)]
527                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
528                    let (head, tail) = S::as_simd_f32s(self.0);
529
530                    let mut acc0 = simd.splat_f32s(1.0);
531                    let mut acc1 = simd.splat_f32s(1.0);
532                    let mut acc2 = simd.splat_f32s(1.0);
533                    let mut acc3 = simd.splat_f32s(1.0);
534
535                    let mut i = 0usize;
536                    while i + 4 <= head.len() {
537                        acc0 = simd.mul_f32s(acc0, head[i]);
538                        acc1 = simd.mul_f32s(acc1, head[i + 1]);
539                        acc2 = simd.mul_f32s(acc2, head[i + 2]);
540                        acc3 = simd.mul_f32s(acc3, head[i + 3]);
541                        i += 4;
542                    }
543                    for &value in &head[i..] {
544                        acc0 = simd.mul_f32s(acc0, value);
545                    }
546
547                    let acc = simd.mul_f32s(simd.mul_f32s(acc0, acc1), simd.mul_f32s(acc2, acc3));
548                    let mut product = simd.reduce_product_f32s(acc);
549                    for &value in tail {
550                        product *= value;
551                    }
552                    product
553                }
554            }
555
556            Some(pulp::Arch::new().dispatch(Product(src)))
557        }
558    }
559
560    impl MaybeSimdSumSquares for f32 {
561        fn try_simd_sum_squares(src: &[f32]) -> Option<f32> {
562            struct SumSquares<'a>(&'a [f32]);
563            impl<'a> WithSimd for SumSquares<'a> {
564                type Output = f32;
565
566                #[inline(always)]
567                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
568                    let (head, tail) = S::as_simd_f32s(self.0);
569                    let mut acc0 = simd.splat_f32s(0.0);
570                    let mut acc1 = simd.splat_f32s(0.0);
571                    let mut acc2 = simd.splat_f32s(0.0);
572                    let mut acc3 = simd.splat_f32s(0.0);
573
574                    let mut i = 0usize;
575                    while i + 4 <= head.len() {
576                        let square0 = simd.mul_f32s(head[i], head[i]);
577                        let square1 = simd.mul_f32s(head[i + 1], head[i + 1]);
578                        let square2 = simd.mul_f32s(head[i + 2], head[i + 2]);
579                        let square3 = simd.mul_f32s(head[i + 3], head[i + 3]);
580                        acc0 = simd.add_f32s(acc0, square0);
581                        acc1 = simd.add_f32s(acc1, square1);
582                        acc2 = simd.add_f32s(acc2, square2);
583                        acc3 = simd.add_f32s(acc3, square3);
584                        i += 4;
585                    }
586                    for &value in &head[i..] {
587                        let square = simd.mul_f32s(value, value);
588                        acc0 = simd.add_f32s(acc0, square);
589                    }
590
591                    let acc = simd.add_f32s(simd.add_f32s(acc0, acc1), simd.add_f32s(acc2, acc3));
592                    let mut sum = simd.reduce_sum_f32s(acc);
593                    for &value in tail {
594                        let square = value * value;
595                        sum += square;
596                    }
597                    sum
598                }
599            }
600
601            Some(pulp::Arch::new().dispatch(SumSquares(src)))
602        }
603    }
604
605    fn try_simd_dot_f32(a: &[f32], b: &[f32]) -> Option<f32> {
606        struct Dot<'a> {
607            a: &'a [f32],
608            b: &'a [f32],
609        }
610        impl<'a> WithSimd for Dot<'a> {
611            type Output = f32;
612
613            #[inline(always)]
614            fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
615                debug_assert_eq!(self.a.len(), self.b.len());
616                let (a_head, a_tail) = S::as_simd_f32s(self.a);
617                let (b_head, b_tail) = S::as_simd_f32s(self.b);
618                debug_assert_eq!(a_head.len(), b_head.len());
619                debug_assert_eq!(a_tail.len(), b_tail.len());
620
621                let mut acc0 = simd.splat_f32s(0.0);
622                let mut acc1 = simd.splat_f32s(0.0);
623                let mut acc2 = simd.splat_f32s(0.0);
624                let mut acc3 = simd.splat_f32s(0.0);
625
626                let mut i = 0usize;
627                while i + 4 <= a_head.len() {
628                    acc0 = simd.mul_add_f32s(a_head[i], b_head[i], acc0);
629                    acc1 = simd.mul_add_f32s(a_head[i + 1], b_head[i + 1], acc1);
630                    acc2 = simd.mul_add_f32s(a_head[i + 2], b_head[i + 2], acc2);
631                    acc3 = simd.mul_add_f32s(a_head[i + 3], b_head[i + 3], acc3);
632                    i += 4;
633                }
634                for j in i..a_head.len() {
635                    acc0 = simd.mul_add_f32s(a_head[j], b_head[j], acc0);
636                }
637
638                let acc = simd.add_f32s(simd.add_f32s(acc0, acc1), simd.add_f32s(acc2, acc3));
639                let mut sum = simd.reduce_sum_f32s(acc);
640                for (&x, &y) in a_tail.iter().zip(b_tail.iter()) {
641                    sum += x * y;
642                }
643                sum
644            }
645        }
646
647        Some(pulp::Arch::new().dispatch(Dot { a, b }))
648    }
649
650    impl MaybeSimdBlockFold for f32 {
651        fn try_simd_add_block(acc: &mut [f32], values: &[f32]) -> bool {
652            if acc.len() < BLOCK_FOLD_MIN_LEN {
653                return false;
654            }
655            struct AddBlock<'a, 'b>(&'a mut [f32], &'b [f32]);
656            impl pulp::WithSimd for AddBlock<'_, '_> {
657                type Output = ();
658                #[inline(always)]
659                fn with_simd<S: Simd>(self, simd: S) {
660                    let (ah, at) = S::as_mut_simd_f32s(self.0);
661                    let (vh, vt) = S::as_simd_f32s(self.1);
662                    // Equal block lengths are the trait contract, so the vector
663                    // prefix and the scalar tail cover the same indices.
664                    for (a, &v) in ah.iter_mut().zip(vh) {
665                        *a = simd.add_f32s(*a, v);
666                    }
667                    for (a, &v) in at.iter_mut().zip(vt) {
668                        *a += v;
669                    }
670                }
671            }
672            pulp::Arch::new().dispatch(AddBlock(acc, values));
673            true
674        }
675
676        fn try_simd_mul_block(acc: &mut [f32], values: &[f32]) -> bool {
677            if acc.len() < BLOCK_FOLD_MIN_LEN {
678                return false;
679            }
680            struct MulBlock<'a, 'b>(&'a mut [f32], &'b [f32]);
681            impl pulp::WithSimd for MulBlock<'_, '_> {
682                type Output = ();
683                #[inline(always)]
684                fn with_simd<S: Simd>(self, simd: S) {
685                    let (ah, at) = S::as_mut_simd_f32s(self.0);
686                    let (vh, vt) = S::as_simd_f32s(self.1);
687                    for (a, &v) in ah.iter_mut().zip(vh) {
688                        *a = simd.mul_f32s(*a, v);
689                    }
690                    for (a, &v) in at.iter_mut().zip(vt) {
691                        *a *= v;
692                    }
693                }
694            }
695            pulp::Arch::new().dispatch(MulBlock(acc, values));
696            true
697        }
698
699        fn try_simd_sum_squares_block(acc: &mut [f32], values: &[f32]) -> bool {
700            if acc.len() < BLOCK_FOLD_MIN_LEN {
701                return false;
702            }
703            struct SumSquaresBlock<'a, 'b>(&'a mut [f32], &'b [f32]);
704            impl pulp::WithSimd for SumSquaresBlock<'_, '_> {
705                type Output = ();
706                #[inline(always)]
707                fn with_simd<S: Simd>(self, simd: S) {
708                    let (ah, at) = S::as_mut_simd_f32s(self.0);
709                    let (vh, vt) = S::as_simd_f32s(self.1);
710                    for (a, &v) in ah.iter_mut().zip(vh) {
711                        let sq = simd.mul_f32s(v, v);
712                        *a = simd.add_f32s(*a, sq);
713                    }
714                    for (a, &v) in at.iter_mut().zip(vt) {
715                        *a += v * v;
716                    }
717                }
718            }
719            pulp::Arch::new().dispatch(SumSquaresBlock(acc, values));
720            true
721        }
722    }
723
724    impl MaybeSimdOps for f64 {
725        fn try_simd_sum(src: &[f64]) -> Option<f64> {
726            struct Sum<'a>(&'a [f64]);
727            impl<'a> WithSimd for Sum<'a> {
728                type Output = f64;
729
730                #[inline(always)]
731                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
732                    let (head, tail) = S::as_simd_f64s(self.0);
733
734                    let mut acc0 = simd.splat_f64s(0.0);
735                    let mut acc1 = simd.splat_f64s(0.0);
736                    let mut acc2 = simd.splat_f64s(0.0);
737                    let mut acc3 = simd.splat_f64s(0.0);
738
739                    let mut i = 0usize;
740                    while i + 4 <= head.len() {
741                        acc0 = simd.add_f64s(acc0, head[i]);
742                        acc1 = simd.add_f64s(acc1, head[i + 1]);
743                        acc2 = simd.add_f64s(acc2, head[i + 2]);
744                        acc3 = simd.add_f64s(acc3, head[i + 3]);
745                        i += 4;
746                    }
747                    for &v in &head[i..] {
748                        acc0 = simd.add_f64s(acc0, v);
749                    }
750
751                    let acc = simd.add_f64s(simd.add_f64s(acc0, acc1), simd.add_f64s(acc2, acc3));
752                    let mut sum = simd.reduce_sum_f64s(acc);
753                    for &x in tail {
754                        sum += x;
755                    }
756                    sum
757                }
758            }
759
760            Some(pulp::Arch::new().dispatch(Sum(src)))
761        }
762
763        fn try_simd_dot(a: &[f64], b: &[f64]) -> Option<f64> {
764            try_simd_dot_f64(a, b)
765        }
766    }
767
768    impl MaybeSimdProduct for f64 {
769        fn try_simd_product(src: &[f64]) -> Option<f64> {
770            struct Product<'a>(&'a [f64]);
771            impl<'a> WithSimd for Product<'a> {
772                type Output = f64;
773
774                #[inline(always)]
775                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
776                    let (head, tail) = S::as_simd_f64s(self.0);
777
778                    let mut acc0 = simd.splat_f64s(1.0);
779                    let mut acc1 = simd.splat_f64s(1.0);
780                    let mut acc2 = simd.splat_f64s(1.0);
781                    let mut acc3 = simd.splat_f64s(1.0);
782
783                    let mut i = 0usize;
784                    while i + 4 <= head.len() {
785                        acc0 = simd.mul_f64s(acc0, head[i]);
786                        acc1 = simd.mul_f64s(acc1, head[i + 1]);
787                        acc2 = simd.mul_f64s(acc2, head[i + 2]);
788                        acc3 = simd.mul_f64s(acc3, head[i + 3]);
789                        i += 4;
790                    }
791                    for &value in &head[i..] {
792                        acc0 = simd.mul_f64s(acc0, value);
793                    }
794
795                    let acc = simd.mul_f64s(simd.mul_f64s(acc0, acc1), simd.mul_f64s(acc2, acc3));
796                    let mut product = simd.reduce_product_f64s(acc);
797                    for &value in tail {
798                        product *= value;
799                    }
800                    product
801                }
802            }
803
804            Some(pulp::Arch::new().dispatch(Product(src)))
805        }
806    }
807
808    impl MaybeSimdSumSquares for f64 {
809        fn try_simd_sum_squares(src: &[f64]) -> Option<f64> {
810            struct SumSquares<'a>(&'a [f64]);
811            impl<'a> WithSimd for SumSquares<'a> {
812                type Output = f64;
813
814                #[inline(always)]
815                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
816                    let (head, tail) = S::as_simd_f64s(self.0);
817                    let mut acc0 = simd.splat_f64s(0.0);
818                    let mut acc1 = simd.splat_f64s(0.0);
819                    let mut acc2 = simd.splat_f64s(0.0);
820                    let mut acc3 = simd.splat_f64s(0.0);
821
822                    let mut i = 0usize;
823                    while i + 4 <= head.len() {
824                        let square0 = simd.mul_f64s(head[i], head[i]);
825                        let square1 = simd.mul_f64s(head[i + 1], head[i + 1]);
826                        let square2 = simd.mul_f64s(head[i + 2], head[i + 2]);
827                        let square3 = simd.mul_f64s(head[i + 3], head[i + 3]);
828                        acc0 = simd.add_f64s(acc0, square0);
829                        acc1 = simd.add_f64s(acc1, square1);
830                        acc2 = simd.add_f64s(acc2, square2);
831                        acc3 = simd.add_f64s(acc3, square3);
832                        i += 4;
833                    }
834                    for &value in &head[i..] {
835                        let square = simd.mul_f64s(value, value);
836                        acc0 = simd.add_f64s(acc0, square);
837                    }
838
839                    let acc = simd.add_f64s(simd.add_f64s(acc0, acc1), simd.add_f64s(acc2, acc3));
840                    let mut sum = simd.reduce_sum_f64s(acc);
841                    for &value in tail {
842                        let square = value * value;
843                        sum += square;
844                    }
845                    sum
846                }
847            }
848
849            Some(pulp::Arch::new().dispatch(SumSquares(src)))
850        }
851    }
852
853    fn try_simd_dot_f64(a: &[f64], b: &[f64]) -> Option<f64> {
854        struct Dot<'a> {
855            a: &'a [f64],
856            b: &'a [f64],
857        }
858        impl<'a> WithSimd for Dot<'a> {
859            type Output = f64;
860
861            #[inline(always)]
862            fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
863                debug_assert_eq!(self.a.len(), self.b.len());
864                let (a_head, a_tail) = S::as_simd_f64s(self.a);
865                let (b_head, b_tail) = S::as_simd_f64s(self.b);
866                debug_assert_eq!(a_head.len(), b_head.len());
867                debug_assert_eq!(a_tail.len(), b_tail.len());
868
869                let mut acc0 = simd.splat_f64s(0.0);
870                let mut acc1 = simd.splat_f64s(0.0);
871                let mut acc2 = simd.splat_f64s(0.0);
872                let mut acc3 = simd.splat_f64s(0.0);
873
874                let mut i = 0usize;
875                while i + 4 <= a_head.len() {
876                    acc0 = simd.mul_add_f64s(a_head[i], b_head[i], acc0);
877                    acc1 = simd.mul_add_f64s(a_head[i + 1], b_head[i + 1], acc1);
878                    acc2 = simd.mul_add_f64s(a_head[i + 2], b_head[i + 2], acc2);
879                    acc3 = simd.mul_add_f64s(a_head[i + 3], b_head[i + 3], acc3);
880                    i += 4;
881                }
882                for j in i..a_head.len() {
883                    acc0 = simd.mul_add_f64s(a_head[j], b_head[j], acc0);
884                }
885
886                let acc = simd.add_f64s(simd.add_f64s(acc0, acc1), simd.add_f64s(acc2, acc3));
887                let mut sum = simd.reduce_sum_f64s(acc);
888                for (&x, &y) in a_tail.iter().zip(b_tail.iter()) {
889                    sum += x * y;
890                }
891                sum
892            }
893        }
894
895        Some(pulp::Arch::new().dispatch(Dot { a, b }))
896    }
897
898    impl MaybeSimdBlockFold for f64 {
899        fn try_simd_add_block(acc: &mut [f64], values: &[f64]) -> bool {
900            if acc.len() < BLOCK_FOLD_MIN_LEN {
901                return false;
902            }
903            struct AddBlock<'a, 'b>(&'a mut [f64], &'b [f64]);
904            impl pulp::WithSimd for AddBlock<'_, '_> {
905                type Output = ();
906                #[inline(always)]
907                fn with_simd<S: Simd>(self, simd: S) {
908                    let (ah, at) = S::as_mut_simd_f64s(self.0);
909                    let (vh, vt) = S::as_simd_f64s(self.1);
910                    // Equal block lengths are the trait contract, so the vector
911                    // prefix and the scalar tail cover the same indices.
912                    for (a, &v) in ah.iter_mut().zip(vh) {
913                        *a = simd.add_f64s(*a, v);
914                    }
915                    for (a, &v) in at.iter_mut().zip(vt) {
916                        *a += v;
917                    }
918                }
919            }
920            pulp::Arch::new().dispatch(AddBlock(acc, values));
921            true
922        }
923
924        fn try_simd_mul_block(acc: &mut [f64], values: &[f64]) -> bool {
925            if acc.len() < BLOCK_FOLD_MIN_LEN {
926                return false;
927            }
928            struct MulBlock<'a, 'b>(&'a mut [f64], &'b [f64]);
929            impl pulp::WithSimd for MulBlock<'_, '_> {
930                type Output = ();
931                #[inline(always)]
932                fn with_simd<S: Simd>(self, simd: S) {
933                    let (ah, at) = S::as_mut_simd_f64s(self.0);
934                    let (vh, vt) = S::as_simd_f64s(self.1);
935                    for (a, &v) in ah.iter_mut().zip(vh) {
936                        *a = simd.mul_f64s(*a, v);
937                    }
938                    for (a, &v) in at.iter_mut().zip(vt) {
939                        *a *= v;
940                    }
941                }
942            }
943            pulp::Arch::new().dispatch(MulBlock(acc, values));
944            true
945        }
946
947        fn try_simd_sum_squares_block(acc: &mut [f64], values: &[f64]) -> bool {
948            if acc.len() < BLOCK_FOLD_MIN_LEN {
949                return false;
950            }
951            struct SumSquaresBlock<'a, 'b>(&'a mut [f64], &'b [f64]);
952            impl pulp::WithSimd for SumSquaresBlock<'_, '_> {
953                type Output = ();
954                #[inline(always)]
955                fn with_simd<S: Simd>(self, simd: S) {
956                    let (ah, at) = S::as_mut_simd_f64s(self.0);
957                    let (vh, vt) = S::as_simd_f64s(self.1);
958                    for (a, &v) in ah.iter_mut().zip(vh) {
959                        let sq = simd.mul_f64s(v, v);
960                        *a = simd.add_f64s(*a, sq);
961                    }
962                    for (a, &v) in at.iter_mut().zip(vt) {
963                        *a += v * v;
964                    }
965                }
966            }
967            pulp::Arch::new().dispatch(SumSquaresBlock(acc, values));
968            true
969        }
970    }
971}
972
973#[cfg(test)]
974#[path = "simd/tests/tests.rs"]
975mod tests;