custos_math/ops/
fns.rs

1use custos::{impl_stack, number::Float, CDatatype, Device, MainMemory, Shape};
2
3#[cfg(feature = "cpu")]
4use custos::CPU;
5
6#[cfg(feature = "stack")]
7use custos::Stack;
8
9use crate::{each_op, Matrix};
10
11#[cfg(feature = "cuda")]
12use crate::cu_str_op;
13#[cfg(feature = "cuda")]
14use custos::CUDA;
15
16#[cfg(feature = "opencl")]
17use crate::opencl::cl_str_op_mat;
18#[cfg(feature = "opencl")]
19use custos::OpenCL;
20
21impl<'a, T: Float, S: Shape, D: FnsOps<T, S, D>> Matrix<'a, T, D, S> {
22    #[inline]
23    pub fn exp(&self) -> Self {
24        self.device().exp(self)
25    }
26
27    #[inline]
28    pub fn ln(&self) -> Self {
29        self.device().ln(self)
30    }
31
32    #[inline]
33    pub fn neg(&self) -> Self {
34        self.device().neg(self)
35    }
36
37    #[inline]
38    pub fn powf(&self, rhs: T) -> Self {
39        self.device().powf(self, rhs)
40    }
41
42    #[inline]
43    pub fn powi(&self, rhs: i32) -> Self {
44        self.device().powi(self, rhs)
45    }
46}
47
48pub trait FnsOps<T, S: Shape = (), D: Device = Self>: Device {
49    fn exp(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
50    fn ln(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
51    fn neg(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S>;
52    fn powf(&self, x: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S>;
53    fn powi(&self, x: &Matrix<T, D, S>, rhs: i32) -> Matrix<T, Self, S>;
54}
55
56#[impl_stack]
57impl<T, D, S> FnsOps<T, S, D> for CPU
58where
59    T: Float,
60    D: MainMemory,
61    S: Shape,
62{
63    #[inline]
64    fn exp(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
65        each_op(self, x, |x| x.exp())
66    }
67
68    #[inline]
69    fn ln(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
70        each_op(self, x, |x| x.ln())
71    }
72
73    #[inline]
74    fn neg(&self, x: &Matrix<T, D, S>) -> Matrix<T, Self, S> {
75        each_op(self, x, |x| -x)
76    }
77
78    #[inline]
79    fn powf(&self, x: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S> {
80        each_op(self, x, |x| x.powf(rhs))
81    }
82
83    #[inline]
84    fn powi(&self, x: &Matrix<T, D, S>, rhs: i32) -> Matrix<T, Self, S> {
85        each_op(self, x, |x| x.powi(rhs))
86    }
87}
88
89#[cfg(feature = "opencl")]
90impl<T: CDatatype> FnsOps<T> for OpenCL {
91    #[inline]
92    fn exp(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
93        cl_str_op_mat(self, x, "exp(x)").unwrap()
94    }
95
96    #[inline]
97    fn ln(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
98        cl_str_op_mat(self, x, "log(x)").unwrap()
99    }
100
101    #[inline]
102    fn neg(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
103        cl_str_op_mat(self, x, "-x").unwrap()
104    }
105
106    #[inline]
107    fn powf(&self, x: &Matrix<T, Self>, rhs: T) -> Matrix<T, Self> {
108        cl_str_op_mat(self, x, &format!("pow(x, {rhs})")).unwrap()
109    }
110
111    #[inline]
112    fn powi(&self, x: &Matrix<T, Self>, rhs: i32) -> Matrix<T, Self> {
113        cl_str_op_mat(self, x, &format!("pow(x, {rhs})")).unwrap()
114    }
115}
116
117#[cfg(feature = "cuda")]
118impl<T: CDatatype> FnsOps<T> for CUDA {
119    #[inline]
120    fn exp(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
121        let out = cu_str_op(self, x, "exp(x)").unwrap();
122        (out, x.dims()).into()
123    }
124
125    #[inline]
126    fn ln(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
127        let out = cu_str_op(self, x, "logf(x)").unwrap();
128        (out, x.dims()).into()
129    }
130
131    #[inline]
132    fn neg(&self, x: &Matrix<T, Self>) -> Matrix<T, Self> {
133        let out = cu_str_op(self, x, "-x").unwrap();
134        (out, x.dims()).into()
135    }
136
137    #[inline]
138    fn powf(&self, x: &Matrix<T, Self>, rhs: T) -> Matrix<T, Self> {
139        let out = cu_str_op(self, x, &format!("powf(x, {rhs})")).unwrap();
140        (out, x.dims()).into()
141    }
142
143    #[inline]
144    fn powi(&self, x: &Matrix<T, Self>, rhs: i32) -> Matrix<T, Self> {
145        let out = cu_str_op(self, x, &format!("powf(x, {rhs})")).unwrap();
146        (out, x.dims()).into()
147    }
148}
149
150#[cfg(test)]
151mod tests {
152
153    #[cfg(feature = "stack")]
154    #[test]
155    fn test_stack_impl() {
156        use custos::{Buffer, Stack};
157
158        use crate::Matrix;
159
160        let data = Buffer::from((Stack, &[3., 1., 5.]));
161        let mat = Matrix { data, dims: (1, 3) };
162
163        mat.ln();
164    }
165}