Skip to main content

dupehell_core/
cpu_affinity.rs

1// DupeHell -- MIT License
2//
3// Synthetic multi-domain dataset generator for record linkage benchmarking.
4// No liability for misuse.
5
6//! Optional opt-in pinning of the process to "performance" cores on
7//! hybrid (P-core/E-core) CPUs.
8//!
9//! On hybrid topologies, Rayon workers that land on an E-core become
10//! stragglers relative to their P-core siblings at every parallel sync
11//! point, which can degrade throughput disproportionately at high record
12//! counts. This restricts the process (and therefore Rayon's default
13//! global pool, spawned lazily on first use) to the highest
14//! `EfficiencyClass` reported by Windows. No-op, returns `false`, on
15//! non-Windows targets or non-hybrid CPUs.
16
17#[cfg(target_os = "windows")]
18pub fn pin_to_p_cores() -> bool {
19    use windows::Win32::Foundation::HANDLE;
20    use windows::Win32::System::SystemInformation::{
21        GetLogicalProcessorInformationEx, RelationProcessorCore,
22        SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX,
23    };
24    use windows::Win32::System::Threading::{GetCurrentProcess, SetProcessAffinityMask};
25
26    // First call to learn the required buffer size.
27    let mut len: u32 = 0;
28    unsafe {
29        let _ = GetLogicalProcessorInformationEx(RelationProcessorCore, None, &mut len);
30    }
31    if len == 0 {
32        return false;
33    }
34
35    let mut buf: Vec<u8> = vec![0u8; len as usize];
36    let ok = unsafe {
37        GetLogicalProcessorInformationEx(
38            RelationProcessorCore,
39            Some(
40                buf.as_mut_ptr()
41                    .cast::<SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>(),
42            ),
43            &mut len,
44        )
45    };
46    if ok.is_err() {
47        return false;
48    }
49
50    // Walk the variable-length list of entries, each describing one
51    // physical core (EfficiencyClass + the logical-processor mask of its
52    // hyperthread siblings, group 0 only -- fine below 64 logical CPUs).
53    let mut offset = 0usize;
54    let mut max_class: u8 = 0;
55    let mut entries: Vec<(u8, u64)> = Vec::new();
56    while offset < buf.len() {
57        let entry = unsafe {
58            &*(buf.as_ptr().add(offset) as *const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)
59        };
60        let size = entry.Size as usize;
61        if size == 0 {
62            break;
63        }
64        let core = unsafe { &entry.Anonymous.Processor };
65        let class = core.EfficiencyClass;
66        let mask = core.GroupMask[0].Mask as u64;
67        max_class = max_class.max(class);
68        entries.push((class, mask));
69        offset += size;
70    }
71
72    // No hybrid topology (single efficiency class) -- nothing to gain.
73    if entries.iter().all(|(class, _)| *class == entries[0].0) {
74        return false;
75    }
76
77    let p_core_mask: u64 = entries
78        .iter()
79        .filter(|(class, _)| *class == max_class)
80        .fold(0u64, |acc, (_, mask)| acc | mask);
81    if p_core_mask == 0 {
82        return false;
83    }
84
85    let handle: HANDLE = unsafe { GetCurrentProcess() };
86    unsafe { SetProcessAffinityMask(handle, p_core_mask as usize) }.is_ok()
87}
88
89#[cfg(not(target_os = "windows"))]
90pub fn pin_to_p_cores() -> bool {
91    false
92}