1
2#[derive(Clone, Debug)]
3pub struct DeviceBuffer<T> {
4 host_shadow: Vec<T>,
5}
6
7impl<T> DeviceBuffer<T> {
8 pub const fn from_host_shadow(host_shadow: Vec<T>) -> Self {
9 Self { host_shadow }
10 }
11
12 pub const fn len(&self) -> usize {
13 self.host_shadow.len()
14 }
15
16 pub const fn is_empty(&self) -> bool {
17 self.host_shadow.len() == 0
18 }
19
20 pub fn host_shadow(&self) -> &[T] {
21 &self.host_shadow
22 }
23}
24
25#[derive(Clone, Debug)]
26pub struct DeviceVector {
27 pub len: usize,
28 pub data: DeviceBuffer<f64>,
29}
30
31#[derive(Clone, Debug)]
32pub struct DeviceMatrix {
33 pub rows: usize,
34 pub cols: usize,
35 pub data: DeviceBuffer<f64>,
36 pub column_major: bool,
37}
38
39impl DeviceMatrix {
40
41 pub const fn bytes(&self) -> usize {
42 self.rows
43 .saturating_mul(self.cols)
44 .saturating_mul(std::mem::size_of::<f64>())
45 }
46}
47
48#[derive(Clone, Debug)]
49pub struct DeviceCsrMatrix {
50 pub rows: usize,
51 pub cols: usize,
52 pub rowptr: DeviceBuffer<i32>,
53 pub colidx: DeviceBuffer<i32>,
54 pub values: DeviceBuffer<f64>,
55}
56
57impl DeviceCsrMatrix {
58 pub fn new(
70 rows: usize,
71 cols: usize,
72 rowptr: DeviceBuffer<i32>,
73 colidx: DeviceBuffer<i32>,
74 values: DeviceBuffer<f64>,
75 ) -> Self {
76 let expected = rows + 1;
77 let mut ptr = rowptr.host_shadow().to_vec();
78 if ptr.len() != expected {
79 let fill = ptr.last().copied().unwrap_or(0);
80 ptr.resize(expected, fill);
81 }
82 Self {
83 rows,
84 cols,
85 rowptr: DeviceBuffer::from_host_shadow(ptr),
86 colidx,
87 values,
88 }
89 }
90
91 pub const fn nnz(&self) -> usize {
92 self.values.len()
93 }
94}