singe_cutensor/mg/
workspace.rs

1use crate::{
2    error::{Error, Result},
3    mg::context::ContextRef,
4    utility::to_i64,
5};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct Workspace {
9    device_sizes: Vec<i64>,
10    host_size: i64,
11}
12
13impl Workspace {
14    pub fn create(device_sizes: Vec<u64>, host_size: u64) -> Result<Self> {
15        Ok(Self {
16            device_sizes: device_sizes
17                .into_iter()
18                .map(|value| to_i64(value, "device workspace size"))
19                .collect::<Result<_>>()?,
20            host_size: to_i64(host_size, "host workspace size")?,
21        })
22    }
23
24    pub(crate) fn from_raw(device_sizes: Vec<i64>, host_size: i64) -> Result<Self> {
25        if device_sizes.iter().any(|&size| size < 0) || host_size < 0 {
26            return Err(Error::OutOfRange {
27                name: "workspace size".into(),
28            });
29        }
30
31        Ok(Self {
32            device_sizes,
33            host_size,
34        })
35    }
36
37    pub fn device_sizes(&self) -> &[i64] {
38        &self.device_sizes
39    }
40
41    pub fn host_size(&self) -> i64 {
42        self.host_size
43    }
44
45    pub fn device_sizes_bytes(&self) -> Result<Vec<usize>> {
46        self.device_sizes
47            .iter()
48            .copied()
49            .map(|size| {
50                usize::try_from(size).map_err(|_| Error::OutOfRange {
51                    name: "device workspace size".into(),
52                })
53            })
54            .collect()
55    }
56
57    pub fn host_size_bytes(&self) -> Result<usize> {
58        usize::try_from(self.host_size).map_err(|_| Error::OutOfRange {
59            name: "host workspace size".into(),
60        })
61    }
62
63    pub(crate) fn validate_for_context(&self, context: &ContextRef) -> Result<()> {
64        if self.device_sizes.len() != context.device_count() {
65            return Err(Error::LengthMismatch {
66                name: "device_workspace_size".into(),
67                expected: context.device_count(),
68                actual: self.device_sizes.len(),
69            });
70        }
71        Ok(())
72    }
73
74    pub(crate) fn device_sizes_ptr(&self) -> *const i64 {
75        self.device_sizes.as_ptr()
76    }
77}
singe_cutensor/mg/workspace.rs

singe_cutensor/mg/
workspace.rs
