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}