custos_math/ops/
col_op.rs

1use crate::{cpu::col_op, Matrix};
2use custos::{number::Number, Device, MainMemory, CPU};
3
4#[cfg(feature = "opencl")]
5use super::cl_to_cpu_lr;
6#[cfg(feature = "opencl")]
7use custos::OpenCL;
8
9pub trait ColOp<T, D: Device = Self>: Device {
10    fn add_col(&self, lhs: &Matrix<T, D>, rhs: &Matrix<T, D>) -> Matrix<T, Self>;
11    fn sub_col(&self, lhs: &Matrix<T, D>, rhs: &Matrix<T, D>) -> Matrix<T, Self>;
12    fn div_col(&self, lhs: &Matrix<T, D>, rhs: &Matrix<T, D>) -> Matrix<T, Self>;
13}
14
15#[cfg(feature = "cpu")]
16impl<T: Number, D: MainMemory> ColOp<T, D> for CPU {
17    #[inline]
18    fn add_col(&self, lhs: &Matrix<T, D>, rhs: &Matrix<T, D>) -> Matrix<T> {
19        col_op(self, lhs, rhs, |c, a, b| *c = a + b)
20    }
21
22    #[inline]
23    fn sub_col(&self, lhs: &Matrix<T, D>, rhs: &Matrix<T, D>) -> Matrix<T> {
24        col_op(self, lhs, rhs, |c, a, b| *c = a - b)
25    }
26
27    #[inline]
28    fn div_col(&self, lhs: &Matrix<T, D>, rhs: &Matrix<T, D>) -> Matrix<T> {
29        col_op(self, lhs, rhs, |c, a, b| *c = a / b)
30    }
31}
32
33#[cfg(feature = "opencl")]
34impl<T: custos::CDatatype> ColOp<T> for OpenCL {
35    #[inline]
36    fn add_col(&self, lhs: &Matrix<T, Self>, rhs: &Matrix<T, Self>) -> Matrix<T, Self> {
37        cl_to_cpu_lr(self, lhs, rhs, |device, lhs, rhs| device.add_col(lhs, rhs))
38    }
39
40    #[inline]
41    fn sub_col(&self, lhs: &Matrix<T, Self>, rhs: &Matrix<T, Self>) -> Matrix<T, Self> {
42        cl_to_cpu_lr(self, lhs, rhs, |device, lhs, rhs| device.sub_col(lhs, rhs))
43    }
44
45    #[inline]
46    fn div_col(&self, lhs: &Matrix<T, Self>, rhs: &Matrix<T, Self>) -> Matrix<T, Self> {
47        cl_to_cpu_lr(self, lhs, rhs, |device, lhs, rhs| device.div_col(lhs, rhs))
48    }
49}
50#[cfg(feature = "cuda")]
51use crate::cu_to_cpu_lr;
52#[cfg(feature = "cuda")]
53use custos::CUDA;
54
55#[cfg(feature = "cuda")]
56impl<T: custos::CDatatype> ColOp<T> for CUDA {
57    #[inline]
58    fn add_col(&self, lhs: &Matrix<T, Self>, rhs: &Matrix<T, Self>) -> Matrix<T, Self> {
59        cu_to_cpu_lr(self, lhs, rhs, |device, lhs, rhs| device.add_col(lhs, rhs))
60    }
61
62    #[inline]
63    fn sub_col(&self, lhs: &Matrix<T, Self>, rhs: &Matrix<T, Self>) -> Matrix<T, Self> {
64        cu_to_cpu_lr(self, lhs, rhs, |device, lhs, rhs| device.sub_col(lhs, rhs))
65    }
66
67    #[inline]
68    fn div_col(&self, lhs: &Matrix<T, Self>, rhs: &Matrix<T, Self>) -> Matrix<T, Self> {
69        cu_to_cpu_lr(self, lhs, rhs, |device, lhs, rhs| device.div_col(lhs, rhs))
70    }
71}