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
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

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use half::f16;

use super::CpuTensor;
use crate::tensor::TensorDeviceMetrics;

#[derive(Debug, Clone, Default)]
pub struct CpuTensorDeviceOptions {
    /// when enabled, whenever tensor called with `with_name`, the name and the
    /// tensor will be recorded in the device. only used in test.
    pub debug_named_tensors: bool,
}

#[derive(Debug)]
pub struct CpuTensorDevice<'a> {
    pub(crate) opts: CpuTensorDeviceOptions,
    pub(crate) metrics: TensorDeviceMetrics,
    pub(crate) debug_tensors: RefCell<HashMap<String, Vec<f32>>>,
    pub(crate) wbuf: RefCell<Option<Vec<f32>>>,
    pub(crate) exp_cache: Vec<f16>,
    _phantom: std::marker::PhantomData<&'a ()>,
}

pub type CpuTensorDeviceRef<'a> = Rc<CpuTensorDevice<'a>>;

impl<'a> CpuTensorDevice<'a> {
    pub fn new() -> CpuTensorDeviceRef<'a> {
        let device = Self {
            opts: CpuTensorDeviceOptions::default(),
            debug_tensors: RefCell::new(HashMap::new()),
            metrics: TensorDeviceMetrics::default(),
            wbuf: RefCell::new(Some(vec![0.0; 32000])),
            exp_cache: Self::init_exp_cache(),
            _phantom: std::marker::PhantomData,
        };
        Rc::new(device)
    }

    pub fn with_options(opts: CpuTensorDeviceOptions) -> CpuTensorDeviceRef<'a> {
        let device = Self {
            opts,
            debug_tensors: RefCell::new(HashMap::new()),
            metrics: TensorDeviceMetrics::default(),
            wbuf: RefCell::new(Some(vec![0.0; 32000])),
            exp_cache: Self::init_exp_cache(),
            _phantom: std::marker::PhantomData,
        };
        Rc::new(device)
    }

    pub fn with_metrics(self: Rc<Self>, metrics: TensorDeviceMetrics) -> CpuTensorDeviceRef<'a> {
        let device = Self {
            opts: self.opts.clone(),
            debug_tensors: self.debug_tensors.clone(),
            wbuf: self.wbuf.clone(),
            exp_cache: self.exp_cache.clone(),
            metrics,
            _phantom: std::marker::PhantomData,
        };
        Rc::new(device)
    }

    pub fn metrics(&self) -> &TensorDeviceMetrics {
        &self.metrics
    }

    pub fn dump_debug_tensor(&self, name: &str) -> Option<Vec<f32>> {
        self.debug_tensors.borrow().get(name).cloned()
    }

    fn init_exp_cache() -> Vec<f16> {
        (0..65536)
            .map(|x| {
                let exp32 = f16::from_bits(x as u16).to_f32().exp();
                f16::from_f32(exp32)
            })
            .collect()
    }

    pub(crate) fn add_debug_tensor(&self, tensor: &CpuTensor<'a>) {
        let buf = tensor.buf().iter_f32().collect::<Vec<_>>();
        self.debug_tensors
            .borrow_mut()
            .insert(tensor.name.clone().unwrap(), buf);
    }
}