From 24d023cf06f02132ebd4fe3decf79ab4cb8d651c Mon Sep 17 00:00:00 2001 From: David Carlier Date: Fri, 17 Jan 2025 07:21:07 +0000 Subject: [PATCH] scx_utils: cleanup gpu topology part code using more modern construct to replace simple match patterns mostly, simplifying loops. --- rust/scx_utils/src/gpu.rs | 79 ++++++++++++++++------------------ rust/scx_utils/src/topology.rs | 26 +++++------ 2 files changed, 48 insertions(+), 57 deletions(-) diff --git a/rust/scx_utils/src/gpu.rs b/rust/scx_utils/src/gpu.rs index b4bbb4d3e..3c8eabc9e 100644 --- a/rust/scx_utils/src/gpu.rs +++ b/rust/scx_utils/src/gpu.rs @@ -29,51 +29,48 @@ pub fn create_gpus() -> BTreeMap> { 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 } diff --git a/rust/scx_utils/src/topology.rs b/rust/scx_utils/src/topology.rs index 59401c0ed..b3bb1b918 100644 --- a/rust/scx_utils/src/topology.rs +++ b/rust/scx_utils/src/topology.rs @@ -270,9 +270,9 @@ impl Topology { #[cfg(feature = "gpu-topology")] pub fn gpus(&self) -> BTreeMap { 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 @@ -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() { @@ -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]*");