cubecl_common/device/handle/
placement.rs

1//! Where a device's server thread runs: under the same last-level cache as the thread issuing
2//! its launches.
3//!
4//! Every launch crosses from the thread that issues it to the device's server thread through
5//! shared memory, hundreds of times a step on a decode path. On a CPU whose cores are split
6//! between several last-level caches (AMD's chiplets, Intel's tiles), two threads under
7//! different caches pay a cross-cache transfer on each of those crossings, and a decode step
8//! ran about 60% slower when the scheduler happened to place them apart. The scheduler may
9//! also move the issuing thread at any time, so the server follows it: the issuing thread
10//! states the CPU it runs on each time it hands over a full buffer of tasks
11//! ([`current_cpu`]), and the server moves to that CPU's cache once the issuer has settled
12//! there ([`ServerPlacement::follow`]).
13
14/// The CPU the calling thread runs on; [`UNKNOWN_CPU`] where the system does not say, and under
15/// Miri, which cannot call `sched_getcpu`.
16pub(crate) fn current_cpu() -> u32 {
17    #[cfg(all(target_os = "linux", not(miri)))]
18    {
19        // SAFETY: `sched_getcpu` takes no arguments.
20        u32::try_from(unsafe { libc::sched_getcpu() }).unwrap_or(UNKNOWN_CPU)
21    }
22    #[cfg(not(all(target_os = "linux", not(miri))))]
23    {
24        UNKNOWN_CPU
25    }
26}
27
28/// What [`current_cpu`] answers where the system does not say.
29pub(crate) const UNKNOWN_CPU: u32 = u32::MAX;
30
31/// The CPUs the process may run on, read once before any thread is kept under one cache.
32///
33/// Threads a device server spawns inherit its cache; those with no launches to stay near, such
34/// as compile workers, spread back over every CPU.
35pub struct ProcessCpus {
36    parallelism: usize,
37    /// `None` where there is only one cache, so no thread is ever kept under one.
38    #[cfg(all(target_os = "linux", not(miri)))]
39    domains: Option<linux::CacheDomains>,
40}
41
42impl ProcessCpus {
43    /// Read on the first call.
44    pub fn get() -> &'static Self {
45        static CPUS: std::sync::OnceLock<ProcessCpus> = std::sync::OnceLock::new();
46        CPUS.get_or_init(|| Self {
47            parallelism: std::thread::available_parallelism().map_or(1, |threads| threads.get()),
48            #[cfg(all(target_os = "linux", not(miri)))]
49            domains: linux::CacheDomains::new(),
50        })
51    }
52
53    /// How many threads the process may run at once.
54    pub fn parallelism(&self) -> usize {
55        self.parallelism
56    }
57
58    /// Lets the calling thread run on every CPU, undoing the cache it inherited from a server.
59    pub fn spread_current_thread(&self) {
60        #[cfg(all(target_os = "linux", not(miri)))]
61        if let Some(domains) = &self.domains {
62            domains.spread_current_thread();
63        }
64    }
65}
66
67/// Buffers in a row an issuer fills from one other last-level cache before the server moves
68/// there. Several threads may issue a device's launches from different caches, each now and
69/// then; moving on every buffer would chase each of them in turn, a migration every time.
70const SETTLED_BUFFERS: u32 = 8;
71
72/// The last-level cache a device's server thread is kept under, and where the issuer has been
73/// filling its buffers from instead, if anywhere.
74#[derive(Default)]
75pub(crate) struct ServerPlacement {
76    /// The cache domain the server is kept under; `None` until the first buffer.
77    kept: Option<usize>,
78    /// A domain other than `kept` the last buffers came from, and how many in a row.
79    elsewhere: Option<(usize, u32)>,
80}
81
82impl ServerPlacement {
83    /// Keeps the calling thread, the server, under `issuer_cpu`'s last-level cache: at once for
84    /// the first buffer, and once [`SETTLED_BUFFERS`] in a row came from a cache other than the
85    /// one it is kept under after that. A no-op on a CPU with one last-level cache over the
86    /// CPUs the process may use, for an unknown CPU, on systems this is not written for, and
87    /// under Miri, which cannot read or set the thread's affinity.
88    pub(crate) fn follow(&mut self, issuer_cpu: u32) {
89        #[cfg(all(target_os = "linux", not(miri)))]
90        {
91            let Some(domains) = &ProcessCpus::get().domains else {
92                return;
93            };
94            if let Some(domain) = domains
95                .of_cpu(issuer_cpu)
96                .and_then(|domain| self.settle(domain))
97            {
98                domains.keep_current_thread(domain);
99            }
100        }
101        #[cfg(not(all(target_os = "linux", not(miri))))]
102        let _ = issuer_cpu;
103    }
104
105    /// Notes that a buffer came from `domain`, and answers the domain the server moves to, if
106    /// it moves.
107    #[cfg_attr(not(all(target_os = "linux", not(miri))), allow(dead_code))]
108    fn settle(&mut self, domain: usize) -> Option<usize> {
109        let moves = match (self.kept, self.elsewhere) {
110            (None, _) => true,
111            (Some(kept), _) if kept == domain => {
112                self.elsewhere = None;
113                false
114            }
115            (Some(_), Some((other, buffers))) if other == domain => {
116                self.elsewhere = Some((domain, buffers + 1));
117                buffers + 1 >= SETTLED_BUFFERS
118            }
119            (Some(_), _) => {
120                self.elsewhere = Some((domain, 1));
121                SETTLED_BUFFERS <= 1
122            }
123        };
124        moves.then(|| {
125            self.kept = Some(domain);
126            self.elsewhere = None;
127            domain
128        })
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn the_first_buffer_places_the_server() {
138        assert_eq!(ServerPlacement::default().settle(3), Some(3));
139    }
140
141    #[test]
142    fn issuers_taking_turns_from_two_caches_leave_the_server_where_it_is() {
143        let mut placement = ServerPlacement::default();
144        placement.settle(1);
145        for _ in 0..100 {
146            assert_eq!(placement.settle(3), None);
147            assert_eq!(placement.settle(1), None);
148        }
149    }
150
151    #[test]
152    fn an_issuer_settled_under_another_cache_draws_the_server_there() {
153        let mut placement = ServerPlacement::default();
154        placement.settle(1);
155        for _ in 1..SETTLED_BUFFERS {
156            assert_eq!(placement.settle(3), None);
157        }
158        assert_eq!(placement.settle(3), Some(3));
159        assert_eq!(placement.settle(3), None);
160    }
161}
162
163#[cfg(all(target_os = "linux", not(miri)))]
164mod linux {
165    use std::{fs, path::Path, string::String, vec::Vec};
166
167    /// The CPUs the process may run on, grouped by the last-level cache they share.
168    pub(super) struct CacheDomains {
169        /// Each CPU's domain, by CPU number; `None` for a CPU the process may not run on.
170        of_cpu: Vec<Option<usize>>,
171        /// Each domain's CPUs.
172        cpus: Vec<CpuSet>,
173    }
174
175    impl CacheDomains {
176        /// The domains of the CPUs the calling thread may run on; `None` where there is only
177        /// one or the system does not say.
178        pub(super) fn new() -> Option<Self> {
179            let allowed = CpuSet::allowed()?;
180            let mut of_cpu = Vec::new();
181            let mut cpus: Vec<CpuSet> = Vec::new();
182            // Each domain's lowest CPU, which names it in sysfs's lists.
183            let mut lowest_cpus: Vec<usize> = Vec::new();
184            for cpu in (0..libc::CPU_SETSIZE as usize).filter(|&cpu| allowed.contains(cpu)) {
185                let lowest = *last_level_cache_cpus(cpu)?.iter().min()?;
186                let domain = match lowest_cpus.iter().position(|&known| known == lowest) {
187                    Some(domain) => domain,
188                    None => {
189                        lowest_cpus.push(lowest);
190                        cpus.push(CpuSet::empty());
191                        cpus.len() - 1
192                    }
193                };
194                cpus[domain].insert(cpu);
195                of_cpu.resize(of_cpu.len().max(cpu + 1), None);
196                of_cpu[cpu] = Some(domain);
197            }
198            (cpus.len() > 1).then_some(Self { of_cpu, cpus })
199        }
200
201        /// The domain of `cpu`, if the process may run on it.
202        pub(super) fn of_cpu(&self, cpu: u32) -> Option<usize> {
203            *self.of_cpu.get(cpu as usize)?
204        }
205
206        /// Keeps the calling thread on `domain`'s CPUs. A refusal leaves the thread where it
207        /// may already run, which is correct, only slower.
208        pub(super) fn keep_current_thread(&self, domain: usize) {
209            let set = &self.cpus[domain].0;
210            // SAFETY: `set` is a valid `cpu_set_t` of the size passed.
211            let kept = unsafe { libc::sched_setaffinity(0, size_of::<libc::cpu_set_t>(), set) };
212            if kept != 0 {
213                log::debug!("The device server stays where it is: sched_setaffinity refused");
214            }
215        }
216
217        /// Lets the calling thread run on every domain's CPUs.
218        pub(super) fn spread_current_thread(&self) {
219            let mut every = CpuSet::empty();
220            for (cpu, domain) in self.of_cpu.iter().enumerate() {
221                if domain.is_some() {
222                    every.insert(cpu);
223                }
224            }
225            // SAFETY: `every` is a valid `cpu_set_t` of the size passed.
226            let spread =
227                unsafe { libc::sched_setaffinity(0, size_of::<libc::cpu_set_t>(), &every.0) };
228            if spread != 0 {
229                log::debug!("The thread stays under its cache: sched_setaffinity refused");
230            }
231        }
232    }
233
234    /// A set of CPUs as the kernel's affinity calls take it.
235    struct CpuSet(libc::cpu_set_t);
236
237    impl CpuSet {
238        fn empty() -> Self {
239            // SAFETY: an all-zero `cpu_set_t` is the empty set.
240            Self(unsafe { core::mem::zeroed() })
241        }
242
243        /// The CPUs the calling thread may run on.
244        fn allowed() -> Option<Self> {
245            let mut set = Self::empty();
246            // SAFETY: `set` is a valid `cpu_set_t` of the size passed.
247            let read =
248                unsafe { libc::sched_getaffinity(0, size_of::<libc::cpu_set_t>(), &mut set.0) };
249            (read == 0).then_some(set)
250        }
251
252        /// Whether `cpu`, below `CPU_SETSIZE`, is in the set.
253        fn contains(&self, cpu: usize) -> bool {
254            // SAFETY: `CPU_ISSET` reads one bit of a valid set.
255            unsafe { libc::CPU_ISSET(cpu, &self.0) }
256        }
257
258        /// Adds `cpu`, below `CPU_SETSIZE`.
259        fn insert(&mut self, cpu: usize) {
260            // SAFETY: `CPU_SET` writes one bit of a valid set.
261            unsafe { libc::CPU_SET(cpu, &mut self.0) }
262        }
263    }
264
265    // SAFETY: a `cpu_set_t` is a plain bit array.
266    unsafe impl Send for CpuSet {}
267    unsafe impl Sync for CpuSet {}
268
269    /// The CPUs sharing `cpu`'s highest-level cache, as sysfs lists them.
270    fn last_level_cache_cpus(cpu: usize) -> Option<Vec<usize>> {
271        let caches = std::format!("/sys/devices/system/cpu/cpu{cpu}/cache");
272        let last = fs::read_dir(&caches)
273            .ok()?
274            .filter_map(|entry| entry.ok())
275            .filter(|entry| entry.file_name().to_string_lossy().starts_with("index"))
276            .filter_map(|entry| {
277                let level = read_trimmed(&entry.path().join("level"))?
278                    .parse::<u32>()
279                    .ok()?;
280                Some((level, entry.path()))
281            })
282            .max_by_key(|(level, _)| *level)?
283            .1;
284        cpu_list(&read_trimmed(&last.join("shared_cpu_list"))?)
285    }
286
287    fn read_trimmed(path: &Path) -> Option<String> {
288        fs::read_to_string(path).ok().map(|text| text.trim().into())
289    }
290
291    /// The CPUs a sysfs list names, `0-7,48-55` being sixteen.
292    fn cpu_list(text: &str) -> Option<Vec<usize>> {
293        let mut cpus = Vec::new();
294        for range in text.split(',').filter(|range| !range.is_empty()) {
295            match range.split_once('-') {
296                Some((first, last)) => {
297                    cpus.extend(first.parse::<usize>().ok()?..=last.parse::<usize>().ok()?)
298                }
299                None => cpus.push(range.parse().ok()?),
300            }
301        }
302        Some(cpus)
303    }
304
305    #[cfg(test)]
306    mod tests {
307        use super::*;
308
309        #[test]
310        fn a_cpu_list_names_its_ranges_and_singles() {
311            assert_eq!(
312                cpu_list("0-3,8,10-11"),
313                Some(std::vec![0, 1, 2, 3, 8, 10, 11])
314            );
315            assert_eq!(cpu_list("5"), Some(std::vec![5]));
316            assert_eq!(cpu_list("0-x"), None);
317        }
318
319        #[test]
320        fn a_server_that_follows_a_cpu_is_kept_under_its_cache() {
321            let Some(domains) = &super::super::ProcessCpus::get().domains else {
322                return;
323            };
324            std::thread::spawn(move || {
325                let cpu = super::super::current_cpu();
326                super::super::ServerPlacement::default().follow(cpu);
327                let domain = domains.of_cpu(cpu).unwrap();
328                let kept = CpuSet::allowed().unwrap();
329                for other in 0..libc::CPU_SETSIZE as usize {
330                    assert_eq!(kept.contains(other), domains.cpus[domain].contains(other));
331                }
332            })
333            .join()
334            .unwrap();
335        }
336
337        /// A compile worker that kept the server's cache would compile on its cores only.
338        #[test]
339        fn a_thread_spawned_by_a_kept_server_spreads_over_every_cpu() {
340            let cpus = super::super::ProcessCpus::get();
341            let Some(domains) = &cpus.domains else {
342                return;
343            };
344            std::thread::spawn(move || {
345                super::super::ServerPlacement::default().follow(super::super::current_cpu());
346                std::thread::spawn(move || {
347                    cpus.spread_current_thread();
348                    let spread = CpuSet::allowed().unwrap();
349                    for cpu in 0..libc::CPU_SETSIZE as usize {
350                        let in_a_domain = domains.of_cpu(cpu as u32).is_some();
351                        assert_eq!(spread.contains(cpu), in_a_domain);
352                    }
353                })
354                .join()
355                .unwrap();
356            })
357            .join()
358            .unwrap();
359        }
360
361        #[test]
362        fn every_allowed_cpu_lies_in_its_own_domain() {
363            let Some(domains) = &super::super::ProcessCpus::get().domains else {
364                return;
365            };
366            let allowed = CpuSet::allowed().unwrap();
367            for cpu in (0..libc::CPU_SETSIZE as usize).filter(|&cpu| allowed.contains(cpu)) {
368                let domain = domains.of_cpu(cpu as u32).unwrap();
369                assert!(domains.cpus[domain].contains(cpu));
370            }
371        }
372    }
373}
cubecl_common/device/handle/placement.rs

cubecl_common/device/handle/
placement.rs
