ruda_tensor/api/
sparse.rs

1use super::{Tensor, TensorPrimitive};
2use crate::{DType, TensorData, ops::SparseOps, tensor::{FloatTensor, Int}};
3
4#[derive(Clone, Debug)]
5pub struct CsrTensor<B: SparseOps> {
6    handle: B::CsrHandle,
7    values: Tensor<B, 1>,
8}
9
10impl<B: SparseOps> CsrTensor<B> {
11    pub fn from_data(data: &B::CsrData, device: &B::Device) -> Result<Self, B::SparseError> {
12        B::csr_from_data(data, device).map(Self::from_handle)
13    }
14
15    pub fn from_handle(handle: B::CsrHandle) -> Self {
16        let values = Tensor::from_primitive(TensorPrimitive::Float(B::csr_values(&handle)));
17        Self { handle, values }
18    }
19
20    pub fn from_parts(handle: B::CsrHandle, values: Tensor<B, 1>) -> Result<Self, B::SparseError> {
21        B::csr_validate_values(&handle, &unquantized(values.clone()))?;
22        Ok(Self { handle, values })
23    }
24
25    pub fn into_parts(self) -> (B::CsrHandle, Tensor<B, 1>) {
26        (self.handle, self.values)
27    }
28
29    pub fn shape(&self) -> [usize; 2] {
30        B::csr_shape(&self.handle)
31    }
32
33    pub fn values(&self) -> Tensor<B, 1> {
34        self.values.clone()
35    }
36
37    pub fn device(&self) -> B::Device {
38        self.values.device()
39    }
40
41    pub fn to_device(&self, device: &B::Device) -> Self {
42        Self {
43            handle: B::csr_to_device(&self.handle, device),
44            values: self.values.clone().to_device(device),
45        }
46    }
47
48    pub async fn to_data(&self) -> Result<B::CsrData, B::SparseError> {
49        B::csr_to_data(&self.handle, unquantized(self.values.clone())).await
50    }
51
52    pub fn with_values(&self, values: Tensor<B, 1>) -> Result<Self, B::SparseError> {
53        let primitive = unquantized(values.clone());
54        B::csr_validate_values(&self.handle, &primitive)?;
55        Ok(Self { handle: self.handle.clone(), values })
56    }
57
58    pub fn matmul(&self, rhs: Tensor<B, 2>) -> Result<Tensor<B, 2>, B::SparseError> {
59        self.matmul_impl(rhs, false)
60    }
61
62    pub fn gather(&self, dense: Tensor<B, 2>) -> Result<Tensor<B, 1>, B::SparseError> {
63        let output = B::csr_gather(&self.handle, unquantized(dense))?;
64        Ok(Tensor::from_primitive(TensorPrimitive::Float(output)))
65    }
66
67    pub fn scatter_add(&self) -> Result<Tensor<B, 2>, B::SparseError> {
68        let output = B::csr_scatter_add(&self.handle, unquantized(self.values.clone()))?;
69        Ok(Tensor::from_primitive(TensorPrimitive::Float(output)))
70    }
71
72    pub fn mul_dense(&self, rhs: Tensor<B, 2>) -> Result<Self, B::SparseError> {
73        self.with_values(self.values.clone() * self.gather(rhs)?)
74    }
75
76    pub fn to_dense(&self) -> Result<Tensor<B, 2>, B::SparseError> {
77        let output = B::csr_to_dense(&self.handle, unquantized(self.values.clone()))?;
78        Ok(Tensor::from_primitive(TensorPrimitive::Float(output)))
79    }
80
81    pub fn transpose_matmul(&self, rhs: Tensor<B, 2>) -> Result<Tensor<B, 2>, B::SparseError> {
82        self.matmul_impl(rhs, true)
83    }
84
85    pub fn transpose(&self) -> Result<Self, B::SparseError> {
86        let (handle, permutation) = B::csr_transpose_with_permutation(&self.handle)?;
87        let nnz = permutation.len();
88        let indices = Tensor::<B, 1, Int>::from_data(
89            TensorData::new(permutation, [nnz]), (&self.device(), DType::I64),
90        );
91        Self::from_parts(handle, self.values.clone().select(0, indices))
92    }
93
94    pub fn add(&self, rhs: &Self) -> Result<Self, B::SparseError> {
95        self.add_scaled(rhs, 1.0, 1.0)
96    }
97
98    pub fn sparse_matmul(&self, rhs: &Self) -> Result<Self, B::SparseError> {
99        let handle = B::csr_product_pattern(&self.handle, &rhs.handle)?;
100        let values = B::csr_sampled_sparse_matmul(
101            &handle, &self.handle, unquantized(self.values.clone()),
102            &rhs.handle, unquantized(rhs.values.clone()), false, false,
103        )?;
104        Self::from_parts(handle, Tensor::from_primitive(TensorPrimitive::Float(values)))
105    }
106
107    pub fn sampled_sparse_matmul(&self, lhs: &Self, rhs: &Self) -> Result<Self, B::SparseError> {
108        let values = B::csr_sampled_sparse_matmul(
109            &self.handle, &lhs.handle, unquantized(lhs.values.clone()),
110            &rhs.handle, unquantized(rhs.values.clone()), false, false,
111        )?;
112        Self::from_parts(self.handle.clone(), Tensor::from_primitive(TensorPrimitive::Float(values)))
113    }
114
115    pub fn add_scaled(&self, rhs: &Self, alpha: f32, beta: f32) -> Result<Self, B::SparseError> {
116        let plan = B::csr_add_prepare(&self.handle, &rhs.handle)?;
117        let values = B::csr_add(&plan, unquantized(self.values.clone()), unquantized(rhs.values.clone()), alpha, beta)?;
118        Self::from_parts(plan.output, Tensor::from_primitive(TensorPrimitive::Float(values)))
119    }
120
121    fn matmul_impl(&self, rhs: Tensor<B, 2>, transpose: bool) -> Result<Tensor<B, 2>, B::SparseError> {
122        let output = B::csr_matmul(
123            &self.handle, unquantized(self.values.clone()), unquantized(rhs), transpose,
124        )?;
125        Ok(Tensor::from_primitive(TensorPrimitive::Float(output)))
126    }
127
128    pub fn sampled_matmul(
129        &self,
130        lhs: Tensor<B, 2>,
131        rhs: Tensor<B, 2>,
132    ) -> Result<Self, B::SparseError> {
133        let values = B::csr_sampled_matmul(&self.handle, unquantized(lhs), unquantized(rhs))?;
134        Ok(Self {
135            handle: self.handle.clone(),
136            values: Tensor::from_primitive(TensorPrimitive::Float(values)),
137        })
138    }
139}
140
141fn unquantized<B: SparseOps, const D: usize>(tensor: Tensor<B, D>) -> FloatTensor<B> {
142    match tensor.into_primitive() {
143        TensorPrimitive::Float(tensor) => tensor,
144        TensorPrimitive::QFloat(_) => panic!("Sparse matrix operations require unquantized tensors"),
145    }
146}
ruda_tensor/api/sparse.rs

ruda_tensor/api/
sparse.rs
