custos_math/ops/
assign.rs

1use custos::{impl_stack, number::Number, Alloc, Buffer, Device, MainMemory, Shape, CPU};
2
3#[cfg(feature = "stack")]
4use custos::Stack;
5
6#[cfg(any(feature = "cuda", feature = "opencl"))]
7use custos::CDatatype;
8
9#[cfg(feature = "opencl")]
10use crate::cl_tew_self;
11#[cfg(feature = "opencl")]
12use custos::OpenCL;
13
14#[cfg(feature = "cuda")]
15use crate::cu_ew_self;
16use crate::{assign_to_lhs, element_wise_op_mut, Matrix};
17
18/// Assignment operations
19/// # Examples
20#[cfg_attr(feature = "cpu", doc = "```")]
21#[cfg_attr(not(feature = "cpu"), doc = "```ignore")]
22/// use custos::{CPU, Read};
23/// use custos_math::{Matrix, AssignOps};
24///
25/// let device = CPU::new();
26/// let mut lhs = Matrix::from((&device, 2, 2, [3, 5, 4, 1]));
27/// let rhs = Matrix::from((&device, 2, 2, [1, 8, 6, 2]));
28///
29/// device.add_assign(&mut lhs, &rhs);
30/// assert_eq!(vec![4, 13, 10, 3], lhs.read());
31///
32/// device.sub_assign(&mut lhs, &rhs);
33/// assert_eq!(vec![3, 5, 4, 1], lhs.read());
34/// ```
35pub trait AssignOps<T, S: Shape = (), D: Device = Self>: Device {
36    /// Add assign
37    /// # Examples
38    #[cfg_attr(feature = "cpu", doc = "```")]
39    #[cfg_attr(not(feature = "cpu"), doc = "```ignore")]
40    /// use custos::{CPU, Read};
41    /// use custos_math::{Matrix, AssignOps};
42    ///
43    /// let device = CPU::new();
44    /// let mut lhs = Matrix::from((&device, 2, 2, [3, 5, 4, 1]));
45    /// let rhs = Matrix::from((&device, 2, 2, [1, 8, 6, 2]));
46    ///
47    /// device.add_assign(&mut lhs, &rhs);
48    /// assert_eq!(vec![4, 13, 10, 3], lhs.read());
49    ///
50    /// device.sub_assign(&mut lhs, &rhs);
51    /// assert_eq!(vec![3, 5, 4, 1], lhs.read());
52    /// ```
53    fn add_assign(&self, lhs: &mut Buffer<T, Self, S>, rhs: &Buffer<T, D, S>);
54    fn sub_assign(&self, lhs: &mut Buffer<T, Self, S>, rhs: &Buffer<T, D, S>);
55    fn mul_assign(&self, lhs: &mut Buffer<T, Self, S>, rhs: &Buffer<T, D, S>);
56}
57
58pub fn ew_op<'a, T, F, D, S, Host>(
59    device: &'a Host,
60    lhs: &Matrix<T, D, S>,
61    rhs: &Matrix<T, D, S>,
62    f: F,
63) -> Matrix<'a, T, Host, S>
64where
65    T: Copy + Default,
66    F: Fn(T, T) -> T,
67    D: MainMemory,
68    S: Shape,
69    Host: for<'b> Alloc<'b, T, S> + MainMemory,
70{
71    let mut out = device.retrieve(lhs.size(), (lhs.node.idx, rhs.node.idx));
72    element_wise_op_mut(lhs, rhs, &mut out, f);
73    (out, lhs.dims()).into()
74}
75
76#[impl_stack]
77impl<T: Number, D: MainMemory, S: Shape> AssignOps<T, S, D> for CPU {
78    #[inline]
79    fn add_assign(&self, lhs: &mut Buffer<T, Self, S>, rhs: &Buffer<T, D, S>) {
80        assign_to_lhs(lhs, rhs, |x, y| *x += y)
81    }
82
83    #[inline]
84    fn sub_assign(&self, lhs: &mut Buffer<T, Self, S>, rhs: &Buffer<T, D, S>) {
85        assign_to_lhs(lhs, rhs, |x, y| *x -= y)
86    }
87
88    #[inline]
89    fn mul_assign(&self, lhs: &mut Buffer<T, Self, S>, rhs: &Buffer<T, D, S>) {
90        assign_to_lhs(lhs, rhs, |x, y| *x *= y)
91    }
92}
93
94#[cfg(feature = "opencl")]
95impl<T: CDatatype> AssignOps<T> for OpenCL {
96    #[inline]
97    fn add_assign(&self, lhs: &mut Buffer<T, OpenCL>, rhs: &Buffer<T, OpenCL>) {
98        cl_tew_self(self, lhs, rhs, "+").unwrap()
99    }
100
101    #[inline]
102    fn sub_assign(&self, lhs: &mut Buffer<T, OpenCL>, rhs: &Buffer<T, OpenCL>) {
103        cl_tew_self(self, lhs, rhs, "-").unwrap()
104    }
105
106    #[inline]
107    fn mul_assign(&self, lhs: &mut Buffer<T, Self, ()>, rhs: &Buffer<T, Self, ()>) {
108        cl_tew_self(self, lhs, rhs, "*").unwrap()
109    }
110}
111
112#[cfg(feature = "cuda")]
113impl<T: CDatatype> AssignOps<T> for custos::CUDA {
114    #[inline]
115    fn add_assign(&self, lhs: &mut Buffer<T, custos::CUDA>, rhs: &Buffer<T, custos::CUDA>) {
116        cu_ew_self(self, lhs, rhs, "+").unwrap();
117    }
118
119    #[inline]
120    fn sub_assign(&self, lhs: &mut Buffer<T, custos::CUDA>, rhs: &Buffer<T, custos::CUDA>) {
121        cu_ew_self(self, lhs, rhs, "-").unwrap();
122    }
123
124    #[inline]
125    fn mul_assign(&self, lhs: &mut Buffer<T, Self, ()>, rhs: &Buffer<T, Self, ()>) {
126        cu_ew_self(self, lhs, rhs, "*").unwrap();
127    }
128}
custos_math/ops/assign.rs

custos_math/ops/
assign.rs
