tensor_rs/tensor_impl/lapack_tensor/
mod.rs

1pub mod compare_tensor;
2pub mod convolution;
3pub mod elemwise;
4pub mod index_slicing;
5pub mod linalg;
6pub mod reduction;
7pub mod blas_api;
8pub mod lapack_api;
9
10use crate::tensor_impl::gen_tensor::GenTensor;
11use crate::tensor_impl::lapack_tensor::blas_api::BlasAPI;
12
13macro_rules! blas_matmul {
14    ($a:ty, $b: ident) => {
15        pub fn $b(
16            x: &GenTensor<$a>,
17            y: &GenTensor<$a>,
18        ) -> GenTensor<$a> {
19            if x.size()[x.size().len()-1] != y.size()[0] {
20                panic!("matmul expect matched size {:?}, {:?}", x.size(), y.size());
21            }
22            if x.size().len() == 1 && y.size().len() == 1 {
23                panic!("Two vector have not matched size for matmul! {:?}, {:?}", x.numel(), y.numel());
24            }
25            let inner = y.size()[0];
26            let mut cap = 1;
27            let mut odim = Vec::new();
28            let mut lloop = 1;
29            let mut rloop = 1;
30            for i in 0..x.size().len()-1 {
31                cap *= x.size()[i];
32                odim.push(x.size()[i]);
33                lloop *= x.size()[i];
34            }
35            for i in 1..y.size().len() {
36                cap *= y.size()[i];
37                odim.push(y.size()[i]);
38                rloop *= y.size()[i];
39            }
40            
41            let mut ret = GenTensor::<$a>::new_move(
42                vec![0.; cap], odim);
43            
44            BlasAPI::<$a>::gemm(false, false,
45                                rloop, lloop, inner,
46                                1., y.get_data(), rloop,
47                                x.get_data(), inner,
48                                1., ret.get_data_mut(), rloop,);
49            ret
50        }
51    }
52}
53
54blas_matmul!(f32, matmul_f32);
55blas_matmul!(f64, matmul_f64);
56
57#[cfg(test)]
58mod tests {
59    use crate::tensor_impl::gen_tensor::GenTensor;
60    use super::*;
61
62    #[test]
63    fn test_matmul() {
64        let v1 = GenTensor::<f32>::new_raw(&[1., 2., 3., 4., 5., 6.], &[2, 3]);
65        let v2 = GenTensor::<f32>::new_raw(&[11., 12., 13., 14., 15., 16., 17., 18., 19.], &[3, 3]);
66        let v3 = matmul_f32(&v1, &v2);
67        let em = GenTensor::<f32>::new_raw(&[90.0, 96.0, 102.0, 216.0, 231.0, 246.0], &[2, 3]);
68        assert_eq!(v3, em);
69    }
70}