dspk_core/
device.rs

1//! Which compute runtime runs the DSP kernels.
2//!
3//! [`ComputeTarget`] is the one name for a runtime across the workspace: the CLI flag, Python
4//! argument, app setting and the `DSPK_RUNTIME` environment variable all resolve to it.
5//! Choosing a target needs no GPU libraries; running work on it is [`crate::compute`] (the
6//! `compute` feature). Which targets can run depends on the enabled runtime features (`wgpu`,
7//! `cpu`, `cuda`, `hip`).
8
9use serde::{Deserialize, Serialize};
10use thiserror::Error;
11
12/// Environment variable that overrides [`ComputeTarget::from_env`]'s default (`wgpu`, `cpu`,
13/// `cuda`, `hip`).
14pub const RUNTIME_ENV: &str = "DSPK_RUNTIME";
15
16/// Environment variable that limits the runtimes tests run on ([`ComputeTarget::for_tests`]):
17/// comma-separated names, e.g. `cpu` on a CI machine that has CUDA compiled in but no GPU.
18pub const TEST_RUNTIMES_ENV: &str = "DSPK_TEST_RUNTIMES";
19
20/// A CubeCL runtime and its default device.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
22#[serde(rename_all = "lowercase")]
23pub enum ComputeTarget {
24    /// WebGPU (Vulkan / Metal / DX12).
25    Wgpu,
26    /// Multi-threaded CPU (MLIR JIT).
27    Cpu,
28    /// NVIDIA CUDA.
29    Cuda,
30    /// AMD ROCm HIP.
31    Hip,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Error)]
35pub enum ComputeError {
36    #[error("unknown compute runtime '{0}' (expected wgpu, cpu, cuda or hip)")]
37    Unknown(String),
38    #[error("compute runtime {0:?} is not compiled in (enable the `{1}` feature)")]
39    Unavailable(ComputeTarget, &'static str),
40    #[error("no compute runtime is compiled in (enable one of the wgpu, cpu, cuda, hip features)")]
41    NoneAvailable,
42}
43
44impl ComputeTarget {
45    pub const ALL: [ComputeTarget; 4] = [Self::Wgpu, Self::Cuda, Self::Hip, Self::Cpu];
46
47    pub fn name(self) -> &'static str {
48        match self {
49            Self::Wgpu => "wgpu",
50            Self::Cpu => "cpu",
51            Self::Cuda => "cuda",
52            Self::Hip => "hip",
53        }
54    }
55
56    pub fn parse(name: &str) -> Result<Self, ComputeError> {
57        match name.trim().to_ascii_lowercase().as_str() {
58            "wgpu" | "webgpu" | "vulkan" | "metal" | "gpu" => Ok(Self::Wgpu),
59            "cpu" => Ok(Self::Cpu),
60            "cuda" => Ok(Self::Cuda),
61            "hip" | "rocm" => Ok(Self::Hip),
62            other => Err(ComputeError::Unknown(other.to_string())),
63        }
64    }
65
66    /// Whether this runtime is compiled in.
67    pub fn is_available(self) -> bool {
68        match self {
69            Self::Wgpu => cfg!(feature = "wgpu"),
70            Self::Cpu => cfg!(feature = "cpu"),
71            Self::Cuda => cfg!(feature = "cuda"),
72            Self::Hip => cfg!(feature = "hip"),
73        }
74    }
75
76    /// Compiled-in runtimes in preference order (GPU runtimes first).
77    pub fn available() -> Vec<Self> {
78        Self::ALL.into_iter().filter(|t| t.is_available()).collect()
79    }
80
81    /// The runtimes tests run on: [`Self::available`], limited to the names in `DSPK_TEST_RUNTIMES`
82    /// when it is set. Compiled in is not the same as usable: a build with CUDA can run on a
83    /// machine without an NVIDIA driver.
84    ///
85    /// # Panics
86    /// If `DSPK_TEST_RUNTIMES` names an unknown runtime (a typo would silently skip tests).
87    pub fn for_tests() -> Vec<Self> {
88        match std::env::var(TEST_RUNTIMES_ENV) {
89            Ok(list) if !list.trim().is_empty() => {
90                let wanted: Vec<Self> = list
91                    .split(',')
92                    .map(|name| Self::parse(name).unwrap_or_else(|e| panic!("{TEST_RUNTIMES_ENV}: {e}")))
93                    .collect();
94                Self::available().into_iter().filter(|t| wanted.contains(t)).collect()
95            }
96            _ => Self::available(),
97        }
98    }
99
100    /// `DSPK_RUNTIME` when set, else the first compiled-in runtime.
101    pub fn from_env() -> Result<Self, ComputeError> {
102        match std::env::var(RUNTIME_ENV) {
103            Ok(name) if !name.trim().is_empty() => Self::parse(&name)?.checked(),
104            _ => Self::available().first().copied().ok_or(ComputeError::NoneAvailable),
105        }
106    }
107
108    /// `self` if it is compiled in, else [`ComputeError::Unavailable`].
109    pub fn checked(self) -> Result<Self, ComputeError> {
110        if self.is_available() { Ok(self) } else { Err(ComputeError::Unavailable(self, self.name())) }
111    }
112}
113
114impl std::fmt::Display for ComputeTarget {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        f.write_str(self.name())
117    }
118}
119
120impl std::str::FromStr for ComputeTarget {
121    type Err = ComputeError;
122    fn from_str(s: &str) -> Result<Self, Self::Err> {
123        Self::parse(s)
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn parses_names_and_reports_missing_runtimes() {
133        assert_eq!(ComputeTarget::parse("ROCm").unwrap(), ComputeTarget::Hip);
134        assert!(ComputeTarget::parse("tpu").is_err());
135        for target in ComputeTarget::ALL.into_iter().filter(|t| !t.is_available()) {
136            assert!(matches!(target.checked(), Err(ComputeError::Unavailable(..))));
137        }
138    }
139}
dspk_core/device.rs

dspk_core/
device.rs
