lak-kernels 0.1.0

BLAS-like linear algebra kernels in fully-safe Rust.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

// asum.rs

use std::ops::{Add, AddAssign};
use std::simd::{Simd, SimdElement}; 
use std::simd::num::SimdFloat; 

use crate::types::VecRef; 
use crate::traits::Abs; 

const LANES: usize = 32; 

/// sum of the absolute value of all elements in x 
///
/// args: 
/// * x: [VecRef] - vector x 
///
/// returns: 
/// * T - sum of absolute values elements in x 
pub fn asum<T>( 
    x: VecRef<'_, T>, 
) -> T 
where 
    T: SimdElement 
        + Copy 
        + Default 
        + AddAssign 
        + Abs 
        + Add<Output=T>, 

    Simd<T, LANES>: SimdFloat<Scalar=T> 
        + AddAssign
        + Add<Output=Simd<T, LANES>>
{
    let x_slice = x.as_slice(); 

    let (x_chunks, x_tail) = x_slice.as_chunks::<LANES>(); 

    let mut accumulator = Simd::<T, LANES>::splat(T::default());

    for &x_chunk in x_chunks.iter() { 
        let x_vec = Simd::from_slice(&x_chunk[0..LANES]); 
        accumulator += x_vec.abs(); 
    }

    let mut sum = T::default(); 
    for &xt in x_tail.iter() { 
        sum += xt.abs(); 
    }

    accumulator.reduce_sum() + sum
}