himada-integrations 0.1.0

Himada framework integrations — Candle, Burn, ONNX Runtime, ndarray
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
100
101
102

use std::cell::RefCell;

use himada_dispatch::*;
use ndarray::{Array1, Array2};

pub trait HimadaArray1 {
    fn himada_dot(&self, other: &Self) -> f64;
    fn himada_reduce_sum(&self) -> f64;
    fn himada_reduce_max(&self) -> f64;
    fn himada_abs_max(&self) -> f64;
    fn himada_argmax(&self) -> usize;
    fn himada_softmax(&self) -> Array1<f64>;
    fn himada_euclidean(&self, other: &Self) -> f64;
    fn himada_cosine_similarity(&self, other: &Self) -> f64;
    fn himada_negate(&self) -> Array1<f64>;
    fn himada_clamp(&self, lo: f64, hi: f64) -> Array1<f64>;
}

pub trait HimadaArray2 {
    fn himada_matmul(&self, other: &Self) -> Array2<f64>;
}

thread_local! {
    static DOT: RefCell<Dispatch<DotKernel>> = RefCell::new(Dispatch::new("ndarray_dot", vec![]));
    static REDUCE_SUM: RefCell<Dispatch<ReduceKernel>> = RefCell::new(Dispatch::new("ndarray_reduce_sum", vec![]));
    static REDUCE_MAX: RefCell<Dispatch<ReduceKernel>> = RefCell::new(Dispatch::new("ndarray_reduce_max", vec![]));
    static ABS_MAX: RefCell<Dispatch<ReduceKernel>> = RefCell::new(Dispatch::new("ndarray_abs_max", vec![]));
    static ARGMAX: RefCell<Dispatch<ArgmaxKernel>> = RefCell::new(Dispatch::new("ndarray_argmax", vec![]));
    static SOFTMAX: RefCell<Dispatch<SoftmaxKernel>> = RefCell::new(Dispatch::new("ndarray_softmax", vec![]));
    static EUCLIDEAN: RefCell<Dispatch<DotKernel>> = RefCell::new(Dispatch::new("ndarray_euclidean", vec![]));
    static COSINE: RefCell<Dispatch<DotKernel>> = RefCell::new(Dispatch::new("ndarray_cosine", vec![]));
    static NEGATE: RefCell<Dispatch<SoftmaxKernel>> = RefCell::new(Dispatch::new("ndarray_negate", vec![]));
    static CLAMP: RefCell<Dispatch<ClampKernel>> = RefCell::new(Dispatch::new("ndarray_clamp", vec![]));
    static MATMUL: RefCell<Dispatch<MatMulKernel>> = RefCell::new(Dispatch::new("ndarray_matmul", vec![]));
}

impl HimadaArray1 for Array1<f64> {
    fn himada_dot(&self, other: &Self) -> f64 {
        DOT.with(|d| d.borrow_mut().compute(self.as_slice().unwrap(), other.as_slice().unwrap()))
    }

    fn himada_reduce_sum(&self) -> f64 {
        REDUCE_SUM.with(|d| d.borrow_mut().compute(self.as_slice().unwrap()))
    }

    fn himada_reduce_max(&self) -> f64 {
        REDUCE_MAX.with(|d| d.borrow_mut().compute(self.as_slice().unwrap()))
    }

    fn himada_abs_max(&self) -> f64 {
        ABS_MAX.with(|d| d.borrow_mut().compute(self.as_slice().unwrap()))
    }

    fn himada_argmax(&self) -> usize {
        ARGMAX.with(|d| d.borrow_mut().compute(self.as_slice().unwrap()))
    }

    fn himada_softmax(&self) -> Array1<f64> {
        let mut out = vec![0.0; self.len()];
        SOFTMAX.with(|d| d.borrow_mut().compute(self.as_slice().unwrap(), &mut out));
        Array1::from_vec(out)
    }

    fn himada_euclidean(&self, other: &Self) -> f64 {
        EUCLIDEAN.with(|d| d.borrow_mut().compute(self.as_slice().unwrap(), other.as_slice().unwrap()))
    }

    fn himada_cosine_similarity(&self, other: &Self) -> f64 {
        COSINE.with(|d| d.borrow_mut().compute(self.as_slice().unwrap(), other.as_slice().unwrap()))
    }

    fn himada_negate(&self) -> Array1<f64> {
        let mut out = vec![0.0; self.len()];
        NEGATE.with(|d| d.borrow_mut().compute(self.as_slice().unwrap(), &mut out));
        Array1::from_vec(out)
    }

    fn himada_clamp(&self, lo: f64, hi: f64) -> Array1<f64> {
        let mut out = vec![0.0; self.len()];
        CLAMP.with(|d| d.borrow_mut().compute(self.as_slice().unwrap(), lo, hi, &mut out));
        Array1::from_vec(out)
    }
}

impl HimadaArray2 for Array2<f64> {
    fn himada_matmul(&self, other: &Self) -> Array2<f64> {
        let n = self.nrows();
        assert_eq!(n, self.ncols());
        assert_eq!(n, other.nrows());
        assert_eq!(n, other.ncols());
        let mut c = vec![0.0; n * n];
        MATMUL.with(|d| {
            d.borrow_mut().compute(
                self.as_slice().unwrap(),
                other.as_slice().unwrap(),
                &mut c,
                n,
            )
        });
        Array2::from_shape_vec((n, n), c).unwrap()
    }
}