acyclib 0.3.0

ML library for directed acyclic tensor graphs.
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use crate::device::{
    Device, OperationError,
    function::DeviceOperation,
    operation::{BlasOperations, GemmConfig},
    tensor::{Shape, TensorRef},
};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MatmulType {
    BatBat,
    NobBat,
    NobNob,
    BatBatRed,
}

#[derive(Clone)]
pub struct Matmul<D: Device> {
    pub cfg: GemmConfig,
    pub input_a: TensorRef<D>,
    pub input_b: TensorRef<D>,
    pub output: TensorRef<D>,
    pub ty: MatmulType,
}

impl<D: Device> DeviceOperation<D> for Matmul<D> {
    fn opname(&self) -> String {
        format!("Matmul({:?})", self.ty)
    }

    fn execute(&self) -> Result<(), OperationError<<D as Device>::DeviceError>> {
        let Matmul { cfg, input_a, input_b, output, ty } = self;

        let input_a = input_a.dense();
        let input_b = input_b.dense();
        let mut output = output.dense_mut();

        if input_a.single_size() != cfg.shape_a.size() || input_b.single_size() != cfg.shape_b.size() {
            return Err(OperationError::InvalidTensorFormat);
        }

        match ty {
            MatmulType::BatBat => {
                let bs = input_a.batch_size();

                if bs != input_b.batch_size() || bs != output.batch_size() {
                    return Err(OperationError::MismatchedBatchSizes);
                }

                output.buf.gebmm(cfg, bs.unwrap_or(1), &input_a.buf, &input_b.buf)?;
            }
            MatmulType::NobNob => {
                if input_a.batch_size().is_some() || input_b.batch_size().is_some() || output.batch_size().is_some() {
                    return Err(OperationError::MismatchedBatchSizes);
                }

                output.buf.gemm(cfg, &input_a.buf, &input_b.buf)?;
            }
            MatmulType::NobBat => {
                if input_a.batch_size().is_some()
                    || input_b.batch_size().is_none()
                    || input_b.batch_size() != output.batch_size()
                {
                    return Err(OperationError::MismatchedBatchSizes);
                }

                if cfg.trans_b {
                    return Err(OperationError::UnsupportedOperation);
                }

                let bs = input_b.batch_size().unwrap();

                let shape_b = Shape::new(cfg.shape_b.rows(), bs * cfg.shape_b.cols());
                let cfg = GemmConfig { shape_b, ..*cfg };
                output.buf.gemm(&cfg, &input_a.buf, &input_b.buf)?;
            }
            MatmulType::BatBatRed => {
                if input_a.batch_size() != input_b.batch_size() || output.batch_size().is_some() {
                    return Err(OperationError::MismatchedBatchSizes);
                }

                if !cfg.trans_b {
                    return Err(OperationError::UnsupportedOperation);
                }

                let bs = input_a.batch_size().unwrap_or(1);

                let cfg = GemmConfig {
                    shape_a: Shape::new(self.cfg.shape_a.rows(), bs * self.cfg.shape_a.cols()),
                    shape_b: Shape::new(self.cfg.shape_b.rows(), bs * self.cfg.shape_b.cols()),
                    ..self.cfg
                };

                output.buf.gemm(&cfg, &input_a.buf, &input_b.buf)?;
            }
        }

        Ok(())
    }
}