1use std::num::NonZeroU16;
2
3use raw_cpuid::{CpuId as RawCpuId, TopologyType};
4
5#[cfg(target_os = "windows")]
6use windows::Win32::{Foundation::*, System::SystemInformation::*};
7
8#[derive(Debug, Clone, Copy, PartialEq)]
9pub enum Vendor {
10 Intel,
11 Amd,
12 Unknown,
13}
14
15pub struct CpuInfo {
16 vendor: Vendor,
17 family: u32,
18 model: u32,
19 stepping: u32,
20 core_count: Option<NonZeroU16>,
21}
22
23impl CpuInfo {
24 pub fn new() -> Option<Self> {
25 let cpuid = RawCpuId::new();
26
27 let vendor = match cpuid.get_vendor_info() {
29 Some(info) => match info.as_str() {
30 "GenuineIntel" => Vendor::Intel,
31 "AuthenticAMD" => Vendor::Amd,
32 _ => Vendor::Unknown,
33 },
34 None => Vendor::Unknown,
35 };
36
37 let (family, model, stepping) = match cpuid.get_feature_info() {
39 Some(info) => (
40 info.family_id() as u32,
41 info.model_id() as u32,
42 info.stepping_id() as u32,
43 ),
44 None => (0, 0, 0),
45 };
46
47 let core_count = cpuid
48 .get_extended_topology_info()
49 .and_then(|mut info| info.find(|level| level.level_type() == TopologyType::Core))
50 .and_then(|level| NonZeroU16::new(level.processors()));
51
52 Some(CpuInfo {
53 vendor,
54 family,
55 model,
56 stepping,
57 core_count,
58 })
59 }
60
61 pub fn get_vendor(&self) -> Vendor {
62 self.vendor
63 }
64
65 pub fn get_family(&self) -> u32 {
66 self.family
67 }
68
69 pub fn get_model(&self) -> u32 {
70 self.model
71 }
72
73 pub fn get_stepping(&self) -> u32 {
74 self.stepping
75 }
76
77 pub fn get_core_count(&self) -> Option<NonZeroU16> {
78 self.core_count
79 }
80
81 pub fn has_temperature_sensor(&self) -> bool {
83 let cpuid = RawCpuId::new();
84 if cpuid.get_feature_info().is_some() {
85 if let Some(thermal_info) = cpuid.get_thermal_power_info() {
88 return thermal_info.has_dts(); }
90 }
91 false
92 }
93
94 pub fn get_name(&self) -> String {
96 let cpuid = RawCpuId::new();
97 if let Some(info) = cpuid.get_processor_brand_string() {
98 info.as_str().to_string()
99 } else {
100 "Unknown CPU".to_string()
101 }
102 }
103
104 pub fn get_core_cpu_mapping() -> anyhow::Result<Vec<Vec<usize>>> {
107 #[cfg(target_os = "windows")]
108 {
109 Self::get_core_cpu_mapping_windows()
110 }
111 #[cfg(target_os = "linux")]
112 {
113 let cpu_info = CpuInfo::new().ok_or_else(|| anyhow::anyhow!("Failed to get CPU info"))?;
114 cpu_info.get_core_cpu_mapping_linux()
115 }
116 #[cfg(not(any(target_os = "windows", target_os = "linux")))]
117 {
118 anyhow::bail!("not implementated");
120 }
121 }
122
123 #[cfg(target_os = "windows")]
125 fn get_core_cpu_mapping_windows() -> anyhow::Result<Vec<Vec<usize>>> {
126 use std::mem;
127
128 let mut mappings = Vec::new();
129
130 let mut buffer_length: u32 = 0;
132 let result =
133 unsafe { GetLogicalProcessorInformationEx(RelationAll, None, &mut buffer_length as _) };
134
135 match result {
136 Ok(()) => return Ok(mappings),
137 Err(e) => {
138 if unsafe { GetLastError() } != ERROR_INSUFFICIENT_BUFFER {
139 anyhow::bail!(
140 "GetLogicalProcessorInformationEx failed (0x{:X}): {}",
141 e.code().0,
142 e.message()
143 )
144 }
145 }
146 }
147
148 let mut buffer: Vec<u8> = vec![0; buffer_length as usize];
150
151 let result = unsafe {
153 GetLogicalProcessorInformationEx(
154 RelationAll,
155 Some(buffer.as_mut_ptr() as _),
156 &mut buffer_length as _,
157 )
158 };
159
160 if let Err(e) = result {
161 anyhow::bail!(
162 "GetLogicalProcessorInformationEx failed (0x{:X}): {}",
163 e.code().0,
164 e.message()
165 )
166 }
167
168 let mut offset: usize = 0;
170 while offset < buffer.len() {
171 use anyhow::Context;
172 let size = mem::size_of::<SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>();
173 let cur_buffer = buffer
174 .get(offset..offset + size)
175 .context("slice info error")?;
176 let info: SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX =
177 unsafe { std::ptr::read_unaligned(cur_buffer.as_ptr() as _) };
178 if info.Relationship == RelationProcessorCore {
179 let mut logical_processors = Vec::new();
180
181 let processor_info = unsafe { &info.Anonymous.Processor };
183
184 let group_mask_offset = offset
185 + mem::offset_of!(
186 SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX,
187 Anonymous.Processor.GroupMask
188 );
189 let group_mask_size = processor_info.GroupCount as usize
190 * mem::size_of_val(&processor_info.GroupMask[0]);
191 let mask_buffer = buffer
192 .get(group_mask_offset..group_mask_offset + group_mask_size)
193 .context("slice group mask error")?;
194
195 for i in 0..processor_info.GroupCount as usize {
197 let group_mask = unsafe {
198 mask_buffer
199 .as_ptr()
200 .cast::<GROUP_AFFINITY>()
201 .add(i)
202 .read_unaligned()
203 };
204 let mut mask = group_mask.Mask;
205
206 while mask != 0 {
208 let bit = mask.trailing_zeros();
210 logical_processors.push(i * 64 + bit as usize);
211
212 mask &= mask - 1;
214 }
215 }
216
217 mappings.push(logical_processors);
218 }
219
220 offset += info.Size as usize;
222 }
223
224 Ok(mappings)
225 }
226
227 #[cfg(target_os = "linux")]
229 fn get_core_cpu_mapping_linux(&self) -> anyhow::Result<Vec<Vec<usize>>> {
230 use std::collections::BTreeMap;
231 use std::fs;
232
233 use anyhow::Context as _;
234
235 const CORE_ID_PATTERN: &str = "/sys/devices/system/cpu/cpu[0-9]*/topology/core_id";
236
237 let mut groups: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
238 for entry in glob::glob(CORE_ID_PATTERN).context("Failed to read glob pattern")? {
239 if let Ok(path) = entry {
240 let cpu_id = path
242 .ancestors()
243 .nth(2)
244 .and_then(|p| p.file_name())
245 .and_then(|s| s.to_str())
246 .and_then(|s| s.trim_start_matches("cpu").parse::<usize>().ok());
247
248 let core_id = fs::read_to_string(&path)
250 .ok()
251 .and_then(|s| s.trim().parse::<usize>().ok());
252
253 if let (Some(cpu), Some(core)) = (cpu_id, core_id) {
254 groups.entry(core).or_default().push(cpu);
255 }
256 }
257 }
258
259 Ok(groups.into_values().collect())
261 }
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267
268 #[test]
269 fn test_get_core_cpu_mapping() {
270 let res = CpuInfo::get_core_cpu_mapping().unwrap();
285 println!("{:?}", res);
286 }
287}