1use super::{Api, BackendEntry, BackendFactory, BackendKind, MemKind, Reason, Source, Status};
3use trueno_gpu::driver::{device_count, CudaContext};
4
5pub struct CudaFactory;
7
8const LIB: &str = "libcuda.so.1";
9
10impl BackendFactory for CudaFactory {
11 fn kind(&self) -> BackendKind {
12 BackendKind::Cuda
13 }
14
15 fn discover(&self) -> Vec<BackendEntry> {
16 if trueno_gpu::driver::sys::CudaDriver::load().is_none() {
20 return vec![BackendEntry::unavailable(
21 BackendKind::Cuda,
22 Api::CudaDriver,
23 Source::Dlopen(LIB.to_string()),
24 Reason::DriverNotFound { path: LIB.to_string() },
25 )];
26 }
27 let n = match device_count() {
28 Ok(n) => n,
29 Err(e) => {
30 return vec![BackendEntry::unavailable(
31 BackendKind::Cuda,
32 Api::CudaDriver,
33 Source::Dlopen(LIB.to_string()),
34 Reason::ProbeFailed { error: format!("cuDeviceGetCount: {e:?}") },
35 )]
36 }
37 };
38 if n == 0 {
39 return vec![BackendEntry::unavailable(
40 BackendKind::Cuda,
41 Api::CudaDriver,
42 Source::Dlopen(LIB.to_string()),
43 Reason::NoDevice,
44 )];
45 }
46 (0..n).map(probe_device).collect()
47 }
48}
49
50#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
51fn probe_device(i: usize) -> BackendEntry {
52 let ctx = match CudaContext::new(i as i32) {
53 Ok(c) => c,
54 Err(e) => {
55 let mut entry = BackendEntry::unavailable(
56 BackendKind::Cuda,
57 Api::CudaDriver,
58 Source::Dlopen(LIB.to_string()),
59 Reason::ProbeFailed { error: format!("cuCtxCreate({i}): {e:?}") },
60 );
61 entry.device_index = Some(i as u32);
62 return entry;
63 }
64 };
65 let name = ctx.device_name().unwrap_or_else(|_| format!("cuda device {i}"));
66 let (free, total) =
67 ctx.memory_info().map(|(f, t)| (Some(f as u64), Some(t as u64))).unwrap_or((None, None));
68 let cc = ctx.compute_capability().ok().map(|(maj, min)| format!("sm_{maj}{min}"));
69 BackendEntry {
70 kind: BackendKind::Cuda,
71 api: Api::CudaDriver,
72 device_index: Some(i as u32),
73 device_uid: Some(super::device_uid("nvidia", &name)),
74 device_name: name,
75 vendor: "NVIDIA".to_string(),
76 vendor_id: Some(0x10de),
77 device_type: "discrete-gpu".to_string(),
78 mem_total: total,
79 mem_free: free,
80 mem_kind: MemKind::Discrete,
81 compute_class: cc,
82 caps: vec!["async".to_string()],
83 source: Source::Dlopen(LIB.to_string()),
84 status: Status::Ready,
85 transport: None,
86 }
87}