mesh_llm_native_runtime/
flavor.rs

1use serde::{Deserialize, Deserializer, Serialize, Serializer};
2use std::{fmt, str::FromStr};
3
4#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
5pub enum NativeRuntimeBackendKind {
6    Cpu,
7    Metal,
8    Cuda,
9    Rocm,
10    Vulkan,
11    Other(String),
12}
13
14pub type NativeRuntimeFlavor = NativeRuntimeBackendKind;
15
16#[derive(Clone, Debug, Eq, PartialEq)]
17pub struct NativeRuntimeFlavorParseError {
18    value: String,
19}
20
21#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
22pub struct NativeRuntimeBackend {
23    pub kind: NativeRuntimeBackendKind,
24    #[serde(default, skip_serializing_if = "Option::is_none")]
25    pub cuda: Option<CudaRuntimeRequirements>,
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub rocm: Option<RocmRuntimeRequirements>,
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub vulkan: Option<VulkanRuntimeRequirements>,
30}
31
32#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
33pub struct CudaRuntimeRequirements {
34    pub toolkit_major: u32,
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub min_driver: Option<String>,
37    #[serde(default, skip_serializing_if = "Vec::is_empty")]
38    pub gpu_arches: Vec<String>,
39}
40
41#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
42pub struct RocmRuntimeRequirements {
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub version: Option<String>,
45    #[serde(default, skip_serializing_if = "Vec::is_empty")]
46    pub gpu_arches: Vec<String>,
47}
48
49#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
50pub struct VulkanRuntimeRequirements {
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub min_api_version: Option<String>,
53}
54
55impl fmt::Display for NativeRuntimeFlavorParseError {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        write!(f, "invalid native runtime backend '{}'", self.value)
58    }
59}
60
61impl std::error::Error for NativeRuntimeFlavorParseError {}
62
63impl NativeRuntimeBackendKind {
64    pub fn as_str(&self) -> &str {
65        match self {
66            Self::Cpu => "cpu",
67            Self::Metal => "metal",
68            Self::Cuda => "cuda",
69            Self::Rocm => "rocm",
70            Self::Vulkan => "vulkan",
71            Self::Other(value) => value.as_str(),
72        }
73    }
74
75    pub fn default_rank(&self) -> i64 {
76        match self {
77            Self::Cuda => 650,
78            Self::Rocm => 600,
79            Self::Metal => 600,
80            Self::Vulkan => 350,
81            Self::Cpu => 100,
82            Self::Other(_) => 0,
83        }
84    }
85}
86
87impl NativeRuntimeBackend {
88    pub fn cpu() -> Self {
89        Self {
90            kind: NativeRuntimeBackendKind::Cpu,
91            cuda: None,
92            rocm: None,
93            vulkan: None,
94        }
95    }
96
97    pub fn metal() -> Self {
98        Self {
99            kind: NativeRuntimeBackendKind::Metal,
100            cuda: None,
101            rocm: None,
102            vulkan: None,
103        }
104    }
105
106    pub fn cuda(toolkit_major: u32, gpu_arches: Vec<String>) -> Self {
107        Self {
108            kind: NativeRuntimeBackendKind::Cuda,
109            cuda: Some(CudaRuntimeRequirements {
110                toolkit_major,
111                min_driver: None,
112                gpu_arches,
113            }),
114            rocm: None,
115            vulkan: None,
116        }
117    }
118
119    pub fn rocm(gpu_arches: Vec<String>) -> Self {
120        Self {
121            kind: NativeRuntimeBackendKind::Rocm,
122            cuda: None,
123            rocm: Some(RocmRuntimeRequirements {
124                version: None,
125                gpu_arches,
126            }),
127            vulkan: None,
128        }
129    }
130
131    pub fn vulkan() -> Self {
132        Self {
133            kind: NativeRuntimeBackendKind::Vulkan,
134            cuda: None,
135            rocm: None,
136            vulkan: Some(VulkanRuntimeRequirements {
137                min_api_version: None,
138            }),
139        }
140    }
141}
142
143impl fmt::Display for NativeRuntimeBackendKind {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        f.write_str(self.as_str())
146    }
147}
148
149impl Serialize for NativeRuntimeBackendKind {
150    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
151    where
152        S: Serializer,
153    {
154        serializer.serialize_str(self.as_str())
155    }
156}
157
158impl<'de> Deserialize<'de> for NativeRuntimeBackendKind {
159    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
160    where
161        D: Deserializer<'de>,
162    {
163        let value = String::deserialize(deserializer)?;
164        Ok(Self::from(value.as_str()))
165    }
166}
167
168impl FromStr for NativeRuntimeBackendKind {
169    type Err = NativeRuntimeFlavorParseError;
170
171    fn from_str(value: &str) -> Result<Self, Self::Err> {
172        let normalized = value.trim().to_ascii_lowercase();
173        if normalized.is_empty() {
174            return Err(NativeRuntimeFlavorParseError {
175                value: value.to_string(),
176            });
177        }
178        Ok(match normalized.as_str() {
179            "cpu" => Self::Cpu,
180            "metal" => Self::Metal,
181            "cuda" | "cuda-blackwell" | "blackwell" => Self::Cuda,
182            "rocm" | "hip" => Self::Rocm,
183            "vulkan" => Self::Vulkan,
184            _ => Self::Other(normalized),
185        })
186    }
187}
188
189impl From<&str> for NativeRuntimeBackendKind {
190    fn from(value: &str) -> Self {
191        value
192            .parse()
193            .unwrap_or_else(|_| Self::Other(value.trim().to_ascii_lowercase()))
194    }
195}
mesh_llm_native_runtime/flavor.rs

mesh_llm_native_runtime/
flavor.rs
