ruda_tensor/ops/
sparse.rs

1use crate::{Backend, TensorMetadata, tensor::FloatTensor};
2use core::fmt::{Debug, Display};
3use alloc::vec::Vec;
4
5#[derive(Clone, Debug)]
6pub struct CsrAddition<H> {
7    pub left: H,
8    pub right: H,
9    pub output: H,
10    pub left_entries: Vec<u32>,
11    pub right_entries: Vec<u32>,
12}
13
14pub trait SparseOps: Backend {
15    type CsrHandle: Clone + Debug + Send + 'static;
16    type CsrData: Clone + Debug + Send + 'static;
17    type SparseError: Debug + Display;
18
19    fn csr_from_data(data: &Self::CsrData, device: &Self::Device) -> Result<Self::CsrHandle, Self::SparseError>;
20    fn csr_to_device(matrix: &Self::CsrHandle, device: &Self::Device) -> Self::CsrHandle;
21    fn csr_transpose_with_permutation(
22        matrix: &Self::CsrHandle,
23    ) -> Result<(Self::CsrHandle, Vec<u32>), Self::SparseError>;
24    async fn csr_to_data(
25        matrix: &Self::CsrHandle,
26        values: FloatTensor<Self>,
27    ) -> Result<Self::CsrData, Self::SparseError>;
28
29    fn csr_shape(matrix: &Self::CsrHandle) -> [usize; 2];
30    fn csr_nnz(matrix: &Self::CsrHandle) -> usize;
31    fn csr_add_prepare(
32        left: &Self::CsrHandle, right: &Self::CsrHandle,
33    ) -> Result<CsrAddition<Self::CsrHandle>, Self::SparseError>;
34    fn csr_add(
35        plan: &CsrAddition<Self::CsrHandle>, left: FloatTensor<Self>, right: FloatTensor<Self>,
36        alpha: f32, beta: f32,
37    ) -> Result<FloatTensor<Self>, Self::SparseError>;
38    fn csr_product_pattern(
39        left: &Self::CsrHandle, right: &Self::CsrHandle,
40    ) -> Result<Self::CsrHandle, Self::SparseError>;
41    fn csr_sampled_sparse_matmul(
42        pattern: &Self::CsrHandle,
43        left: &Self::CsrHandle, left_values: FloatTensor<Self>,
44        right: &Self::CsrHandle, right_values: FloatTensor<Self>,
45        transpose_left: bool, transpose_right: bool,
46    ) -> Result<FloatTensor<Self>, Self::SparseError>;
47    fn csr_validate_operand<T: TensorMetadata>(
48        matrix: &Self::CsrHandle,
49        operand: &T,
50        device: &Self::Device,
51        shape: &[usize],
52    ) -> Result<(), Self::SparseError>;
53    fn csr_values(matrix: &Self::CsrHandle) -> FloatTensor<Self>;
54    fn csr_validate_values(
55        matrix: &Self::CsrHandle,
56        values: &FloatTensor<Self>,
57    ) -> Result<(), Self::SparseError>;
58
59    fn csr_gather(
60        matrix: &Self::CsrHandle, dense: FloatTensor<Self>,
61    ) -> Result<FloatTensor<Self>, Self::SparseError>;
62
63    fn csr_scatter_add(
64        matrix: &Self::CsrHandle, values: FloatTensor<Self>,
65    ) -> Result<FloatTensor<Self>, Self::SparseError>;
66
67    fn csr_to_dense(
68        matrix: &Self::CsrHandle,
69        values: FloatTensor<Self>,
70    ) -> Result<FloatTensor<Self>, Self::SparseError>;
71
72    fn csr_to_dense_backward(
73        matrix: &Self::CsrHandle,
74        grad: FloatTensor<Self>,
75    ) -> Result<FloatTensor<Self>, Self::SparseError>;
76
77    fn csr_matmul(
78        matrix: &Self::CsrHandle,
79        values: FloatTensor<Self>,
80        rhs: FloatTensor<Self>,
81        transpose: bool,
82    ) -> Result<FloatTensor<Self>, Self::SparseError>;
83
84    fn csr_sampled_matmul(
85        matrix: &Self::CsrHandle,
86        lhs: FloatTensor<Self>,
87        rhs: FloatTensor<Self>,
88    ) -> Result<FloatTensor<Self>, Self::SparseError>;
89}
ruda_tensor/ops/sparse.rs

ruda_tensor/ops/
sparse.rs
