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
52
53
54
55
56
57
58
59
60
61
62
63

// iamax.rs 

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

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

const LANES: usize = 16; 

/// returns [usize] index for largest absolute value element in x 
///
/// args: 
/// * x: [VecRef] 
///
/// returns: 
/// * [usize] - index for largest absolute value element in x. 
/// * 0       - if input vector empty. 
pub fn iamax<T>( 
    x: VecRef<'_, T>
) -> usize
where 
    T: SimdElement
        + Copy 
        + Default
        + PartialOrd
        + Abs,

    Simd<T, LANES>: SimdFloat<Scalar=T>
{ 
    let x_slice = x.as_slice(); 
    let (x_chunks, x_tail) = x_slice.as_chunks::<LANES>(); 

    let mut max = T::default(); 
    let mut max_idx = 0; 

    for (idx, &x_chunk) in x_chunks.iter().enumerate() { 
        let x_vec = Simd::from_array(x_chunk).abs();
        let x_vec_max = x_vec.reduce_max(); 

        if x_vec_max > max { 
            for lane in 0..LANES { 
                let x_val = x_vec[lane]; 

                if x_val > max { 
                    max = x_val; 
                    max_idx = idx * LANES + lane 
                }
            }
        }
    } 

    let simd_len = x_chunks.len() * LANES; 
    for (idx, &xv) in x_tail.iter().enumerate() { 
        let x_val = xv.abs(); 
        if x_val > max { 
            max = x_val; 
            max_idx = simd_len + idx;
        }
    }

    max_idx
}