pytorch_cpuinfo/
wrappers.rs

1use super::ffi::*;
2use once_cell::sync::Lazy;
3use std::marker::PhantomData;
4
5pub struct Cluster {
6    /// Index of the first logical processor in the cluster
7    pub id: u32,
8    /// Clock rate (non-Turbo) of the cores in the cluster, in Hz
9    pub frequency: u64,
10    /// Either value of CPUID leaf 1 EAX register on x86 OR value of Main ID Register
11    /// on ARM
12    pub cpuid: u32,
13}
14
15pub struct Package {
16    pub name: String,
17    pub clusters: Vec<Cluster>,
18}
19
20struct Wrapper {
21    pub packages: Vec<Package>,
22    pub phantom: PhantomData<Vec<Package>>,
23}
24
25impl Drop for Wrapper {
26    fn drop(&mut self) {
27        unsafe { cpuinfo_deinitialize() };
28    }
29}
30
31static GLOBAL_SOCKETS: Lazy<Wrapper> = Lazy::new(|| {
32    unsafe { cpuinfo_initialize() };
33
34    let mut sockets: Vec<Package> = vec![];
35
36    let packages_count = unsafe { cpuinfo_get_packages_count() };
37
38    let packages_ptr = unsafe { cpuinfo_get_packages() };
39
40    let packages = unsafe { std::slice::from_raw_parts(packages_ptr, packages_count as usize) };
41
42    for p in packages {
43        let name = unsafe { std::ffi::CStr::from_ptr(p.name.as_ptr()) };
44
45        let mut clusters = vec![];
46
47        for cid in p.cluster_start..(p.cluster_start + p.cluster_count) {
48            let maybe_cluster = unsafe { cpuinfo_get_cluster(cid).as_ref() };
49            match maybe_cluster {
50                Some(c) => {
51                    #[allow(unused_assignments)]
52                    let mut cpuid = 0 as u32;
53                    cfg_if::cfg_if! {
54                        if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
55                            cpuid = c.cpuid;
56                        } else if #[cfg(any(target_arch = "arm", target_arch = "aarch64"))] {
57                            cpuid = c.midr;
58                        }
59                    };
60                    let cluster = Cluster {
61                        id: c.cluster_id,
62                        frequency: c.frequency,
63                        cpuid,
64                    };
65
66                    clusters.push(cluster);
67                }
68                None => {}
69            }
70        }
71
72        let package = Package {
73            name: name
74                .to_owned()
75                .to_str()
76                .expect("failed to construct string")
77                .to_owned(),
78            clusters,
79        };
80        sockets.push(package);
81    }
82
83    return Wrapper {
84        packages: sockets,
85        phantom: PhantomData::default(),
86    };
87});
88
89pub fn get_packages() -> &'static [Package] {
90    return &GLOBAL_SOCKETS.packages;
91}
92
93#[cfg(test)]
94mod test {
95    use crate::get_packages;
96
97    #[test]
98    fn basic_info() {
99        let packages = get_packages();
100        assert!(!packages.is_empty());
101        assert!(!packages.first().unwrap().name.is_empty());
102
103        let package = packages.first().unwrap();
104
105        assert!(!package.clusters.is_empty());
106
107        let cluster = package.clusters.first().unwrap();
108
109        assert_ne!(cluster.cpuid, 0);
110        assert_eq!(cluster.frequency, 0);
111    }
112}
pytorch_cpuinfo/wrappers.rs

pytorch_cpuinfo/
wrappers.rs
