custos_math/ops/
scalar_assign.rs

1use crate::{assign_to_lhs_scalar, Matrix};
2use core::ops::{AddAssign, DivAssign, MulAssign, SubAssign};
3use custos::{impl_stack, CDatatype, Device, MainMemory, Shape, CPU};
4
5#[cfg(feature = "stack")]
6use custos::Stack;
7
8#[cfg(feature = "opencl")]
9use crate::cl_assign_scalar;
10#[cfg(feature = "opencl")]
11use custos::OpenCL;
12
13#[cfg(feature = "cuda")]
14use crate::cu_assign_scalar;
15#[cfg(feature = "cuda")]
16use custos::CUDA;
17
18impl<'a, T, S, D> AddAssign<T> for Matrix<'a, T, D, S>
19where
20    S: Shape,
21    D: ScalarAssign<T, S>,
22{
23    #[inline]
24    fn add_assign(&mut self, rhs: T) {
25        self.device().adds_assign(self, rhs)
26    }
27}
28
29impl<T, S, D> MulAssign<T> for Matrix<'_, T, D, S>
30where
31    S: Shape,
32    D: ScalarAssign<T, S>,
33{
34    #[inline]
35    fn mul_assign(&mut self, rhs: T) {
36        self.device().muls_assign(self, rhs);
37    }
38}
39
40impl<T, S, D> DivAssign<T> for Matrix<'_, T, D, S>
41where
42    S: Shape,
43    D: ScalarAssign<T, S>,
44{
45    #[inline]
46    fn div_assign(&mut self, rhs: T) {
47        self.device().divs_assign(self, rhs);
48    }
49}
50
51pub trait ScalarAssign<T, S: Shape = (), D: Device = Self>: Device {
52    fn adds_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T);
53    fn muls_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T);
54    fn divs_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T);
55    fn subs_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T);
56}
57
58#[impl_stack]
59impl<T, D: MainMemory, S: Shape> ScalarAssign<T, S, D> for CPU
60where
61    T: Copy + AddAssign + MulAssign + DivAssign + SubAssign,
62{
63    #[inline]
64    fn adds_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T) {
65        assign_to_lhs_scalar(lhs, rhs, |x, y| *x += y);
66    }
67
68    #[inline]
69    fn muls_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T) {
70        assign_to_lhs_scalar(lhs, rhs, |x, y| *x *= y);
71    }
72
73    #[inline]
74    fn divs_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T) {
75        assign_to_lhs_scalar(lhs, rhs, |x, y| *x /= y);
76    }
77
78    #[inline]
79    fn subs_assign(&self, lhs: &mut Matrix<T, D, S>, rhs: T) {
80        assign_to_lhs_scalar(lhs, rhs, |x, y| *x -= y);
81    }
82}
83
84#[cfg(feature = "opencl")]
85impl<T: CDatatype> ScalarAssign<T> for OpenCL {
86    #[inline]
87    fn adds_assign(&self, lhs: &mut Matrix<T, Self>, rhs: T) {
88        cl_assign_scalar(self, lhs, rhs, "+").unwrap();
89    }
90
91    #[inline]
92    fn muls_assign(&self, lhs: &mut Matrix<T, Self>, rhs: T) {
93        cl_assign_scalar(self, lhs, rhs, "*").unwrap();
94    }
95
96    fn divs_assign(&self, lhs: &mut Matrix<T, Self>, rhs: T) {
97        cl_assign_scalar(self, lhs, rhs, "/").unwrap();
98    }
99
100    fn subs_assign(&self, lhs: &mut Matrix<T, Self>, rhs: T) {
101        cl_assign_scalar(self, lhs, rhs, "-").unwrap();
102    }
103}
104
105#[cfg(feature = "cuda")]
106impl<T: CDatatype> ScalarAssign<T> for CUDA {
107    #[inline]
108    fn adds_assign(&self, lhs: &mut Matrix<T, CUDA>, rhs: T) {
109        cu_assign_scalar(self, lhs, rhs, "+").unwrap();
110    }
111
112    #[inline]
113    fn muls_assign(&self, lhs: &mut Matrix<T, CUDA>, rhs: T) {
114        cu_assign_scalar(self, lhs, rhs, "*").unwrap();
115    }
116
117    #[inline]
118    fn divs_assign(&self, lhs: &mut Matrix<T, CUDA>, rhs: T) {
119        cu_assign_scalar(self, lhs, rhs, "/").unwrap();
120    }
121
122    #[inline]
123    fn subs_assign(&self, lhs: &mut Matrix<T, CUDA>, rhs: T) {
124        cu_assign_scalar(self, lhs, rhs, "-").unwrap();
125    }
126}
custos_math/ops/scalar_assign.rs

custos_math/ops/
scalar_assign.rs
