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

strided_basic/
simd.rs
