rlst 0.10.2

A Rust native linear algebra library.
Documentation 
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

//! Implementation of matrix multiplication.
//!

use crate::{
    base_types::TransMode,
    traits::{
        accessors::{RawAccess, RawAccessMut},
        base_operations::{Shape, Stride},
        linalg::base::Gemm,
    },
};

/// Matrix mulitplication
///
/// This provides a convenient interface to BLAS3 Gemm.
///
/// Compute the product `C = alpha * transa(A) * transb(B)  + beta * C`.
/// `transa` and `transb` are allowed to have all enum values from [TransMode].
pub fn matrix_multiply<
    Item: Gemm,
    MatA: RawAccess<Item = Item> + Shape<2> + Stride<2>,
    MatB: RawAccess<Item = Item> + Shape<2> + Stride<2>,
    MatC: RawAccessMut<Item = Item> + Shape<2> + Stride<2>,
>(
    transa: TransMode,
    transb: TransMode,
    alpha: Item,
    mat_a: &MatA,
    mat_b: &MatB,
    beta: Item,
    mat_c: &mut MatC,
) {
    let m = mat_c.shape()[0];
    let n = mat_c.shape()[1];

    let a_shape = match transa {
        TransMode::NoTrans => mat_a.shape(),
        TransMode::ConjNoTrans => mat_a.shape(),
        TransMode::Trans => [mat_a.shape()[1], mat_a.shape()[0]],
        TransMode::ConjTrans => [mat_a.shape()[1], mat_a.shape()[0]],
    };

    let b_shape = match transb {
        TransMode::NoTrans => mat_b.shape(),
        TransMode::ConjNoTrans => mat_b.shape(),
        TransMode::Trans => [mat_b.shape()[1], mat_b.shape()[0]],
        TransMode::ConjTrans => [mat_b.shape()[1], mat_b.shape()[0]],
    };

    assert_eq!(m, a_shape[0], "Wrong dimension. {} != {}", m, a_shape[0]);
    assert_eq!(n, b_shape[1], "Wrong dimension. {} != {}", n, b_shape[1]);
    assert_eq!(
        a_shape[1], b_shape[0],
        "Wrong dimension. {} != {}",
        a_shape[1], b_shape[0]
    );

    let [rsa, csa] = mat_a.stride();
    let [rsb, csb] = mat_b.stride();
    let [rsc, csc] = mat_c.stride();

    <Item as Gemm>::gemm(
        transa,
        transb,
        m,
        n,
        a_shape[1],
        alpha,
        mat_a.data().unwrap(),
        rsa,
        csa,
        mat_b.data().unwrap(),
        rsb,
        csb,
        beta,
        mat_c.data_mut().unwrap(),
        rsc,
        csc,
    )
}