crabml 0.1.0

crabml core package
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

use std::simd::f32x32;

use crate::backends::cpu::buf::CpuTensorBuf;
use crate::error::Result;
use crate::tensor::TensorStrider;

// both buf1 and buf2 have to be owned, the dtype should be the same
pub fn mul_inplace<'a>(
    buf1: &mut CpuTensorBuf<'a>,
    buf2: &CpuTensorBuf<'a>,
    strider1: &TensorStrider,
    strider2: &TensorStrider,
) -> Result<()> {
    assert!(buf1.is_owned());
    assert!(strider1.shape() == strider2.shape());
    assert!(strider1.is_contiguous());
    assert!(strider2.is_contiguous());
    assert!(buf1.dtype() == buf2.dtype());

    for (ia, ib) in buf1.iter_f32_mut().zip(buf2.iter_f32()) {
        *ia *= ib;
    }

    Ok(())
}

#[allow(unused)]
fn mul_inplace_vec_f32(a: &mut [f32], b: &[f32]) {
    let ac = a.as_chunks_mut::<32>().0;
    let bc = b.as_chunks::<32>().0;
    ac.iter_mut().zip(bc).for_each(|(a, b)| {
        let mut va = f32x32::from_slice(a);
        let vb = f32x32::from_slice(b);
        va *= vb;
        va.copy_to_slice(a);
    });
}