mmtk/scheduler/
affinity.rs

1use super::worker::ThreadId;
2use crate::util::options::AffinityKind;
3#[cfg(target_os = "linux")]
4use libc::{cpu_set_t, sched_getaffinity, sched_setaffinity, CPU_COUNT, CPU_SET, CPU_ZERO};
5
6/// Represents the ID of a logical CPU on a system.
7pub type CoreId = u16;
8
9// XXX: Maybe in the future we can use a library such as https://github.com/Elzair/core_affinity_rs
10// to have an OS agnostic way of setting thread affinity.
11#[cfg(target_os = "linux")]
12/// Return the total number of cores allocated to the program.
13pub fn get_total_num_cpus() -> u16 {
14    use std::mem::MaybeUninit;
15    unsafe {
16        let mut cs = MaybeUninit::zeroed().assume_init();
17        CPU_ZERO(&mut cs);
18        sched_getaffinity(0, std::mem::size_of::<cpu_set_t>(), &mut cs);
19        CPU_COUNT(&cs) as u16
20    }
21}
22
23#[cfg(not(target_os = "linux"))]
24/// Return the total number of cores allocated to the program.
25pub fn get_total_num_cpus() -> u16 {
26    unimplemented!()
27}
28
29impl AffinityKind {
30    /// Resolve affinity of GC thread. Has a side-effect of calling into the kernel to set the
31    /// thread affinity. Note that we assume that each GC thread is equivalent to an OS or hardware
32    /// thread.
33    pub fn resolve_affinity(&self, thread: ThreadId) {
34        match self {
35            AffinityKind::OsDefault => {}
36            AffinityKind::AllInSet(cpuset) => {
37                // Bind the current thread to all the cores in the set
38                debug!("Set affinity for thread {} to cpuset {:?}", thread, cpuset);
39                bind_current_thread_to_cpuset(cpuset.as_slice());
40            }
41            AffinityKind::RoundRobin(cpuset) => {
42                let cpu = cpuset[thread % cpuset.len()];
43                debug!("Set affinity for thread {} to core {}", thread, cpu);
44                bind_current_thread_to_core(cpu);
45            }
46        }
47    }
48}
49
50#[cfg(target_os = "linux")]
51/// Bind the current thread to the specified core.
52fn bind_current_thread_to_core(cpu: CoreId) {
53    use std::mem::MaybeUninit;
54    unsafe {
55        let mut cs = MaybeUninit::zeroed().assume_init();
56        CPU_ZERO(&mut cs);
57        CPU_SET(cpu as usize, &mut cs);
58        sched_setaffinity(0, std::mem::size_of::<cpu_set_t>(), &cs);
59    }
60}
61
62#[cfg(not(target_os = "linux"))]
63/// Bind the current thread to the specified core.
64fn bind_current_thread_to_core(_cpu: CoreId) {
65    unimplemented!()
66}
67
68#[cfg(any(target_os = "linux", target_os = "android"))]
69/// Bind the current thread to the specified core.
70fn bind_current_thread_to_cpuset(cpuset: &[CoreId]) {
71    use std::mem::MaybeUninit;
72    unsafe {
73        let mut cs = MaybeUninit::zeroed().assume_init();
74        CPU_ZERO(&mut cs);
75        for cpu in cpuset {
76            CPU_SET(*cpu as usize, &mut cs);
77        }
78        sched_setaffinity(0, std::mem::size_of::<cpu_set_t>(), &cs);
79    }
80}
81
82#[cfg(not(any(target_os = "linux", target_os = "android")))]
83/// Bind the current thread to the specified core.
84fn bind_current_thread_to_cpuset(_cpuset: &[CoreId]) {
85    unimplemented!()
86}
mmtk/scheduler/affinity.rs

mmtk/scheduler/
affinity.rs
