Skip to main content

max / everycycle

6.6 KB · 193 lines History Blame Raw
1 //! Vulkan device enumeration with a discrete-compute filter.
2 //!
3 //! The goal is to produce a `BoxReportPayload` skeleton for one host:
4 //! the list of bus addresses present, per-device driver versions, and
5 //! the host identity. Real measurements come later. The filter rejects
6 //! anything that isn't a discrete GPU with a non-trivial device-local
7 //! heap, which keeps BMC display chips (ASPEED, Matrox) and integrated
8 //! GPUs out of the list without ever consulting a vendor table.
9
10 use std::ffi::CStr;
11
12 use ash::{Entry, Instance, vk};
13 use everycycle_hal::{BoxReportPayload, DriverBinding, HostIdentity};
14
15 use crate::host::host_identity;
16
17 /// One discrete compute device the enumeration found.
18 #[derive(Clone, Debug)]
19 pub struct EnumeratedDevice {
20 pub bus_address: String,
21 pub pci_ids: String,
22 pub driver_version: String,
23 pub device_name: String,
24 pub physical_bytes: u64,
25 }
26
27 /// Result of enumerating one host.
28 #[derive(Clone, Debug)]
29 pub struct EnumeratedHost {
30 pub host: HostIdentity,
31 pub devices: Vec<EnumeratedDevice>,
32 }
33
34 impl EnumeratedHost {
35 /// Build a `BoxReportPayload` skeleton from this enumeration.
36 ///
37 /// `probed_at` should be an RFC 3339 timestamp; the caller supplies
38 /// it because there are several sensible time sources and no
39 /// reason for this layer to choose. `ambient_c` is similarly the
40 /// caller's responsibility once a thermal probe lands.
41 #[must_use]
42 pub fn into_box_payload(
43 self,
44 probed_at: impl Into<String>,
45 ambient_c: f32,
46 probe_suite_version: impl Into<String>,
47 ) -> BoxReportPayload {
48 let present_devices = self.devices.iter().map(|d| d.bus_address.clone()).collect();
49 let driver_versions = self
50 .devices
51 .iter()
52 .map(|d| DriverBinding {
53 bus_address: d.bus_address.clone(),
54 driver_version: d.driver_version.clone(),
55 })
56 .collect();
57 BoxReportPayload {
58 report_version: "1".into(),
59 probed_at: probed_at.into(),
60 host: self.host,
61 present_devices,
62 driver_versions,
63 ambient_c,
64 probe_suite_version: probe_suite_version.into(),
65 }
66 }
67 }
68
69 #[derive(Debug)]
70 pub enum EnumError {
71 VulkanLoad(ash::LoadingError),
72 Vulkan(vk::Result),
73 }
74
75 impl core::fmt::Display for EnumError {
76 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
77 match self {
78 Self::VulkanLoad(e) => write!(f, "could not load Vulkan loader: {e:?}"),
79 Self::Vulkan(r) => write!(f, "Vulkan call failed: {r:?}"),
80 }
81 }
82 }
83
84 impl std::error::Error for EnumError {}
85
86 const MIN_DEVICE_LOCAL_BYTES: u64 = 1_000_000_000;
87
88 /// Enumerate all discrete compute devices on the current host.
89 pub fn enumerate() -> Result<EnumeratedHost, EnumError> {
90 let entry = unsafe { Entry::load() }.map_err(EnumError::VulkanLoad)?;
91 let app_info = vk::ApplicationInfo::default().api_version(vk::API_VERSION_1_3);
92 let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
93 let instance =
94 unsafe { entry.create_instance(&create_info, None) }.map_err(EnumError::Vulkan)?;
95
96 let result = (|| -> Result<EnumeratedHost, EnumError> {
97 let devices =
98 unsafe { instance.enumerate_physical_devices() }.map_err(EnumError::Vulkan)?;
99 let mut out = Vec::new();
100 for dev in devices {
101 if let Some(d) = inspect_device(&instance, dev) {
102 out.push(d);
103 }
104 }
105 Ok(EnumeratedHost {
106 host: host_identity(),
107 devices: out,
108 })
109 })();
110
111 unsafe { instance.destroy_instance(None) };
112 result
113 }
114
115 fn inspect_device(instance: &Instance, dev: vk::PhysicalDevice) -> Option<EnumeratedDevice> {
116 let props = unsafe { instance.get_physical_device_properties(dev) };
117 if props.device_type != vk::PhysicalDeviceType::DISCRETE_GPU {
118 return None;
119 }
120
121 let mem = unsafe { instance.get_physical_device_memory_properties(dev) };
122 let physical_bytes = mem
123 .memory_heaps
124 .iter()
125 .take(mem.memory_heap_count as usize)
126 .filter(|h| h.flags.contains(vk::MemoryHeapFlags::DEVICE_LOCAL))
127 .map(|h| h.size)
128 .max()
129 .unwrap_or(0);
130 if physical_bytes < MIN_DEVICE_LOCAL_BYTES {
131 return None;
132 }
133
134 let device_name = unsafe { CStr::from_ptr(props.device_name.as_ptr()) }
135 .to_string_lossy()
136 .into_owned();
137 // PCI vendor and device IDs are 16-bit by spec; Vulkan reports them
138 // widened to u32. Format as 4 hex digits without truncating casts.
139 let pci_ids = format!("{:04x}:{:04x}", props.vendor_id, props.device_id);
140
141 let bus_address = pci_bus_address(instance, dev).unwrap_or_else(|| "0000:00:00.0".to_string());
142 let driver_version = driver_string(instance, dev);
143
144 Some(EnumeratedDevice {
145 bus_address,
146 pci_ids,
147 driver_version,
148 device_name,
149 physical_bytes,
150 })
151 }
152
153 fn pci_bus_address(instance: &Instance, dev: vk::PhysicalDevice) -> Option<String> {
154 if !device_supports_extension(instance, dev, ash::ext::pci_bus_info::NAME) {
155 return None;
156 }
157 let mut pci = vk::PhysicalDevicePCIBusInfoPropertiesEXT::default();
158 let mut props2 = vk::PhysicalDeviceProperties2::default().push_next(&mut pci);
159 unsafe { instance.get_physical_device_properties2(dev, &mut props2) };
160 Some(format!(
161 "{:04x}:{:02x}:{:02x}.{:x}",
162 pci.pci_domain, pci.pci_bus, pci.pci_device, pci.pci_function
163 ))
164 }
165
166 fn driver_string(instance: &Instance, dev: vk::PhysicalDevice) -> String {
167 let mut driver = vk::PhysicalDeviceDriverProperties::default();
168 let mut props2 = vk::PhysicalDeviceProperties2::default().push_next(&mut driver);
169 unsafe { instance.get_physical_device_properties2(dev, &mut props2) };
170 let name = unsafe { CStr::from_ptr(driver.driver_name.as_ptr()) }
171 .to_string_lossy()
172 .into_owned();
173 let info = unsafe { CStr::from_ptr(driver.driver_info.as_ptr()) }
174 .to_string_lossy()
175 .into_owned();
176 match (name.is_empty(), info.is_empty()) {
177 (true, true) => "unknown".to_string(),
178 (false, true) => name,
179 (true, false) => info,
180 (false, false) => format!("{name} ({info})"),
181 }
182 }
183
184 fn device_supports_extension(instance: &Instance, dev: vk::PhysicalDevice, want: &CStr) -> bool {
185 let Ok(exts) = (unsafe { instance.enumerate_device_extension_properties(dev) }) else {
186 return false;
187 };
188 exts.iter().any(|ext| {
189 let name = unsafe { CStr::from_ptr(ext.extension_name.as_ptr()) };
190 name == want
191 })
192 }
193