ruSPARSE 0.1.1

Ruda's sparse-computing library, reusing shared tensors, runtimes, and device backends.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use super::*;
use crate::SparseVector;

/// A sparse vector using the same device storage and index contract as one CSR row.
#[derive(Clone, Debug)]
pub struct SparseVectorTensor<R: Runtime> {
    matrix: CsrTensor<R>,
}

impl<R: Runtime> SparseVectorTensor<R> {
    pub fn from_sparse(vector: SparseVector<'_>, device: &R::Device) -> Result<Self, SparseError> {
        let base = vector.index_base().value();
        let end = dimension(vector.nnz(), "SpVV nnz")?
            .checked_add(base)
            .ok_or(SparseError::SizeOverflow("SpVV index range"))?;
        let offsets = [base, end];
        let matrix = CsrMatrix::new(
            1,
            vector.size(),
            &offsets,
            vector.indices(),
            vector.values(),
            vector.index_base(),
        )?;
        Ok(Self {
            matrix: CsrTensor::from_csr(matrix, Operation::None, device)?,
        })
    }

    pub fn size(&self) -> usize {
        self.matrix.columns
    }
    pub fn nnz(&self) -> usize {
        self.matrix.nnz
    }
    pub fn index_base(&self) -> IndexBase {
        self.matrix.base
    }

    pub fn device(&self) -> &R::Device {
        self.matrix.device()
    }

    pub fn indices(&self) -> RudaTensor<R> {
        self.matrix.column_indices()
    }

    pub fn values(&self) -> RudaTensor<R> {
        self.matrix.values()
    }

    pub fn with_values(&self, values: RudaTensor<R>) -> Result<Self, SparseError> {
        Ok(Self { matrix: self.matrix.with_values(values)? })
    }

    pub fn to_device(&self, device: &R::Device) -> Self
    where
        R::Device: PartialEq,
    {
        Self { matrix: self.matrix.to_device(device) }
    }

    pub async fn to_sparse(&self) -> Result<crate::SparseVectorOwned, SparseError> {
        let matrix = self.matrix.to_csr().await?;
        let matrix = matrix.as_ref();
        crate::SparseVectorOwned::new(
            matrix.columns(),
            matrix.column_indices().to_vec(),
            matrix.values().to_vec(),
            matrix.index_base(),
        )
    }

    pub async fn to_runtime_via_host<S: Runtime>(
        &self,
        device: &S::Device,
    ) -> Result<SparseVectorTensor<S>, SparseError> {
        Ok(SparseVectorTensor { matrix: self.matrix.to_runtime_via_host::<S>(device).await? })
    }
}

/// Sparse/dense dot product. Real FP32 conjugation is the identity; duplicates are accumulated.
pub fn spvv<R: Runtime>(
    _operation: Operation,
    x: &SparseVectorTensor<R>,
    y: RudaTensor<R>,
) -> Result<RudaTensor<R>, SparseError> {
    let matrix = &x.matrix;
    matrix.validate_dense(&y, &[matrix.columns])?;
    matrix.grid(1)?;
    let output = empty_device_contiguous_dtype(
        matrix.values.client.clone(),
        matrix.values.device.clone(),
        [1].into(),
        DType::F32,
    );
    dot::launch::<R>(
        &matrix.values.client,
        RudaCount::Static(1, 1, 1),
        RudaDim::new_1d(32),
        matrix.indices.clone().into_array_arg(),
        matrix.values.clone().into_array_arg(),
        into_contiguous(y).into_array_arg(),
        output.clone().into_array_arg(),
        matrix.nnz as u32,
        matrix.base.value(),
        include_str!("spvv.rs").to_owned(),
    );
    Ok(output)
}

#[ruda(launch)]
fn dot(
    indices: &Array<u32>,
    values: &Array<f32>,
    y: &Array<f32>,
    output: &mut Array<f32>,
    nnz: u32,
    base: u32,
    #[comptime] _source: String,
) {
    let lane = UNIT_POS as usize;
    let mut entry = lane;
    let mut sum = 0f32;
    while entry < nnz as usize {
        sum = fma(values[entry], y[(indices[entry] - base) as usize], sum);
        entry += 32;
    }
    sum = plane_sum(sum);
    if lane == 0 {
        output[0] = sum;
    }
}