Skip to main content

scx_mitosis/
topology.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2
3// This software may be used and distributed according to the terms of the
4// GNU General Public License version 2.
5
6use std::collections::BTreeMap;
7
8use anyhow::{anyhow, bail, Result};
9use scx_utils::{Cpumask, Topology};
10
11use crate::bpf_intf;
12use crate::bpf_skel::types::llc_cpumask;
13use crate::bpf_skel::OpenBpfSkel;
14
15const MAX_LLCS: usize = bpf_intf::consts_MAX_LLCS as usize;
16const CPUMASK_LONG_ENTRIES: usize = bpf_intf::consts_CPUMASK_LONG_ENTRIES as usize;
17
18pub struct MitosisTopology {
19    pub cpu_to_llc: BTreeMap<usize, usize>,
20    pub llc_to_cpus: BTreeMap<usize, Cpumask>,
21}
22
23impl MitosisTopology {
24    pub fn new(topology: &Topology) -> Self {
25        let mut cpu_to_llc = BTreeMap::new();
26        for cpu in 0..*scx_utils::NR_CPUS_POSSIBLE {
27            cpu_to_llc.insert(
28                cpu,
29                topology
30                    .all_cpus
31                    .get(&cpu)
32                    .map(|cpu| cpu.llc_id)
33                    .unwrap_or(0),
34            );
35        }
36
37        let mut llc_to_cpus = BTreeMap::new();
38        for llc in topology.all_llcs.values() {
39            llc_to_cpus.insert(llc.id, llc.span.clone());
40        }
41
42        MitosisTopology {
43            cpu_to_llc,
44            llc_to_cpus,
45        }
46    }
47
48    pub fn apply(&self, skel: &mut OpenBpfSkel) -> Result<()> {
49        // TODO: rodata
50        let bss = skel
51            .maps
52            .bss_data
53            .as_mut()
54            .ok_or_else(|| anyhow!("bss_data not available"))?;
55        self.apply_topology(&mut bss.cpu_to_llc, &mut bss.llc_to_cpus)
56    }
57
58    fn apply_topology(
59        &self,
60        cpu_to_llc: &mut [u32],
61        llc_to_cpus: &mut [llc_cpumask; MAX_LLCS],
62    ) -> Result<()> {
63        for (&cpu, &llc_id) in &self.cpu_to_llc {
64            if cpu >= cpu_to_llc.len() {
65                bail!("invalid cpu {cpu}");
66            }
67            cpu_to_llc[cpu] = llc_id as u32;
68        }
69
70        for (&llc, cpumask) in &self.llc_to_cpus {
71            if llc >= llc_to_cpus.len() {
72                bail!("invalid llc {llc}");
73            }
74
75            let raw_span = cpumask.as_raw_slice();
76            // Is this overkill?
77            if raw_span.len() > llc_to_cpus[llc].bits.len() {
78                bail!(
79                    "invalid span {llc}: {} > {}",
80                    raw_span.len(),
81                    llc_to_cpus[llc].bits.len()
82                );
83            }
84
85            llc_to_cpus[llc].bits = [0; CPUMASK_LONG_ENTRIES];
86            llc_to_cpus[llc].bits[..raw_span.len()].copy_from_slice(raw_span);
87        }
88        Ok(())
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    fn check_llc_cpu(llc_to_cpus: &[llc_cpumask; MAX_LLCS], llc: usize, cpu: usize) -> bool {
97        let long_idx = cpu / 64;
98        let bit_idx = cpu % 64;
99        llc_to_cpus[llc].bits[long_idx] & (1u64 << bit_idx) != 0
100    }
101
102    #[test]
103    fn test_topology() {
104        let (topology, _) = scx_utils::testutils::make_test_topo(
105            1,  // nodes
106            16, // llcs
107            8,  // cores
108            2,  // cpus
109        );
110        let mut cpu_to_llc_map = BTreeMap::new();
111        for cpu in 0..256 {
112            cpu_to_llc_map.insert(cpu, cpu / 16);
113        }
114        let llc_to_cpus_map = topology
115            .all_llcs
116            .values()
117            .map(|llc| (llc.id, llc.span.clone()))
118            .collect();
119        let mitosis_topology = MitosisTopology {
120            cpu_to_llc: cpu_to_llc_map,
121            llc_to_cpus: llc_to_cpus_map,
122        };
123
124        let mut cpu_to_llc = [0u32; bpf_intf::consts_MAX_CPUS as usize];
125        let mut llc_to_cpus = [llc_cpumask {
126            bits: [0; CPUMASK_LONG_ENTRIES],
127        }; MAX_LLCS];
128
129        mitosis_topology
130            .apply_topology(&mut cpu_to_llc, &mut llc_to_cpus)
131            .unwrap();
132
133        assert_eq!(cpu_to_llc[0], 0);
134        assert_eq!(cpu_to_llc[15], 0);
135        assert_eq!(cpu_to_llc[16], 1);
136        assert_eq!(cpu_to_llc[255], 15);
137
138        for llc in 0..MAX_LLCS {
139            let first_cpu = llc * 16;
140            let last_cpu = first_cpu + 15;
141
142            assert!(check_llc_cpu(&llc_to_cpus, llc, first_cpu));
143            assert!(check_llc_cpu(&llc_to_cpus, llc, last_cpu));
144
145            if llc > 0 {
146                assert!(!check_llc_cpu(&llc_to_cpus, llc, first_cpu - 1));
147            }
148            if llc + 1 < MAX_LLCS {
149                assert!(!check_llc_cpu(&llc_to_cpus, llc, last_cpu + 1));
150            }
151        }
152    }
153}