custos_math/ops/
scalar.rs

1use crate::{cpu::scalar_apply, Matrix};
2use custos::{impl_stack, number::Number, CDatatype, Device, MainMemory, Shape};
3
4#[cfg(feature = "cpu")]
5use custos::CPU;
6
7#[cfg(feature = "stack")]
8use custos::Stack;
9
10#[cfg(feature = "opencl")]
11use crate::opencl::cl_scalar_op_mat;
12#[cfg(feature = "opencl")]
13use custos::OpenCL;
14
15#[cfg(feature = "cuda")]
16use crate::cuda::cu_scalar_op;
17#[cfg(feature = "cuda")]
18use custos::CUDA;
19
20impl<'a, T, D, S> Matrix<'a, T, D, S>
21where
22    D: AdditionalOps<T, S>,
23    S: Shape,
24{
25    #[inline]
26    pub fn adds(&self, rhs: T) -> Self {
27        self.device().adds(self, rhs)
28    }
29
30    #[inline]
31    pub fn subs(&self, rhs: T) -> Self {
32        self.device().subs(self, rhs)
33    }
34
35    #[inline]
36    pub fn muls(&self, rhs: T) -> Self {
37        self.device().muls(self, rhs)
38    }
39
40    #[inline]
41    pub fn divs(&self, rhs: T) -> Self {
42        self.device().divs(self, rhs)
43    }
44}
45
46pub trait AdditionalOps<T, S: Shape = (), D: Device = Self>: Device {
47    fn adds(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S>;
48    fn subs(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S>;
49    fn muls(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S>;
50    fn divs(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S>;
51}
52
53#[cfg(feature = "cuda")]
54impl<T: CDatatype> AdditionalOps<T> for CUDA {
55    #[inline]
56    fn adds(&self, lhs: &Matrix<T, CUDA>, rhs: T) -> Matrix<T, CUDA> {
57        (cu_scalar_op(self, lhs, rhs, "+").unwrap(), lhs.dims()).into()
58    }
59
60    #[inline]
61    fn muls(&self, lhs: &Matrix<T, CUDA>, rhs: T) -> Matrix<T, CUDA> {
62        (cu_scalar_op(self, lhs, rhs, "*").unwrap(), lhs.dims()).into()
63    }
64
65    #[inline]
66    fn divs(&self, lhs: &Matrix<T, CUDA>, rhs: T) -> Matrix<T, CUDA> {
67        (cu_scalar_op(self, lhs, rhs, "/").unwrap(), lhs.dims()).into()
68    }
69
70    fn subs(&self, lhs: &Matrix<T, Self>, rhs: T) -> Matrix<T, Self, ()> {
71        (cu_scalar_op(self, lhs, rhs, "-").unwrap(), lhs.dims()).into()
72    }
73}
74
75#[cfg(feature = "opencl")]
76impl<T: CDatatype> AdditionalOps<T> for OpenCL {
77    #[inline]
78    fn adds(&self, lhs: &Matrix<T, Self>, rhs: T) -> Matrix<T, Self> {
79        cl_scalar_op_mat(self, lhs, rhs, "+").unwrap()
80    }
81
82    #[inline]
83    fn subs(&self, lhs: &Matrix<T, Self, ()>, rhs: T) -> Matrix<T, Self, ()> {
84        cl_scalar_op_mat(self, lhs, rhs, "-").unwrap()
85    }
86
87    #[inline]
88    fn muls(&self, lhs: &Matrix<T, Self>, rhs: T) -> Matrix<T, Self> {
89        cl_scalar_op_mat(self, lhs, rhs, "*").unwrap()
90    }
91
92    #[inline]
93    fn divs(&self, lhs: &Matrix<T, Self>, rhs: T) -> Matrix<T, Self> {
94        cl_scalar_op_mat(self, lhs, rhs, "/").unwrap()
95    }
96}
97
98#[impl_stack]
99impl<T: Number, D: MainMemory, S: Shape> AdditionalOps<T, S, D> for CPU {
100    #[inline]
101    fn adds(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S> {
102        scalar_apply(self, lhs, rhs, |c, a, b| *c = a + b)
103    }
104
105    #[inline]
106    fn subs(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S> {
107        scalar_apply(self, lhs, rhs, |c, a, b| *c = a - b)
108    }
109
110    #[inline]
111    fn muls(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S> {
112        scalar_apply(self, lhs, rhs, |c, a, b| *c = a * b)
113    }
114
115    #[inline]
116    fn divs(&self, lhs: &Matrix<T, D, S>, rhs: T) -> Matrix<T, Self, S> {
117        scalar_apply(self, lhs, rhs, |c, a, b| *c = a / b)
118    }
119}
custos_math/ops/scalar.rs

custos_math/ops/
scalar.rs
