1use crate::error::{CoreError, CoreResult};
7
8#[derive(Debug, Clone)]
10pub struct GpuInfo {
11 pub name: String,
13 pub vendor: GpuVendor,
15 pub memory_total: usize,
17 pub memory_available: usize,
19 pub memorybandwidth_gbps: f64,
21 pub compute_units: usize,
23 pub base_clock_mhz: usize,
25 pub memory_clock_mhz: usize,
27 pub compute_capability: ComputeCapability,
29 pub features: GpuFeatures,
31 pub performance: GpuPerformance,
33}
34
35impl GpuInfo {
36 pub fn detect() -> CoreResult<Self> {
38 #[cfg(feature = "gpu")]
39 {
40 if let Ok(gpu) = Self::detect_cuda() {
42 return Ok(gpu);
43 }
44
45 if let Ok(gpu) = Self::detect_opencl() {
46 return Ok(gpu);
47 }
48
49 if let Ok(gpu) = Self::detect_vulkan() {
50 return Ok(gpu);
51 }
52 }
53
54 #[cfg(target_os = "linux")]
56 if let Ok(gpu) = Self::detect_linux() {
57 return Ok(gpu);
58 }
59
60 #[cfg(target_os = "windows")]
61 if let Ok(gpu) = Self::detect_windows() {
62 return Ok(gpu);
63 }
64
65 #[cfg(target_os = "macos")]
66 if let Ok(gpu) = Self::detect_macos() {
67 return Ok(gpu);
68 }
69
70 Err(CoreError::ComputationError(
71 crate::error::ErrorContext::new("No GPU detected"),
72 ))
73 }
74
75 #[cfg(feature = "gpu")]
82 fn detect_cuda() -> CoreResult<Self> {
83 Err(CoreError::ComputationError(
84 crate::error::ErrorContext::new("CUDA detection not implemented"),
85 ))
86 }
87
88 #[cfg(feature = "gpu")]
90 fn detect_opencl() -> CoreResult<Self> {
91 Err(CoreError::ComputationError(
93 crate::error::ErrorContext::new("OpenCL detection not implemented"),
94 ))
95 }
96
97 #[cfg(feature = "gpu")]
99 fn detect_vulkan() -> CoreResult<Self> {
100 Err(CoreError::ComputationError(
102 crate::error::ErrorContext::new("Vulkan detection not implemented"),
103 ))
104 }
105
106 #[cfg(target_os = "linux")]
108 fn detect_linux() -> CoreResult<Self> {
109 use std::fs;
110
111 if let Ok(entries) = fs::read_dir("/sys/class/drm") {
113 for entry in entries.flatten() {
114 let path = entry.path();
115 if let Some(name) = path.file_name() {
116 if name.to_string_lossy().starts_with("card") {
117 let device_path = path.join("device");
119 if let Ok(vendor) = fs::read_to_string(device_path.join("vendor")) {
120 if let Ok(device) = fs::read_to_string(device_path.join("device")) {
121 let vendor_id = vendor.trim();
122 let device_id = device.trim();
123
124 return Ok(Self::create_from_pci_ids(vendor_id, device_id));
125 }
126 }
127 }
128 }
129 }
130 }
131
132 Err(CoreError::ComputationError(
133 crate::error::ErrorContext::new("No GPU detected on Linux"),
134 ))
135 }
136
137 #[cfg(target_os = "windows")]
139 fn detect_windows() -> CoreResult<Self> {
140 Err(CoreError::ComputationError(
142 crate::error::ErrorContext::new("Windows GPU detection not implemented"),
143 ))
144 }
145
146 #[cfg(target_os = "macos")]
148 fn detect_macos() -> CoreResult<Self> {
149 #[cfg(target_arch = "aarch64")]
151 {
152 Ok(Self {
153 name: "Apple GPU".to_string(),
154 vendor: GpuVendor::Apple,
155 memory_total: 8 * 1024 * 1024 * 1024, memory_available: 6 * 1024 * 1024 * 1024,
157 memorybandwidth_gbps: 200.0,
158 compute_units: 8,
159 base_clock_mhz: 1000,
160 memory_clock_mhz: 2000,
161 compute_capability: ComputeCapability::Metal,
162 features: GpuFeatures {
163 unified_memory: true,
164 double_precision: true,
165 half_precision: true,
166 tensor_cores: false,
167 ray_tracing: false,
168 },
169 performance: GpuPerformance {
170 fp32_gflops: 2600.0,
171 fp16_gflops: 5200.0,
172 memorybandwidth_gbps: 200.0,
173 efficiency_score: 0.9,
174 },
175 })
176 }
177 #[cfg(not(target_arch = "aarch64"))]
178 {
179 Err(CoreError::ComputationError(
180 crate::error::ErrorContext::new("macOS GPU detection not implemented"),
181 ))
182 }
183 }
184
185 #[allow(dead_code)]
187 fn from_pci_ids(vendor_id: u16, _device_id: &str) -> Self {
188 let vendor = match vendor_id {
189 0x10de => GpuVendor::Nvidia,
190 0x1002 => GpuVendor::Amd,
191 0x8086 => GpuVendor::Intel,
192 _ => GpuVendor::Unknown,
193 };
194
195 let (name, memory_gb, compute_units) = match vendor_id {
198 0x10de => ("NVIDIA GPU".to_string(), 8, 2048),
199 0x1002 => ("AMD GPU".to_string(), 8, 64),
200 0x8086 => ("Intel GPU".to_string(), 4, 96),
201 _ => ("Unknown GPU".to_string(), 4, 32),
202 };
203
204 Self {
205 name,
206 vendor,
207 memory_total: memory_gb * 1024 * 1024 * 1024,
208 memory_available: (memory_gb * 1024 * 1024 * 1024 * 3) / 4, memorybandwidth_gbps: 500.0,
210 compute_units,
211 base_clock_mhz: 1500,
212 memory_clock_mhz: 7000,
213 compute_capability: ComputeCapability::Unknown,
214 features: GpuFeatures::default(),
215 performance: GpuPerformance::default(),
216 }
217 }
218
219 pub fn performance_score(&self) -> f64 {
221 let memory_score = (self.memory_total as f64 / (24.0 * 1024.0 * 1024.0 * 1024.0)).min(1.0); let compute_score = (self.compute_units as f64 / 4096.0).min(1.0); let bandwidth_score = (self.memorybandwidth_gbps / 1000.0).min(1.0); let efficiency_score = self.performance.efficiency_score;
225
226 (memory_score + compute_score + bandwidth_score + efficiency_score) / 4.0
227 }
228
229 pub fn optimal_workgroup_size(&self) -> usize {
231 match self.vendor {
232 GpuVendor::Nvidia => 256, GpuVendor::Amd => 64, GpuVendor::Intel => 128, GpuVendor::Apple => 32, GpuVendor::Unknown => 64,
237 }
238 }
239
240 pub fn is_compute_capable(&self) -> bool {
242 self.memory_total >= 2 * 1024 * 1024 * 1024 && self.compute_units >= 32 }
245
246 pub fn is_ml_capable(&self) -> bool {
248 self.is_compute_capable() && (self.features.tensor_cores || self.features.half_precision)
249 }
250
251 pub fn create_from_pci_ids(vendor_id: &str, device_id: &str) -> Self {
253 let vendor_id = vendor_id.strip_prefix("0x").unwrap_or(vendor_id);
255
256 let vendor = match vendor_id {
257 "10de" => GpuVendor::Nvidia,
258 "1002" => GpuVendor::Amd,
259 "8086" => GpuVendor::Intel,
260 _ => GpuVendor::Unknown,
261 };
262
263 let name = match vendor {
265 GpuVendor::Nvidia => format!("NVIDIA GPU {}", device_id),
266 GpuVendor::Amd => format!("AMD GPU {}", device_id),
267 GpuVendor::Intel => format!("Intel GPU {}", device_id),
268 GpuVendor::Apple => format!("Apple GPU {}", device_id),
269 GpuVendor::Unknown => format!("Unknown GPU {}", device_id),
270 };
271
272 Self {
275 name,
276 vendor,
277 memory_total: (8u64 * 1024 * 1024 * 1024) as usize, memory_available: (8u64 * 1024 * 1024 * 1024) as usize,
279 memorybandwidth_gbps: 400.0,
280 compute_capability: ComputeCapability::Cuda(7, 0), compute_units: 128,
282 base_clock_mhz: 1500,
283 memory_clock_mhz: 1750, features: GpuFeatures::default(),
285 performance: GpuPerformance::default(),
286 }
287 }
288}
289
290#[derive(Debug, Clone, Copy, PartialEq, Eq)]
292pub enum GpuVendor {
293 Nvidia,
295 Amd,
297 Intel,
299 Apple,
301 Unknown,
303}
304
305#[derive(Debug, Clone, Copy, PartialEq, Eq)]
307pub enum ComputeCapability {
308 Cuda(u32, u32), OpenCL(u32, u32), Vulkan(u32, u32), Metal,
316 DirectCompute,
318 Unknown,
320}
321
322#[derive(Debug, Clone)]
324pub struct GpuFeatures {
325 pub unified_memory: bool,
327 pub double_precision: bool,
329 pub half_precision: bool,
331 pub tensor_cores: bool,
333 pub ray_tracing: bool,
335}
336
337impl Default for GpuFeatures {
338 fn default() -> Self {
339 Self {
340 unified_memory: false,
341 double_precision: true,
342 half_precision: false,
343 tensor_cores: false,
344 ray_tracing: false,
345 }
346 }
347}
348
349#[derive(Debug, Clone)]
351pub struct GpuPerformance {
352 pub fp32_gflops: f64,
354 pub fp16_gflops: f64,
356 pub memorybandwidth_gbps: f64,
358 pub efficiency_score: f64,
360}
361
362impl Default for GpuPerformance {
363 fn default() -> Self {
364 Self {
365 fp32_gflops: 1000.0,
366 fp16_gflops: 2000.0,
367 memorybandwidth_gbps: 500.0,
368 efficiency_score: 0.7,
369 }
370 }
371}
372
373#[derive(Debug, Clone)]
375pub struct MultiGpuInfo {
376 pub gpus: Vec<GpuInfo>,
378 pub total_memory: usize,
380 pub p2p_capable: bool,
382 pub multi_gpuconfig: MultiGpuConfig,
384}
385
386impl MultiGpuInfo {
387 pub fn detect() -> CoreResult<Self> {
389 let mut gpus = Vec::new();
390
391 if let Ok(gpu) = GpuInfo::detect() {
394 gpus.push(gpu);
395 }
396
397 let total_memory = gpus.iter().map(|gpu| gpu.memory_total).sum();
398
399 Ok(Self {
400 gpus,
401 total_memory,
402 p2p_capable: false,
403 multi_gpuconfig: MultiGpuConfig::Single,
404 })
405 }
406
407 pub fn best_compute_gpu(&self) -> Option<&GpuInfo> {
409 self.gpus
410 .iter()
411 .filter(|gpu| gpu.is_compute_capable())
412 .max_by(|a, b| {
413 a.performance_score()
414 .partial_cmp(&b.performance_score())
415 .expect("Operation failed")
416 })
417 }
418
419 pub fn total_compute_units(&self) -> usize {
421 self.gpus.iter().map(|gpu| gpu.compute_units).sum()
422 }
423}
424
425#[derive(Debug, Clone, Copy, PartialEq, Eq)]
427pub enum MultiGpuConfig {
428 Single,
430 Sli,
432 CrossFire,
434 NvLink,
436 Independent,
438}
439
440impl Default for GpuInfo {
441 fn default() -> Self {
442 Self {
443 name: "Default GPU".to_string(),
444 vendor: GpuVendor::Unknown,
445 memory_total: (4u64 * 1024 * 1024 * 1024) as usize, memory_available: (3u64 * 1024 * 1024 * 1024) as usize, memorybandwidth_gbps: 200.0,
448 compute_units: 512,
449 base_clock_mhz: 1000,
450 memory_clock_mhz: 4000,
451 compute_capability: ComputeCapability::Unknown,
452 features: GpuFeatures::default(),
453 performance: GpuPerformance::default(),
454 }
455 }
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461
462 #[test]
463 fn test_gpu_vendor() {
464 assert_eq!(GpuVendor::Nvidia, GpuVendor::Nvidia);
465 assert_ne!(GpuVendor::Nvidia, GpuVendor::Amd);
466 }
467
468 #[test]
469 fn test_compute_capability() {
470 let cuda_cap = ComputeCapability::Cuda(7, 5);
471 assert_eq!(cuda_cap, ComputeCapability::Cuda(7, 5));
472 assert_ne!(cuda_cap, ComputeCapability::Metal);
473 }
474
475 #[test]
476 fn test_gpu_features() {
477 let features = GpuFeatures {
478 unified_memory: true,
479 tensor_cores: true,
480 ..Default::default()
481 };
482
483 assert!(features.unified_memory);
484 assert!(features.tensor_cores);
485 assert!(!features.ray_tracing);
486 }
487
488 #[test]
489 fn test_gpu_performance() {
490 let perf = GpuPerformance::default();
491 assert!(perf.fp32_gflops > 0.0);
492 assert!(perf.efficiency_score >= 0.0 && perf.efficiency_score <= 1.0);
493 }
494
495 #[test]
496 fn test_pci_id_parsing() {
497 let gpu = GpuInfo::create_from_pci_ids("0x10de", "0x1234");
498 assert_eq!(gpu.vendor, GpuVendor::Nvidia);
499 assert!(gpu.name.contains("NVIDIA"));
500 }
501
502 #[test]
503 fn test_multi_gpu_config() {
504 assert_eq!(MultiGpuConfig::Single, MultiGpuConfig::Single);
505 assert_ne!(MultiGpuConfig::Single, MultiGpuConfig::Sli);
506 }
507
508 #[test]
509 fn test_optimal_workgroup_size() {
510 let nvidia_gpu = GpuInfo {
511 vendor: GpuVendor::Nvidia,
512 ..Default::default()
513 };
514 assert_eq!(nvidia_gpu.optimal_workgroup_size(), 256);
515
516 let amd_gpu = GpuInfo {
517 vendor: GpuVendor::Amd,
518 ..Default::default()
519 };
520 assert_eq!(amd_gpu.optimal_workgroup_size(), 64);
521 }
522}