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
213#[inline(always)]
214unsafe fn mul_transposed_scalar_rhs_source_contiguous<T>(
215    dst: *mut T,
216    src: *const T,
217    scalar: T,
218    inner_len: usize,
219    row_len: usize,
220    src_fast_stride: isize,
221    src_row_stride: isize,
222) where
223    T: Copy + std::ops::Mul<Output = T>,
224{
225    let mut inner0 = 0usize;
226
227    while inner0 < inner_len {
228        let inner_count = TRANSPOSE_TILE.min(inner_len - inner0);
229        let mut row0 = 0usize;
230
231        while row0 < row_len {
232            let row_count = TRANSPOSE_TILE.min(row_len - row0);
233
234            for inner in 0..inner_count {
235                let src_base = src.offset((inner0 + inner) as isize * src_fast_stride);
236                for row in 0..row_count {
237                    let row_index = row0 + row;
238                    let src_offset = row_index as isize * src_row_stride;
239                    *dst.add(row_index * inner_len + inner0 + inner) =
240                        *src_base.offset(src_offset) * scalar;
241                }
242            }
243
244            row0 += TRANSPOSE_TILE;
245        }
246
247        inner0 += TRANSPOSE_TILE;
248    }
249}
250
251#[inline(always)]
252unsafe fn mul_transposed_scalar_lhs_source_contiguous<T>(
253    dst: *mut T,
254    scalar: T,
255    src: *const T,
256    inner_len: usize,
257    row_len: usize,
258    src_fast_stride: isize,
259    src_row_stride: isize,
260) where
261    T: Copy + std::ops::Mul<Output = T>,
262{
263    let mut inner0 = 0usize;
264
265    while inner0 < inner_len {
266        let inner_count = TRANSPOSE_TILE.min(inner_len - inner0);
267        let mut row0 = 0usize;
268
269        while row0 < row_len {
270            let row_count = TRANSPOSE_TILE.min(row_len - row0);
271
272            for inner in 0..inner_count {
273                let src_base = src.offset((inner0 + inner) as isize * src_fast_stride);
274                for row in 0..row_count {
275                    let row_index = row0 + row;
276                    let src_offset = row_index as isize * src_row_stride;
277                    *dst.add(row_index * inner_len + inner0 + inner) =
278                        scalar * *src_base.offset(src_offset);
279                }
280            }
281
282            row0 += TRANSPOSE_TILE;
283        }
284
285        inner0 += TRANSPOSE_TILE;
286    }
287}
288
289#[inline]
290pub(crate) unsafe fn try_mul_transposed_scalar_rhs_2d<D: 'static, A: 'static, B: 'static>(
291    dst: *mut D,
292    src: *const A,
293    scalar: *const B,
294    inner_len: usize,
295    row_len: usize,
296    src_fast_stride: isize,
297    src_row_stride: isize,
298) -> bool {
299    use std::any::TypeId;
300
301    macro_rules! try_same_type {
302        ($ty:ty) => {
303            if TypeId::of::<D>() == TypeId::of::<$ty>()
304                && TypeId::of::<A>() == TypeId::of::<$ty>()
305                && TypeId::of::<B>() == TypeId::of::<$ty>()
306            {
307                // One runtime target-feature dispatch per tile sweep (tensor4all/strided-rs#340).
308                dispatch_if_large(inner_len * row_len, || unsafe {
309                    mul_transposed_scalar_rhs_source_contiguous(
310                        dst.cast::<$ty>(),
311                        src.cast::<$ty>(),
312                        *scalar.cast::<$ty>(),
313                        inner_len,
314                        row_len,
315                        src_fast_stride,
316                        src_row_stride,
317                    );
318                });
319                return true;
320            }
321        };
322    }
323
324    try_same_type!(f64);
325    try_same_type!(f32);
326    try_same_type!(num_complex::Complex64);
327    try_same_type!(num_complex::Complex32);
328
329    false
330}
331
332#[inline]
333pub(crate) unsafe fn try_mul_transposed_scalar_lhs_2d<D: 'static, A: 'static, B: 'static>(
334    dst: *mut D,
335    scalar: *const A,
336    src: *const B,
337    inner_len: usize,
338    row_len: usize,
339    src_fast_stride: isize,
340    src_row_stride: isize,
341) -> bool {
342    use std::any::TypeId;
343
344    macro_rules! try_same_type {
345        ($ty:ty) => {
346            if TypeId::of::<D>() == TypeId::of::<$ty>()
347                && TypeId::of::<A>() == TypeId::of::<$ty>()
348                && TypeId::of::<B>() == TypeId::of::<$ty>()
349            {
350                // One runtime target-feature dispatch per tile sweep (tensor4all/strided-rs#340).
351                dispatch_if_large(inner_len * row_len, || unsafe {
352                    mul_transposed_scalar_lhs_source_contiguous(
353                        dst.cast::<$ty>(),
354                        *scalar.cast::<$ty>(),
355                        src.cast::<$ty>(),
356                        inner_len,
357                        row_len,
358                        src_fast_stride,
359                        src_row_stride,
360                    );
361                });
362                return true;
363            }
364        };
365    }
366
367    try_same_type!(f64);
368    try_same_type!(f32);
369    try_same_type!(num_complex::Complex64);
370    try_same_type!(num_complex::Complex32);
371
372    false
373}
374
375/// Trait for types that may have SIMD-accelerated sum/dot operations.
376///
377/// Default implementations return `None` (no SIMD available).
378/// f32/f64 override these with SIMD kernels when the `simd` feature is enabled.
379pub trait MaybeSimdOps: Copy + Sized {
380    fn try_simd_sum(_src: &[Self]) -> Option<Self> {
381        None
382    }
383    fn try_simd_dot(_a: &[Self], _b: &[Self]) -> Option<Self> {
384        None
385    }
386}
387
388pub(crate) trait MaybeSimdProduct: Copy + Sized {
389    fn try_simd_product(_src: &[Self]) -> Option<Self> {
390        None
391    }
392}
393
394pub(crate) trait MaybeSimdSumSquares: Copy + Sized {
395    fn try_simd_sum_squares(_src: &[Self]) -> Option<Self> {
396        None
397    }
398}
399
400/// Minimum accumulator-block length worth a SIMD dispatch.
401#[cfg(feature = "simd")]
402const BLOCK_FOLD_MIN_LEN: usize = 16;
403
404/// SIMD fold of one unit-stride value run into adjacent accumulators.
405///
406/// This is a reduction-kernel implementation hook, not part of the public
407/// surface. Every method requires `acc.len() == values.len()`; the reduction
408/// callers always pass the same block length for both, so the vector loop and
409/// the scalar tail cover exactly the same indices. A method returns false when
410/// no SIMD path applies and leaves `acc` untouched.
411pub(crate) trait MaybeSimdBlockFold: Copy + Sized {
412    fn try_simd_add_block(_acc: &mut [Self], _values: &[Self]) -> bool {
413        false
414    }
415    fn try_simd_mul_block(_acc: &mut [Self], _values: &[Self]) -> bool {
416        false
417    }
418    fn try_simd_sum_squares_block(_acc: &mut [Self], _values: &[Self]) -> bool {
419        false
420    }
421}
422
423// Default (no-op) impls for integer types and Complex
424macro_rules! impl_no_simd {
425    ($($t:ty),*) => {
426        $(
427            impl MaybeSimdOps for $t {}
428            impl MaybeSimdProduct for $t {}
429            impl MaybeSimdSumSquares for $t {}
430            impl MaybeSimdBlockFold for $t {}
431        )*
432    };
433}
434
435impl_no_simd!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);
436
437impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdOps
438    for num_complex::Complex<T>
439{
440}
441impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdProduct
442    for num_complex::Complex<T>
443{
444}
445impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdSumSquares
446    for num_complex::Complex<T>
447{
448}
449impl<T: num_traits::Num + Copy + Clone + std::ops::Neg<Output = T>> MaybeSimdBlockFold
450    for num_complex::Complex<T>
451{
452}
453
454// f32/f64: SIMD-accelerated when feature enabled, no-op otherwise
455#[cfg(not(feature = "simd"))]
456impl MaybeSimdOps for f32 {}
457#[cfg(not(feature = "simd"))]
458impl MaybeSimdProduct for f32 {}
459#[cfg(not(feature = "simd"))]
460impl MaybeSimdSumSquares for f32 {}
461
462#[cfg(not(feature = "simd"))]
463impl MaybeSimdOps for f64 {}
464#[cfg(not(feature = "simd"))]
465impl MaybeSimdProduct for f64 {}
466#[cfg(not(feature = "simd"))]
467impl MaybeSimdSumSquares for f64 {}
468#[cfg(not(feature = "simd"))]
469impl MaybeSimdBlockFold for f32 {}
470#[cfg(not(feature = "simd"))]
471impl MaybeSimdBlockFold for f64 {}
472
473#[cfg(feature = "simd")]
474mod simd_impls {
475    use super::{
476        MaybeSimdBlockFold, MaybeSimdOps, MaybeSimdProduct, MaybeSimdSumSquares, BLOCK_FOLD_MIN_LEN,
477    };
478    use pulp::{Simd, WithSimd};
479
480    impl MaybeSimdOps for f32 {
481        fn try_simd_sum(src: &[f32]) -> Option<f32> {
482            struct Sum<'a>(&'a [f32]);
483            impl<'a> WithSimd for Sum<'a> {
484                type Output = f32;
485
486                #[inline(always)]
487                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
488                    let (head, tail) = S::as_simd_f32s(self.0);
489
490                    let mut acc0 = simd.splat_f32s(0.0);
491                    let mut acc1 = simd.splat_f32s(0.0);
492                    let mut acc2 = simd.splat_f32s(0.0);
493                    let mut acc3 = simd.splat_f32s(0.0);
494
495                    let mut i = 0usize;
496                    while i + 4 <= head.len() {
497                        acc0 = simd.add_f32s(acc0, head[i]);
498                        acc1 = simd.add_f32s(acc1, head[i + 1]);
499                        acc2 = simd.add_f32s(acc2, head[i + 2]);
500                        acc3 = simd.add_f32s(acc3, head[i + 3]);
501                        i += 4;
502                    }
503                    for &v in &head[i..] {
504                        acc0 = simd.add_f32s(acc0, v);
505                    }
506
507                    let acc = simd.add_f32s(simd.add_f32s(acc0, acc1), simd.add_f32s(acc2, acc3));
508                    let mut sum = simd.reduce_sum_f32s(acc);
509                    for &x in tail {
510                        sum += x;
511                    }
512                    sum
513                }
514            }
515
516            Some(pulp::Arch::new().dispatch(Sum(src)))
517        }
518
519        fn try_simd_dot(a: &[f32], b: &[f32]) -> Option<f32> {
520            try_simd_dot_f32(a, b)
521        }
522    }
523
524    impl MaybeSimdProduct for f32 {
525        fn try_simd_product(src: &[f32]) -> Option<f32> {
526            struct Product<'a>(&'a [f32]);
527            impl<'a> WithSimd for Product<'a> {
528                type Output = f32;
529
530                #[inline(always)]
531                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
532                    let (head, tail) = S::as_simd_f32s(self.0);
533
534                    let mut acc0 = simd.splat_f32s(1.0);
535                    let mut acc1 = simd.splat_f32s(1.0);
536                    let mut acc2 = simd.splat_f32s(1.0);
537                    let mut acc3 = simd.splat_f32s(1.0);
538
539                    let mut i = 0usize;
540                    while i + 4 <= head.len() {
541                        acc0 = simd.mul_f32s(acc0, head[i]);
542                        acc1 = simd.mul_f32s(acc1, head[i + 1]);
543                        acc2 = simd.mul_f32s(acc2, head[i + 2]);
544                        acc3 = simd.mul_f32s(acc3, head[i + 3]);
545                        i += 4;
546                    }
547                    for &value in &head[i..] {
548                        acc0 = simd.mul_f32s(acc0, value);
549                    }
550
551                    let acc = simd.mul_f32s(simd.mul_f32s(acc0, acc1), simd.mul_f32s(acc2, acc3));
552                    let mut product = simd.reduce_product_f32s(acc);
553                    for &value in tail {
554                        product *= value;
555                    }
556                    product
557                }
558            }
559
560            Some(pulp::Arch::new().dispatch(Product(src)))
561        }
562    }
563
564    impl MaybeSimdSumSquares for f32 {
565        fn try_simd_sum_squares(src: &[f32]) -> Option<f32> {
566            struct SumSquares<'a>(&'a [f32]);
567            impl<'a> WithSimd for SumSquares<'a> {
568                type Output = f32;
569
570                #[inline(always)]
571                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
572                    let (head, tail) = S::as_simd_f32s(self.0);
573                    let mut acc0 = simd.splat_f32s(0.0);
574                    let mut acc1 = simd.splat_f32s(0.0);
575                    let mut acc2 = simd.splat_f32s(0.0);
576                    let mut acc3 = simd.splat_f32s(0.0);
577
578                    let mut i = 0usize;
579                    while i + 4 <= head.len() {
580                        let square0 = simd.mul_f32s(head[i], head[i]);
581                        let square1 = simd.mul_f32s(head[i + 1], head[i + 1]);
582                        let square2 = simd.mul_f32s(head[i + 2], head[i + 2]);
583                        let square3 = simd.mul_f32s(head[i + 3], head[i + 3]);
584                        acc0 = simd.add_f32s(acc0, square0);
585                        acc1 = simd.add_f32s(acc1, square1);
586                        acc2 = simd.add_f32s(acc2, square2);
587                        acc3 = simd.add_f32s(acc3, square3);
588                        i += 4;
589                    }
590                    for &value in &head[i..] {
591                        let square = simd.mul_f32s(value, value);
592                        acc0 = simd.add_f32s(acc0, square);
593                    }
594
595                    let acc = simd.add_f32s(simd.add_f32s(acc0, acc1), simd.add_f32s(acc2, acc3));
596                    let mut sum = simd.reduce_sum_f32s(acc);
597                    for &value in tail {
598                        let square = value * value;
599                        sum += square;
600                    }
601                    sum
602                }
603            }
604
605            Some(pulp::Arch::new().dispatch(SumSquares(src)))
606        }
607    }
608
609    fn try_simd_dot_f32(a: &[f32], b: &[f32]) -> Option<f32> {
610        struct Dot<'a> {
611            a: &'a [f32],
612            b: &'a [f32],
613        }
614        impl<'a> WithSimd for Dot<'a> {
615            type Output = f32;
616
617            #[inline(always)]
618            fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
619                debug_assert_eq!(self.a.len(), self.b.len());
620                let (a_head, a_tail) = S::as_simd_f32s(self.a);
621                let (b_head, b_tail) = S::as_simd_f32s(self.b);
622                debug_assert_eq!(a_head.len(), b_head.len());
623                debug_assert_eq!(a_tail.len(), b_tail.len());
624
625                let mut acc0 = simd.splat_f32s(0.0);
626                let mut acc1 = simd.splat_f32s(0.0);
627                let mut acc2 = simd.splat_f32s(0.0);
628                let mut acc3 = simd.splat_f32s(0.0);
629
630                let mut i = 0usize;
631                while i + 4 <= a_head.len() {
632                    acc0 = simd.mul_add_f32s(a_head[i], b_head[i], acc0);
633                    acc1 = simd.mul_add_f32s(a_head[i + 1], b_head[i + 1], acc1);
634                    acc2 = simd.mul_add_f32s(a_head[i + 2], b_head[i + 2], acc2);
635                    acc3 = simd.mul_add_f32s(a_head[i + 3], b_head[i + 3], acc3);
636                    i += 4;
637                }
638                for j in i..a_head.len() {
639                    acc0 = simd.mul_add_f32s(a_head[j], b_head[j], acc0);
640                }
641
642                let acc = simd.add_f32s(simd.add_f32s(acc0, acc1), simd.add_f32s(acc2, acc3));
643                let mut sum = simd.reduce_sum_f32s(acc);
644                for (&x, &y) in a_tail.iter().zip(b_tail.iter()) {
645                    sum += x * y;
646                }
647                sum
648            }
649        }
650
651        Some(pulp::Arch::new().dispatch(Dot { a, b }))
652    }
653
654    impl MaybeSimdBlockFold for f32 {
655        fn try_simd_add_block(acc: &mut [f32], values: &[f32]) -> bool {
656            if acc.len() < BLOCK_FOLD_MIN_LEN {
657                return false;
658            }
659            struct AddBlock<'a, 'b>(&'a mut [f32], &'b [f32]);
660            impl pulp::WithSimd for AddBlock<'_, '_> {
661                type Output = ();
662                #[inline(always)]
663                fn with_simd<S: Simd>(self, simd: S) {
664                    let (ah, at) = S::as_mut_simd_f32s(self.0);
665                    let (vh, vt) = S::as_simd_f32s(self.1);
666                    // Equal block lengths are the trait contract, so the vector
667                    // prefix and the scalar tail cover the same indices.
668                    for (a, &v) in ah.iter_mut().zip(vh) {
669                        *a = simd.add_f32s(*a, v);
670                    }
671                    for (a, &v) in at.iter_mut().zip(vt) {
672                        *a += v;
673                    }
674                }
675            }
676            pulp::Arch::new().dispatch(AddBlock(acc, values));
677            true
678        }
679
680        fn try_simd_mul_block(acc: &mut [f32], values: &[f32]) -> bool {
681            if acc.len() < BLOCK_FOLD_MIN_LEN {
682                return false;
683            }
684            struct MulBlock<'a, 'b>(&'a mut [f32], &'b [f32]);
685            impl pulp::WithSimd for MulBlock<'_, '_> {
686                type Output = ();
687                #[inline(always)]
688                fn with_simd<S: Simd>(self, simd: S) {
689                    let (ah, at) = S::as_mut_simd_f32s(self.0);
690                    let (vh, vt) = S::as_simd_f32s(self.1);
691                    for (a, &v) in ah.iter_mut().zip(vh) {
692                        *a = simd.mul_f32s(*a, v);
693                    }
694                    for (a, &v) in at.iter_mut().zip(vt) {
695                        *a *= v;
696                    }
697                }
698            }
699            pulp::Arch::new().dispatch(MulBlock(acc, values));
700            true
701        }
702
703        fn try_simd_sum_squares_block(acc: &mut [f32], values: &[f32]) -> bool {
704            if acc.len() < BLOCK_FOLD_MIN_LEN {
705                return false;
706            }
707            struct SumSquaresBlock<'a, 'b>(&'a mut [f32], &'b [f32]);
708            impl pulp::WithSimd for SumSquaresBlock<'_, '_> {
709                type Output = ();
710                #[inline(always)]
711                fn with_simd<S: Simd>(self, simd: S) {
712                    let (ah, at) = S::as_mut_simd_f32s(self.0);
713                    let (vh, vt) = S::as_simd_f32s(self.1);
714                    for (a, &v) in ah.iter_mut().zip(vh) {
715                        let sq = simd.mul_f32s(v, v);
716                        *a = simd.add_f32s(*a, sq);
717                    }
718                    for (a, &v) in at.iter_mut().zip(vt) {
719                        *a += v * v;
720                    }
721                }
722            }
723            pulp::Arch::new().dispatch(SumSquaresBlock(acc, values));
724            true
725        }
726    }
727
728    impl MaybeSimdOps for f64 {
729        fn try_simd_sum(src: &[f64]) -> Option<f64> {
730            struct Sum<'a>(&'a [f64]);
731            impl<'a> WithSimd for Sum<'a> {
732                type Output = f64;
733
734                #[inline(always)]
735                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
736                    let (head, tail) = S::as_simd_f64s(self.0);
737
738                    let mut acc0 = simd.splat_f64s(0.0);
739                    let mut acc1 = simd.splat_f64s(0.0);
740                    let mut acc2 = simd.splat_f64s(0.0);
741                    let mut acc3 = simd.splat_f64s(0.0);
742
743                    let mut i = 0usize;
744                    while i + 4 <= head.len() {
745                        acc0 = simd.add_f64s(acc0, head[i]);
746                        acc1 = simd.add_f64s(acc1, head[i + 1]);
747                        acc2 = simd.add_f64s(acc2, head[i + 2]);
748                        acc3 = simd.add_f64s(acc3, head[i + 3]);
749                        i += 4;
750                    }
751                    for &v in &head[i..] {
752                        acc0 = simd.add_f64s(acc0, v);
753                    }
754
755                    let acc = simd.add_f64s(simd.add_f64s(acc0, acc1), simd.add_f64s(acc2, acc3));
756                    let mut sum = simd.reduce_sum_f64s(acc);
757                    for &x in tail {
758                        sum += x;
759                    }
760                    sum
761                }
762            }
763
764            Some(pulp::Arch::new().dispatch(Sum(src)))
765        }
766
767        fn try_simd_dot(a: &[f64], b: &[f64]) -> Option<f64> {
768            try_simd_dot_f64(a, b)
769        }
770    }
771
772    impl MaybeSimdProduct for f64 {
773        fn try_simd_product(src: &[f64]) -> Option<f64> {
774            struct Product<'a>(&'a [f64]);
775            impl<'a> WithSimd for Product<'a> {
776                type Output = f64;
777
778                #[inline(always)]
779                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
780                    let (head, tail) = S::as_simd_f64s(self.0);
781
782                    let mut acc0 = simd.splat_f64s(1.0);
783                    let mut acc1 = simd.splat_f64s(1.0);
784                    let mut acc2 = simd.splat_f64s(1.0);
785                    let mut acc3 = simd.splat_f64s(1.0);
786
787                    let mut i = 0usize;
788                    while i + 4 <= head.len() {
789                        acc0 = simd.mul_f64s(acc0, head[i]);
790                        acc1 = simd.mul_f64s(acc1, head[i + 1]);
791                        acc2 = simd.mul_f64s(acc2, head[i + 2]);
792                        acc3 = simd.mul_f64s(acc3, head[i + 3]);
793                        i += 4;
794                    }
795                    for &value in &head[i..] {
796                        acc0 = simd.mul_f64s(acc0, value);
797                    }
798
799                    let acc = simd.mul_f64s(simd.mul_f64s(acc0, acc1), simd.mul_f64s(acc2, acc3));
800                    let mut product = simd.reduce_product_f64s(acc);
801                    for &value in tail {
802                        product *= value;
803                    }
804                    product
805                }
806            }
807
808            Some(pulp::Arch::new().dispatch(Product(src)))
809        }
810    }
811
812    impl MaybeSimdSumSquares for f64 {
813        fn try_simd_sum_squares(src: &[f64]) -> Option<f64> {
814            struct SumSquares<'a>(&'a [f64]);
815            impl<'a> WithSimd for SumSquares<'a> {
816                type Output = f64;
817
818                #[inline(always)]
819                fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
820                    let (head, tail) = S::as_simd_f64s(self.0);
821                    let mut acc0 = simd.splat_f64s(0.0);
822                    let mut acc1 = simd.splat_f64s(0.0);
823                    let mut acc2 = simd.splat_f64s(0.0);
824                    let mut acc3 = simd.splat_f64s(0.0);
825
826                    let mut i = 0usize;
827                    while i + 4 <= head.len() {
828                        let square0 = simd.mul_f64s(head[i], head[i]);
829                        let square1 = simd.mul_f64s(head[i + 1], head[i + 1]);
830                        let square2 = simd.mul_f64s(head[i + 2], head[i + 2]);
831                        let square3 = simd.mul_f64s(head[i + 3], head[i + 3]);
832                        acc0 = simd.add_f64s(acc0, square0);
833                        acc1 = simd.add_f64s(acc1, square1);
834                        acc2 = simd.add_f64s(acc2, square2);
835                        acc3 = simd.add_f64s(acc3, square3);
836                        i += 4;
837                    }
838                    for &value in &head[i..] {
839                        let square = simd.mul_f64s(value, value);
840                        acc0 = simd.add_f64s(acc0, square);
841                    }
842
843                    let acc = simd.add_f64s(simd.add_f64s(acc0, acc1), simd.add_f64s(acc2, acc3));
844                    let mut sum = simd.reduce_sum_f64s(acc);
845                    for &value in tail {
846                        let square = value * value;
847                        sum += square;
848                    }
849                    sum
850                }
851            }
852
853            Some(pulp::Arch::new().dispatch(SumSquares(src)))
854        }
855    }
856
857    fn try_simd_dot_f64(a: &[f64], b: &[f64]) -> Option<f64> {
858        struct Dot<'a> {
859            a: &'a [f64],
860            b: &'a [f64],
861        }
862        impl<'a> WithSimd for Dot<'a> {
863            type Output = f64;
864
865            #[inline(always)]
866            fn with_simd<S: Simd>(self, simd: S) -> Self::Output {
867                debug_assert_eq!(self.a.len(), self.b.len());
868                let (a_head, a_tail) = S::as_simd_f64s(self.a);
869                let (b_head, b_tail) = S::as_simd_f64s(self.b);
870                debug_assert_eq!(a_head.len(), b_head.len());
871                debug_assert_eq!(a_tail.len(), b_tail.len());
872
873                let mut acc0 = simd.splat_f64s(0.0);
874                let mut acc1 = simd.splat_f64s(0.0);
875                let mut acc2 = simd.splat_f64s(0.0);
876                let mut acc3 = simd.splat_f64s(0.0);
877
878                let mut i = 0usize;
879                while i + 4 <= a_head.len() {
880                    acc0 = simd.mul_add_f64s(a_head[i], b_head[i], acc0);
881                    acc1 = simd.mul_add_f64s(a_head[i + 1], b_head[i + 1], acc1);
882                    acc2 = simd.mul_add_f64s(a_head[i + 2], b_head[i + 2], acc2);
883                    acc3 = simd.mul_add_f64s(a_head[i + 3], b_head[i + 3], acc3);
884                    i += 4;
885                }
886                for j in i..a_head.len() {
887                    acc0 = simd.mul_add_f64s(a_head[j], b_head[j], acc0);
888                }
889
890                let acc = simd.add_f64s(simd.add_f64s(acc0, acc1), simd.add_f64s(acc2, acc3));
891                let mut sum = simd.reduce_sum_f64s(acc);
892                for (&x, &y) in a_tail.iter().zip(b_tail.iter()) {
893                    sum += x * y;
894                }
895                sum
896            }
897        }
898
899        Some(pulp::Arch::new().dispatch(Dot { a, b }))
900    }
901
902    impl MaybeSimdBlockFold for f64 {
903        fn try_simd_add_block(acc: &mut [f64], values: &[f64]) -> bool {
904            if acc.len() < BLOCK_FOLD_MIN_LEN {
905                return false;
906            }
907            struct AddBlock<'a, 'b>(&'a mut [f64], &'b [f64]);
908            impl pulp::WithSimd for AddBlock<'_, '_> {
909                type Output = ();
910                #[inline(always)]
911                fn with_simd<S: Simd>(self, simd: S) {
912                    let (ah, at) = S::as_mut_simd_f64s(self.0);
913                    let (vh, vt) = S::as_simd_f64s(self.1);
914                    // Equal block lengths are the trait contract, so the vector
915                    // prefix and the scalar tail cover the same indices.
916                    for (a, &v) in ah.iter_mut().zip(vh) {
917                        *a = simd.add_f64s(*a, v);
918                    }
919                    for (a, &v) in at.iter_mut().zip(vt) {
920                        *a += v;
921                    }
922                }
923            }
924            pulp::Arch::new().dispatch(AddBlock(acc, values));
925            true
926        }
927
928        fn try_simd_mul_block(acc: &mut [f64], values: &[f64]) -> bool {
929            if acc.len() < BLOCK_FOLD_MIN_LEN {
930                return false;
931            }
932            struct MulBlock<'a, 'b>(&'a mut [f64], &'b [f64]);
933            impl pulp::WithSimd for MulBlock<'_, '_> {
934                type Output = ();
935                #[inline(always)]
936                fn with_simd<S: Simd>(self, simd: S) {
937                    let (ah, at) = S::as_mut_simd_f64s(self.0);
938                    let (vh, vt) = S::as_simd_f64s(self.1);
939                    for (a, &v) in ah.iter_mut().zip(vh) {
940                        *a = simd.mul_f64s(*a, v);
941                    }
942                    for (a, &v) in at.iter_mut().zip(vt) {
943                        *a *= v;
944                    }
945                }
946            }
947            pulp::Arch::new().dispatch(MulBlock(acc, values));
948            true
949        }
950
951        fn try_simd_sum_squares_block(acc: &mut [f64], values: &[f64]) -> bool {
952            if acc.len() < BLOCK_FOLD_MIN_LEN {
953                return false;
954            }
955            struct SumSquaresBlock<'a, 'b>(&'a mut [f64], &'b [f64]);
956            impl pulp::WithSimd for SumSquaresBlock<'_, '_> {
957                type Output = ();
958                #[inline(always)]
959                fn with_simd<S: Simd>(self, simd: S) {
960                    let (ah, at) = S::as_mut_simd_f64s(self.0);
961                    let (vh, vt) = S::as_simd_f64s(self.1);
962                    for (a, &v) in ah.iter_mut().zip(vh) {
963                        let sq = simd.mul_f64s(v, v);
964                        *a = simd.add_f64s(*a, sq);
965                    }
966                    for (a, &v) in at.iter_mut().zip(vt) {
967                        *a += v * v;
968                    }
969                }
970            }
971            pulp::Arch::new().dispatch(SumSquaresBlock(acc, values));
972            true
973        }
974    }
975}
976
977#[cfg(test)]
978#[path = "simd/tests/tests.rs"]
979mod tests;