Skip to main content

scx_mitosis/
mitosis_topology_utils.rs

1use crate::bpf_intf;
2use anyhow::{bail, Result};
3use scx_utils::Topology;
4use std::collections::HashMap;
5use std::io::{self, BufRead, BufReader};
6use std::path::Path;
7
8use crate::bpf_skel::OpenBpfSkel;
9
10const CPUMASK_LONG_ENTRIES: usize = bpf_intf::consts_CPUMASK_LONG_ENTRIES as usize;
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
13pub enum MapKind {
14    CpuToLLC,
15    LLCToCpus,
16}
17
18impl std::str::FromStr for MapKind {
19    type Err = anyhow::Error;
20    fn from_str(s: &str) -> Result<Self> {
21        match s {
22            "cpu_to_llc" => Ok(MapKind::CpuToLLC),
23            "llc_to_cpus" => Ok(MapKind::LLCToCpus),
24            _ => bail!("unknown map {s}"),
25        }
26    }
27}
28
29impl std::fmt::Display for MapKind {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        f.write_str(match self {
32            MapKind::CpuToLLC => "cpu_to_llc",
33            MapKind::LLCToCpus => "llc_to_cpus",
34        })
35    }
36}
37
38#[allow(dead_code)]
39pub const SUPPORTED_MAPS: &[MapKind] = &[MapKind::CpuToLLC, MapKind::LLCToCpus];
40
41/// Parse lines of the form `cpu,llc` from the provided reader.
42fn parse_cpu_llc_map<R: BufRead>(reader: R) -> Result<Vec<(usize, usize)>> {
43    let mut pairs = Vec::new();
44    for line in reader.lines() {
45        let line = line?;
46        let line = line.trim();
47        // Ignore blank lines and comments
48        if line.is_empty() || line.starts_with('#') {
49            continue;
50        }
51        let mut parts = line.split(',');
52        let cpu = parts
53            .next()
54            .ok_or_else(|| anyhow::anyhow!("missing cpu"))?
55            .trim()
56            .parse::<usize>()?;
57        let llc = parts
58            .next()
59            .ok_or_else(|| anyhow::anyhow!("missing llc"))?
60            .trim()
61            .parse::<usize>()?;
62        pairs.push((cpu, llc));
63    }
64    Ok(pairs)
65}
66
67/// Read CPU/LLC pairs either from a file or standard input.
68fn read_cpu_llc_map(path: &str) -> Result<Vec<(usize, usize)>> {
69    if path == "-" {
70        let stdin = io::stdin();
71        let reader = BufReader::new(stdin.lock());
72        parse_cpu_llc_map(reader)
73    } else {
74        let file = std::fs::File::open(Path::new(path))?;
75        let reader = BufReader::new(file);
76        parse_cpu_llc_map(reader)
77    }
78}
79
80/// Update global arrays for LLC topology before BPF program load.
81/// This function writes directly to the skeleton's BSS section.
82pub fn populate_topology_maps(
83    skel: &mut OpenBpfSkel,
84    map: MapKind,
85    file: Option<String>,
86) -> Result<()> {
87    match map {
88        MapKind::CpuToLLC => {
89            let map_entries = if let Some(path) = file {
90                read_cpu_llc_map(&path)?
91            } else {
92                let topo = Topology::new()?;
93                (0..*scx_utils::NR_CPUS_POSSIBLE)
94                    // Use 0 if a CPU is missing from the topology
95                    .map(|cpu| (cpu, topo.all_cpus.get(&cpu).map(|c| c.llc_id).unwrap_or(0)))
96                    .collect()
97            };
98            let bss = skel
99                .maps
100                .bss_data
101                .as_mut()
102                .ok_or_else(|| anyhow::anyhow!("bss_data not available"))?;
103            for (cpu, llc) in map_entries {
104                if cpu >= bss.cpu_to_llc.len() {
105                    bail!("invalid cpu {cpu}");
106                }
107                bss.cpu_to_llc[cpu] = llc as u32;
108            }
109        }
110        MapKind::LLCToCpus => {
111            if file.is_some() {
112                anyhow::bail!("Loading llc_to_cpus from file is not supported yet");
113            }
114
115            let topo = Topology::new()?;
116
117            // Group CPUs by LLC cache ID
118            let mut llc_to_cpus: HashMap<usize, Vec<usize>> = HashMap::new();
119            for cpu in topo.all_cpus.values() {
120                llc_to_cpus.entry(cpu.llc_id).or_default().push(cpu.id);
121            }
122
123            // For each LLC cache, create a cpumask and populate the array
124            let bss = skel
125                .maps
126                .bss_data
127                .as_mut()
128                .ok_or_else(|| anyhow::anyhow!("bss_data not available"))?;
129            for (llc_id, cpus) in llc_to_cpus {
130                // Create a cpumask structure that matches the BPF side
131                let mut cpumask_longs = [0u64; CPUMASK_LONG_ENTRIES];
132
133                // Set bits for each CPU in this LLC cache
134                for cpu in cpus {
135                    let long_idx = cpu / 64;
136                    let bit_idx = cpu % 64;
137                    if long_idx < CPUMASK_LONG_ENTRIES {
138                        cpumask_longs[long_idx] |= 1u64 << bit_idx;
139                    }
140                }
141                if llc_id >= bss.llc_to_cpus.len() {
142                    bail!("invalid llc_id {llc_id}");
143                }
144
145                bss.llc_to_cpus[llc_id].bits = cpumask_longs;
146            }
147        }
148    }
149    Ok(())
150}
151
152/// Display CPU to LLC cache relationships discovered from the host topology.
153#[allow(dead_code)]
154pub fn print_topology() -> Result<()> {
155    let topo = Topology::new()?;
156    println!("Number LLC caches: {}", topo.all_llcs.len());
157    println!("CPU -> LLC id:");
158    for cpu in topo.all_cpus.values() {
159        println!("cpu {} -> {}", cpu.id, cpu.llc_id);
160    }
161    println!("\nLLC id -> [cpus]:");
162    let mut by_llc: std::collections::BTreeMap<usize, Vec<usize>> =
163        std::collections::BTreeMap::new();
164    for cpu in topo.all_cpus.values() {
165        by_llc.entry(cpu.llc_id).or_default().push(cpu.id);
166    }
167    for (llc, mut cpus) in by_llc {
168        cpus.sort_unstable();
169        println!("{llc} -> {:?}", cpus);
170    }
171    Ok(())
172}