Skip to content

Commit

Permalink
scx_utils: cleanup gpu topology part
Browse files Browse the repository at this point in the history
code using more modern construct to replace simple match patterns
mostly, simplifying loops.
  • Loading branch information
devnexen committed Jan 23, 2025
1 parent dfe03fe commit 24d023c
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 57 deletions.
79 changes: 38 additions & 41 deletions rust/scx_utils/src/gpu.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,51 +29,48 @@ pub fn create_gpus() -> BTreeMap<usize, Vec<Gpu>> {
let Ok(nvml) = Nvml::init_with_flags(InitFlags::NO_GPUS) else {
return BTreeMap::new();
};
match nvml.device_count() {
Ok(nvidia_gpu_count) => {
for i in 0..nvidia_gpu_count {
let Ok(nvidia_gpu) = nvml.device_by_index(i) else {
continue;
};
let graphics_boost_clock = nvidia_gpu
.max_customer_boost_clock(Clock::Graphics)
.unwrap_or(0);
let sm_boost_clock = nvidia_gpu.max_customer_boost_clock(Clock::SM).unwrap_or(0);
let Ok(memory_info) = nvidia_gpu.memory_info() else {
continue;
};
let Ok(pci_info) = nvidia_gpu.pci_info() else {
continue;
};
let Ok(index) = nvidia_gpu.index() else {
continue;
};
if let Ok(nvidia_gpu_count) = nvml.device_count() {
for i in 0..nvidia_gpu_count {
let Ok(nvidia_gpu) = nvml.device_by_index(i) else {
continue;
};
let graphics_boost_clock = nvidia_gpu
.max_customer_boost_clock(Clock::Graphics)
.unwrap_or(0);
let sm_boost_clock = nvidia_gpu.max_customer_boost_clock(Clock::SM).unwrap_or(0);
let Ok(memory_info) = nvidia_gpu.memory_info() else {
continue;
};
let Ok(pci_info) = nvidia_gpu.pci_info() else {
continue;
};
let Ok(index) = nvidia_gpu.index() else {
continue;
};

// The NVML library doesn't return a PCIe bus ID compatible with sysfs. It includes
// uppercase bus ID values and an extra four leading 0s.
let bus_id = pci_info.bus_id.to_lowercase();
let fixed_bus_id = bus_id.strip_prefix("0000").unwrap_or("");
let numa_path = format!("/sys/bus/pci/devices/{}/numa_node", fixed_bus_id);
let numa_node = read_file_usize(&Path::new(&numa_path)).unwrap_or(0);
// The NVML library doesn't return a PCIe bus ID compatible with sysfs. It includes
// uppercase bus ID values and an extra four leading 0s.
let bus_id = pci_info.bus_id.to_lowercase();
let fixed_bus_id = bus_id.strip_prefix("0000").unwrap_or("");
let numa_path = format!("/sys/bus/pci/devices/{}/numa_node", fixed_bus_id);
let numa_node = read_file_usize(&Path::new(&numa_path)).unwrap_or(0);

let gpu = Gpu {
index: GpuIndex::Nvidia { nvml_id: index },
node_id: numa_node as usize,
max_graphics_clock: graphics_boost_clock as usize,
max_sm_clock: sm_boost_clock as usize,
memory: memory_info.total,
};
if !gpus.contains_key(&numa_node) {
gpus.insert(numa_node, vec![gpu]);
continue;
}
if let Some(gpus) = gpus.get_mut(&numa_node) {
gpus.push(gpu);
}
let gpu = Gpu {
index: GpuIndex::Nvidia { nvml_id: index },
node_id: numa_node as usize,
max_graphics_clock: graphics_boost_clock as usize,
max_sm_clock: sm_boost_clock as usize,
memory: memory_info.total,
};
if !gpus.contains_key(&numa_node) {
gpus.insert(numa_node, vec![gpu]);
continue;
}
if let Some(gpus) = gpus.get_mut(&numa_node) {
gpus.push(gpu);
}
}
_ => {}
};
}

gpus
}
26 changes: 10 additions & 16 deletions rust/scx_utils/src/topology.rs
Original file line number Diff line number Diff line change
Expand Up @@ -270,9 +270,9 @@ impl Topology {
#[cfg(feature = "gpu-topology")]
pub fn gpus(&self) -> BTreeMap<GpuIndex, &Gpu> {
let mut gpus = BTreeMap::new();
for (_, node) in &self.nodes {
for node in self.nodes.values() {
for (idx, gpu) in &node.gpus {
gpus.insert(idx.clone(), gpu);
gpus.insert(*idx, gpu);
}
}
gpus
Expand Down Expand Up @@ -582,14 +582,11 @@ fn create_default_node(
#[cfg(feature = "gpu-topology")]
{
let system_gpus = create_gpus();
match system_gpus.get(&0) {
Some(gpus) => {
for gpu in gpus {
node.gpus.insert(gpu.index, gpu.clone());
}
if let Some(gpus) = system_gpus.get(&0) {
for gpu in gpus {
node.gpus.insert(gpu.index, gpu.clone());
}
_ => {}
};
}
}

if !Path::new("/sys/devices/system/cpu").exists() {
Expand Down Expand Up @@ -647,14 +644,11 @@ fn create_numa_nodes(

#[cfg(feature = "gpu-topology")]
{
match system_gpus.get(&node_id) {
Some(gpus) => {
for gpu in gpus {
node.gpus.insert(gpu.index, gpu.clone());
}
if let Some(gpus) = system_gpus.get(&node_id) {
for gpu in gpus {
node.gpus.insert(gpu.index, gpu.clone());
}
_ => {}
};
}
}

let cpu_pattern = numa_path.join("cpu[0-9]*");
Expand Down

0 comments on commit 24d023c

Please sign in to comment.