Skip to main content

oxidelake_core/
hardware.rs

1//! Hardware backend and device identifiers.
2
3use std::fmt;
4use std::str::FromStr;
5
6use serde::{Deserialize, Serialize};
7
8use crate::EngineError;
9
10/// The execution backends OxideLake knows about.
11///
12/// `CpuSimd` is always available and is the correctness reference; the GPU
13/// backends are opt-in cargo features and are selected at runtime by the
14/// hardware detector.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
16pub enum BackendKind {
17    /// Vectorized CPU execution through Arrow/DataFusion kernels and rayon.
18    CpuSimd,
19    /// NVIDIA CUDA (feature `cuda`).
20    Cuda,
21    /// Apple Metal on unified memory (feature `metal`, macOS only).
22    Metal,
23}
24
25impl BackendKind {
26    /// Every backend kind, in detection priority order (GPU first).
27    pub const ALL: [BackendKind; 3] = [BackendKind::Cuda, BackendKind::Metal, BackendKind::CpuSimd];
28
29    /// The canonical lowercase name used in environment variables, `EXPLAIN`
30    /// tags and logs.
31    pub const fn as_str(self) -> &'static str {
32        match self {
33            BackendKind::CpuSimd => "cpu",
34            BackendKind::Cuda => "cuda",
35            BackendKind::Metal => "metal",
36        }
37    }
38
39    /// `true` for backends that execute on a device rather than the host CPU.
40    pub const fn is_gpu(self) -> bool {
41        !matches!(self, BackendKind::CpuSimd)
42    }
43}
44
45impl fmt::Display for BackendKind {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        f.write_str(self.as_str())
48    }
49}
50
51impl FromStr for BackendKind {
52    type Err = EngineError;
53
54    fn from_str(s: &str) -> Result<Self, Self::Err> {
55        match s.trim().to_ascii_lowercase().as_str() {
56            "cpu" | "cpusimd" | "cpu-simd" | "cpu_simd" => Ok(BackendKind::CpuSimd),
57            "cuda" => Ok(BackendKind::Cuda),
58            "metal" => Ok(BackendKind::Metal),
59            other => Err(EngineError::plan(format!(
60                "unknown backend '{other}'; expected one of cpu, cuda, metal"
61            ))),
62        }
63    }
64}
65
66/// Identifies one device of a backend, e.g. `cuda:0`.
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
68pub struct DeviceId {
69    /// The backend the device belongs to.
70    pub backend: BackendKind,
71    /// The device ordinal within that backend (always `0` for the CPU).
72    pub ordinal: u32,
73}
74
75impl DeviceId {
76    /// The host CPU.
77    pub const CPU: DeviceId = DeviceId {
78        backend: BackendKind::CpuSimd,
79        ordinal: 0,
80    };
81
82    /// Builds a device identifier.
83    pub const fn new(backend: BackendKind, ordinal: u32) -> Self {
84        Self { backend, ordinal }
85    }
86}
87
88impl fmt::Display for DeviceId {
89    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90        write!(f, "{}:{}", self.backend, self.ordinal)
91    }
92}
93
94#[cfg(test)]
95#[allow(clippy::unwrap_used, clippy::expect_used)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn parses_all_names_case_insensitively() {
101        assert_eq!("cpu".parse::<BackendKind>().unwrap(), BackendKind::CpuSimd);
102        assert_eq!(" CUDA ".parse::<BackendKind>().unwrap(), BackendKind::Cuda);
103        assert_eq!("Metal".parse::<BackendKind>().unwrap(), BackendKind::Metal);
104        assert!("tpu".parse::<BackendKind>().is_err());
105    }
106
107    #[test]
108    fn display_round_trips_through_parse() {
109        for kind in BackendKind::ALL {
110            assert_eq!(kind.as_str().parse::<BackendKind>().unwrap(), kind);
111        }
112    }
113
114    #[test]
115    fn device_id_display() {
116        assert_eq!(DeviceId::new(BackendKind::Cuda, 1).to_string(), "cuda:1");
117        assert_eq!(DeviceId::CPU.to_string(), "cpu:0");
118        assert!(!BackendKind::CpuSimd.is_gpu());
119        assert!(BackendKind::Metal.is_gpu());
120    }
121}