monarch_rdma/
device_selection.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! PCI topology parsing and device discovery utilities for RDMA device selection.
10//!
11//! ibverbs-specific selection logic lives in [`crate::backend::ibverbs::device_selection`].
12
13use std::fmt;
14use std::fs;
15use std::path::Path;
16use std::path::PathBuf;
17
18use regex::Regex;
19
20/// PCI address of the CUDA device with ordinal `idx`, read from
21/// `/proc/driver/nvidia/gpus/*/information` (the NVIDIA driver keys each
22/// GPU's bus id there by its device minor, which equals the CUDA ordinal).
23pub fn get_cuda_pci_address(idx: u32) -> Option<PCIAddress> {
24    let gpu_proc_dir = "/proc/driver/nvidia/gpus";
25    if !Path::new(gpu_proc_dir).exists() {
26        return None;
27    }
28
29    let minor_regex =
30        Regex::new(r"Device Minor:\s*(\d+)").expect("should compile: regex literal is valid");
31    for entry in fs::read_dir(gpu_proc_dir).ok()? {
32        let entry = entry.ok()?;
33        let info_file = entry.path().join("information");
34
35        if let Ok(content) = fs::read_to_string(&info_file)
36            && let Some(captures) = minor_regex.captures(&content)
37            && let Ok(device_minor) = captures
38                .get(1)
39                .expect("should be present: capture group 1 matched")
40                .as_str()
41                .parse::<u32>()
42            && device_minor == idx
43        {
44            return PCIAddress::parse(&entry.file_name().to_string_lossy().to_lowercase());
45        }
46    }
47    None
48}
49
50/// A PCI address, e.g. `0000:07:00.0`, as found under
51/// `/sys/bus/pci/devices`.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
53pub struct PCIAddress {
54    pub domain: u16,
55    pub bus: u8,
56    pub device: u8,
57    pub function: u8,
58}
59
60impl PCIAddress {
61    /// Parse a `dddd:bb:dd.f` lowercase-hex PCI address.
62    pub fn parse(s: &str) -> Option<Self> {
63        let (domain, rest) = s.split_once(':')?;
64        let (bus, rest) = rest.split_once(':')?;
65        let (device, function) = rest.split_once('.')?;
66        Some(Self {
67            domain: u16::from_str_radix(domain, 16).ok()?,
68            bus: u8::from_str_radix(bus, 16).ok()?,
69            device: u8::from_str_radix(device, 16).ok()?,
70            function: u8::from_str_radix(function, 16).ok()?,
71        })
72    }
73
74    /// This device's sysfs directory, `/sys/bus/pci/devices/<bdf>`.
75    pub fn sysfs_path(&self) -> PathBuf {
76        Path::new("/sys/bus/pci/devices").join(self.to_string())
77    }
78}
79
80impl fmt::Display for PCIAddress {
81    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82        write!(
83            f,
84            "{:04x}:{:02x}:{:02x}.{:x}",
85            self.domain, self.bus, self.device, self.function
86        )
87    }
88}
89
90/// A source of memory for a transfer. A `None` index means "any device of
91/// this kind": the location is then ranked against the best of all CPU
92/// nodes or all GPUs.
93#[derive(
94    Debug,
95    Clone,
96    Copy,
97    PartialEq,
98    Eq,
99    Hash,
100    serde::Serialize,
101    serde::Deserialize
102)]
103pub enum MemoryLocation {
104    /// CPU memory on the given NUMA node, or any CPU node if `None`.
105    Cpu(Option<u32>),
106    /// GPU memory on the given CUDA device ordinal, or any GPU if `None`.
107    Gpu(Option<u32>),
108}
109
110/// Locality of a path between two PCI endpoints, ordered best to worst.
111/// A path's type is its worst (least local) segment.
112#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
113pub enum PathType {
114    /// Within a single PCIe switch.
115    Pix,
116    /// Across multiple PCIe switches under one host bridge.
117    Pxb,
118    /// Up through the CPU host bridge, within one NUMA node.
119    Phb,
120    /// Across the inter-socket / cross-NUMA interconnect.
121    Sys,
122    /// No path between the endpoints.
123    Dis,
124}
125
126/// A classified path between two PCI endpoints: its [`PathType`] and the
127/// bottleneck (minimum) PCIe link bandwidth along it, in MB/s.
128#[derive(Debug, Clone, Copy, PartialEq, Eq)]
129pub struct PciPath {
130    pub path_type: PathType,
131    pub bottleneck_mbytes_per_sec: u32,
132}
133
134impl PciPath {
135    /// Whether this path is preferable to `other`: more local (lower
136    /// [`PathType`]) wins, and among equally-local paths the higher
137    /// bottleneck bandwidth wins.
138    pub fn is_better_than(&self, other: &PciPath) -> bool {
139        (self.path_type, other.bottleneck_mbytes_per_sec)
140            < (other.path_type, self.bottleneck_mbytes_per_sec)
141    }
142}
143
144/// The [`PathType`] and bottleneck bandwidth of the PCIe path between two
145/// PCI endpoints.
146///
147/// Walks each endpoint's sysfs ancestor chain toward the root complex,
148/// finds their lowest common ancestor, and takes the minimum link
149/// bandwidth along the way. When the endpoints share no PCIe ancestor the
150/// path runs through the CPU: same NUMA node → [`PathType::Phb`],
151/// different nodes → [`PathType::Sys`], unknown → [`PathType::Dis`].
152pub fn pci_path(a: &PCIAddress, b: &PCIAddress) -> PciPath {
153    classify(
154        &ancestor_chain(a),
155        numa_node(a),
156        &ancestor_chain(b),
157        numa_node(b),
158    )
159}
160
161/// Path from CPU memory to the device at `addr`. With a specific NUMA
162/// node, the device is [`PathType::Phb`] when it sits on that node and
163/// [`PathType::Sys`] otherwise; with `None` (any CPU) it is judged against
164/// its own node and so is always [`PathType::Phb`]. Bandwidth is the
165/// device's PCIe chain bottleneck.
166pub fn cpu_path(addr: &PCIAddress, numa: Option<u32>) -> PciPath {
167    let bottleneck_mbytes_per_sec = min_link_mbytes_per_sec(&ancestor_chain(addr));
168    let path_type = match numa {
169        Some(node) if numa_node(addr) != Some(node) => PathType::Sys,
170        _ => PathType::Phb,
171    };
172    PciPath {
173        path_type,
174        bottleneck_mbytes_per_sec,
175    }
176}
177
178/// One node in a device's PCIe ancestor chain: its resolved sysfs path
179/// and the bandwidth (MB/s) of the PCIe link immediately upstream of it.
180struct PciHop {
181    sysfs: PathBuf,
182    link_mbytes_per_sec: u32,
183}
184
185/// Classify the path between two ancestor chains (device → root complex).
186/// Pure over its inputs, so it is unit-tested without touching sysfs.
187fn classify(a: &[PciHop], numa_a: Option<u32>, b: &[PciHop], numa_b: Option<u32>) -> PciPath {
188    if let Some((ia, ib)) = common_ancestor(a, b) {
189        // The path traverses the links upstream of each hop below the common
190        // ancestor (`a[..ia]` then `b[..ib]`). The ancestor's own upstream
191        // link runs further up the tree and is not part of the path, so it is
192        // excluded.
193        let bottleneck_mbytes_per_sec = a[..ia]
194            .iter()
195            .chain(&b[..ib])
196            .map(|h| h.link_mbytes_per_sec)
197            .min()
198            .unwrap_or(0);
199        // A PCIe switch spans two sysfs hops (its downstream and upstream
200        // ports), so meeting within two hops of the common ancestor means
201        // both devices sit under one switch; farther means multiple.
202        let single_switch = ia <= 2 && ib <= 2;
203        let path_type = if single_switch {
204            PathType::Pix
205        } else {
206            PathType::Pxb
207        };
208        return PciPath {
209            path_type,
210            bottleneck_mbytes_per_sec,
211        };
212    }
213    // No shared PCIe ancestor: the path runs through the CPU.
214    let bottleneck_mbytes_per_sec = min_link_mbytes_per_sec(a).min(min_link_mbytes_per_sec(b));
215    let path_type = match (numa_a, numa_b) {
216        (Some(x), Some(y)) if x == y => PathType::Phb,
217        (Some(_), Some(_)) => PathType::Sys,
218        _ => PathType::Dis,
219    };
220    PciPath {
221        path_type,
222        bottleneck_mbytes_per_sec,
223    }
224}
225
226/// Indices of the deepest sysfs path common to both chains, scanning each
227/// from its device end.
228fn common_ancestor(a: &[PciHop], b: &[PciHop]) -> Option<(usize, usize)> {
229    a.iter().enumerate().find_map(|(ia, na)| {
230        b.iter()
231            .position(|nb| nb.sysfs == na.sysfs)
232            .map(|ib| (ia, ib))
233    })
234}
235
236/// Minimum upstream link bandwidth (MB/s) across `hops`. A hop whose
237/// bandwidth couldn't be read is 0 and drags the whole range to 0, so a
238/// path with an unmeasurable link is treated as the worst case. 0 when
239/// `hops` is empty.
240fn min_link_mbytes_per_sec(hops: &[PciHop]) -> u32 {
241    hops.iter()
242        .map(|h| h.link_mbytes_per_sec)
243        .min()
244        .unwrap_or(0)
245}
246
247/// Walk `addr`'s PCIe ancestor chain from the device up toward the root
248/// complex, resolving sysfs symlinks. Empty if the device's sysfs entry
249/// can't be resolved.
250fn ancestor_chain(addr: &PCIAddress) -> Vec<PciHop> {
251    let mut chain = Vec::new();
252    let mut current = match fs::canonicalize(addr.sysfs_path()) {
253        Ok(p) => p,
254        Err(_) => return chain,
255    };
256    loop {
257        let link_mbytes_per_sec = link_bandwidth_mbytes_per_sec(&current);
258        let parent = current.parent().map(Path::to_path_buf);
259        chain.push(PciHop {
260            sysfs: current,
261            link_mbytes_per_sec,
262        });
263        match parent {
264            Some(parent) if is_pci_bdf(&parent) => current = parent,
265            _ => break,
266        }
267    }
268    chain
269}
270
271/// Whether `path`'s final component parses as a PCI address
272/// (`dddd:bb:dd.f`), e.g. `0000:00:01.0`. The root complex (`pci0000:00`)
273/// does not, which stops the ancestor walk there.
274fn is_pci_bdf(path: &Path) -> bool {
275    path.file_name()
276        .and_then(|n| n.to_str())
277        .is_some_and(|n| PCIAddress::parse(n).is_some())
278}
279
280/// NUMA node of the device at `addr`, from `<sysfs>/numa_node`. A `-1`
281/// ("unknown") maps to `None`.
282fn numa_node(addr: &PCIAddress) -> Option<u32> {
283    let raw = fs::read_to_string(addr.sysfs_path().join("numa_node")).ok()?;
284    u32::try_from(raw.trim().parse::<i32>().ok()?).ok()
285}
286
287/// Bandwidth (MB/s) of the PCIe link immediately upstream of the device at
288/// `sysfs`, from its own `max_link_speed` / `max_link_width`. An unreadable
289/// value is 0, so the link is treated as the worst case.
290fn link_bandwidth_mbytes_per_sec(sysfs: &Path) -> u32 {
291    let speed = read_speed_mbits_per_lane(sysfs);
292    let width = read_link_width(sysfs);
293    // `speed` is effective megabits/s per lane (PCIe line-encoding overhead
294    // is already folded into the table), so `speed * width` is the link's
295    // total megabits/s; dividing by 8 converts bits to bytes → MB/s.
296    speed.saturating_mul(width) / 8
297}
298
299/// PCIe lane count from `<dir>/max_link_width`, or 0 if unreadable.
300fn read_link_width(dir: &Path) -> u32 {
301    fs::read_to_string(dir.join("max_link_width"))
302        .ok()
303        .and_then(|s| s.trim().parse().ok())
304        .unwrap_or(0)
305}
306
307/// Per-lane PCIe rate (Mbit/s) from `<dir>/max_link_speed`, or 0 if unreadable.
308fn read_speed_mbits_per_lane(dir: &Path) -> u32 {
309    fs::read_to_string(dir.join("max_link_speed"))
310        .ok()
311        .map(|s| pcie_speed_mbits_per_lane(&s))
312        .unwrap_or(0)
313}
314
315/// Per-lane PCIe bandwidth (Mbit/s) for a `max_link_speed` string such as
316/// `"16 GT/s PCIe"`, with line-encoding overhead folded in. `rate * lanes
317/// / 8` gives the link's MB/s. The values and the Gen3 fallback mirror
318/// NCCL's `kvDictPciGen` (graph/topo.cc).
319fn pcie_speed_mbits_per_lane(speed: &str) -> u32 {
320    // Match the leading "<rate> GT/s" token; the kernel may append a
321    // trailing "PCIe" and prints either "8" or "8.0" style rates.
322    match speed.split_whitespace().next().unwrap_or("") {
323        "2.5" => 1500,          // Gen1
324        "5" | "5.0" => 3000,    // Gen2
325        "8" | "8.0" => 6000,    // Gen3
326        "16" | "16.0" => 12000, // Gen4
327        "32" | "32.0" => 24000, // Gen5
328        "64" | "64.0" => 48000, // Gen6
329        _ => 6000,
330    }
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336
337    #[test]
338    fn test_pci_address_parse_and_display() {
339        let addr = PCIAddress::parse("0000:07:00.0").unwrap();
340        assert_eq!(
341            addr,
342            PCIAddress {
343                domain: 0,
344                bus: 7,
345                device: 0,
346                function: 0,
347            }
348        );
349        assert_eq!(addr.to_string(), "0000:07:00.0");
350        // Hex components round-trip through Display.
351        assert_eq!(
352            PCIAddress::parse("00ff:1a:1f.7").unwrap().to_string(),
353            "00ff:1a:1f.7"
354        );
355        assert_eq!(PCIAddress::parse("not-an-address"), None);
356        assert_eq!(PCIAddress::parse("0000:07:00"), None);
357    }
358
359    #[test]
360    fn test_path_type_orders_best_to_worst() {
361        assert!(PathType::Pix < PathType::Pxb);
362        assert!(PathType::Pxb < PathType::Phb);
363        assert!(PathType::Phb < PathType::Sys);
364        assert!(PathType::Sys < PathType::Dis);
365    }
366
367    #[test]
368    fn test_pci_path_is_better_than() {
369        let pix_slow = PciPath {
370            path_type: PathType::Pix,
371            bottleneck_mbytes_per_sec: 1000,
372        };
373        let pix_fast = PciPath {
374            path_type: PathType::Pix,
375            bottleneck_mbytes_per_sec: 2000,
376        };
377        let phb_fast = PciPath {
378            path_type: PathType::Phb,
379            bottleneck_mbytes_per_sec: 9000,
380        };
381        // A more local path wins regardless of bandwidth.
382        assert!(pix_slow.is_better_than(&phb_fast));
383        assert!(!phb_fast.is_better_than(&pix_slow));
384        // Equal locality: higher bandwidth wins.
385        assert!(pix_fast.is_better_than(&pix_slow));
386        assert!(!pix_slow.is_better_than(&pix_fast));
387        // A path is not strictly better than itself.
388        assert!(!pix_fast.is_better_than(&pix_fast));
389    }
390
391    #[test]
392    fn test_pcie_speed_mbits_per_lane() {
393        assert_eq!(pcie_speed_mbits_per_lane("2.5 GT/s PCIe"), 1500);
394        assert_eq!(pcie_speed_mbits_per_lane("5 GT/s"), 3000);
395        assert_eq!(pcie_speed_mbits_per_lane("8.0 GT/s"), 6000);
396        assert_eq!(pcie_speed_mbits_per_lane("16 GT/s PCIe"), 12000);
397        assert_eq!(pcie_speed_mbits_per_lane("32 GT/s"), 24000);
398        assert_eq!(pcie_speed_mbits_per_lane("64 GT/s"), 48000);
399        // Unrecognized / empty defaults to Gen3.
400        assert_eq!(pcie_speed_mbits_per_lane("garbage"), 6000);
401        assert_eq!(pcie_speed_mbits_per_lane(""), 6000);
402    }
403
404    fn hop(sysfs: &str, link_mbytes_per_sec: u32) -> PciHop {
405        PciHop {
406            sysfs: PathBuf::from(sysfs),
407            link_mbytes_per_sec,
408        }
409    }
410
411    #[test]
412    fn test_classify_pix_single_switch() {
413        // Both devices meet at a shared switch within two hops: PIX, with
414        // the bottleneck being the slowest link to that switch.
415        let a = vec![
416            hop("/d/a", 12000),
417            hop("/d/sw_a", 12000),
418            hop("/d/sw", 8000),
419        ];
420        let b = vec![hop("/d/b", 12000), hop("/d/sw_b", 6000), hop("/d/sw", 8000)];
421        let path = classify(&a, Some(0), &b, Some(0));
422        assert_eq!(path.path_type, PathType::Pix);
423        assert_eq!(path.bottleneck_mbytes_per_sec, 6000);
424    }
425
426    #[test]
427    fn test_classify_excludes_common_ancestor_upstream() {
428        // The common ancestor's own upstream link runs further up the tree and
429        // is not part of the path between `a` and `b`, so a slow link there
430        // must not lower the bottleneck.
431        let a = vec![hop("/d/a", 12000), hop("/d/sw", 2000)];
432        let b = vec![hop("/d/b", 12000), hop("/d/sw", 2000)];
433        let path = classify(&a, Some(0), &b, Some(0));
434        assert_eq!(path.path_type, PathType::Pix);
435        assert_eq!(
436            path.bottleneck_mbytes_per_sec, 12000,
437            "the common ancestor's 2000 MB/s upstream link is above the path and must be excluded"
438        );
439    }
440
441    #[test]
442    fn test_classify_pxb_multiple_switches() {
443        // The chains meet only at the root, several hops up one side: PXB.
444        let a = vec![
445            hop("/d/a", 12000),
446            hop("/d/sw_a1", 12000),
447            hop("/d/sw_a2", 12000),
448            hop("/d/sw_a3", 12000),
449            hop("/d/root", 16000),
450        ];
451        let b = vec![
452            hop("/d/b", 10000),
453            hop("/d/sw_b", 10000),
454            hop("/d/root", 16000),
455        ];
456        let path = classify(&a, Some(0), &b, Some(0));
457        assert_eq!(path.path_type, PathType::Pxb);
458        assert_eq!(path.bottleneck_mbytes_per_sec, 10000);
459    }
460
461    #[test]
462    fn test_classify_through_cpu() {
463        // No shared PCIe ancestor — the path goes through the CPU.
464        let a = vec![hop("/d/a", 4000), hop("/d/root_a", 16000)];
465        let b = vec![hop("/d/b", 8000), hop("/d/root_b", 16000)];
466        // Same NUMA node → PHB; bottleneck is the slowest link overall.
467        let phb = classify(&a, Some(0), &b, Some(0));
468        assert_eq!(phb.path_type, PathType::Phb);
469        assert_eq!(phb.bottleneck_mbytes_per_sec, 4000);
470        // Different NUMA nodes → SYS.
471        assert_eq!(classify(&a, Some(0), &b, Some(1)).path_type, PathType::Sys);
472        // Unknown NUMA node → DIS.
473        assert_eq!(classify(&a, None, &b, Some(0)).path_type, PathType::Dis);
474    }
475}
monarch_rdma/device_selection.rs

monarch_rdma/
device_selection.rs
