feature: Support AMDGPU Data Collection (#1641)

* gpu: support amdgpu tracking Co-authored-by: lvxnull2 <[email protected]> * gpu: dependency-free amdgpu parsing gpu: fix clippy issues Co-authored-by: lvxnull2 <[email protected]> * gpu: change memory usage percentage to be scaled to total memory instead of current memory usage gpu: requested syntax changes Co-authored-by: lvxnull2 <[email protected]> --------- Co-authored-by: lvxnull2 <[email protected]>
ClementTsang · Dec 20, 2024 · 479276b · 479276b
1 parent 797179b
commit 479276b
Show file tree

Hide file tree

Showing 8 changed files with 1,203 additions and 12 deletions.
diff --git a/docs/content/configuration/command-line-options.md b/docs/content/configuration/command-line-options.md
@@ -79,9 +79,9 @@ see information on these options by running `btm -h`, or run `btm --help` to dis
 
 ## GPU Options
 
-| Option          | Behaviour                                                 |
-| --------------- | --------------------------------------------------------- |
-| `--disable_gpu` | Disable collecting and displaying NVIDIA GPU information. |
+| Option          | Behaviour                                                         |
+| --------------- | ----------------------------------------------------------------- |
+| `--disable_gpu` | Disable collecting and displaying NVIDIA and AMD GPU information. |
 
 ## Style Options
 

diff --git a/docs/content/configuration/config-file/flags.md b/docs/content/configuration/config-file/flags.md
@@ -44,7 +44,7 @@ each time:
 | `network_use_binary_prefix`  | Boolean                                                                                                            | Displays the network widget with binary prefixes.                       |
 | `network_use_bytes`          | Boolean                                                                                                            | Displays the network widget using bytes.                                |
 | `network_use_log`            | Boolean                                                                                                            | Displays the network widget with a log scale.                           |
-| `disable_gpu`                | Boolean                                                                                                            | Disable NVIDIA GPU data collection.                                     |
+| `disable_gpu`                | Boolean                                                                                                            | Disable NVIDIA and AMD GPU data collection.                             |
 | `retention`                  | String (human readable time, such as "10m", "1h", etc.)                                                            | How much data is stored at once in terms of time.                       |
 | `unnormalized_cpu`           | Boolean                                                                                                            | Show process CPU% without normalizing over the number of cores.         |
 | `expanded`                   | Boolean                                                                                                            | Expand the default widget upon starting the app.                        |

diff --git a/docs/content/usage/widgets/memory.md b/docs/content/usage/widgets/memory.md
@@ -13,7 +13,7 @@ If the total RAM or swap available is 0, then it is automatically hidden from th
 
 One can also adjust the displayed time range through either the keyboard or mouse, with a range of 30s to 600s.
 
-This widget can also be configured to display Nvidia GPU memory usage (`--disable_gpu` on Linux/Windows to disable) or cache memory usage (`--enable_cache_memory`).
+This widget can also be configured to display Nvidia and AMD GPU memory usage (`--disable_gpu` on Linux/Windows to disable) or cache memory usage (`--enable_cache_memory`).
 
 ## Key bindings
 

diff --git a/docs/content/usage/widgets/temperature.md b/docs/content/usage/widgets/temperature.md
@@ -10,7 +10,7 @@ The temperature widget provides a table of temperature sensors and their current
 
 The temperature widget provides the sensor name as well as its current temperature.
 
-This widget can also be configured to display Nvidia GPU temperatures (`--disable_gpu` on Linux/Windows to disable).
+This widget can also be configured to display Nvidia and AMD GPU temperatures (`--disable_gpu` on Linux/Windows to disable).
 
 ## Key bindings
 

diff --git a/src/data_collection.rs b/src/data_collection.rs
@@ -3,6 +3,9 @@
 #[cfg(feature = "nvidia")]
 pub mod nvidia;
 
+#[cfg(all(target_os = "linux", feature = "gpu"))]
+pub mod amd;
+
 #[cfg(feature = "battery")]
 pub mod batteries;
 
@@ -347,6 +350,10 @@ impl DataCollector {
     #[inline]
     fn update_gpus(&mut self) {
         if self.widgets_to_harvest.use_gpu {
+            let mut local_gpu: Vec<(String, memory::MemHarvest)> = Vec::new();
+            let mut local_gpu_pids: Vec<HashMap<u32, (u64, u32)>> = Vec::new();
+            let mut local_gpu_total_mem: u64 = 0;
+
             #[cfg(feature = "nvidia")]
             if let Some(data) = nvidia::get_nvidia_vecs(
                 &self.temperature_type,
@@ -360,14 +367,41 @@ impl DataCollector {
                         self.data.temperature_sensors = Some(temp);
                     }
                 }
-                if let Some(mem) = data.memory {
-                    self.data.gpu = Some(mem);
+                if let Some(mut mem) = data.memory {
+                    local_gpu.append(&mut mem);
+                }
+                if let Some(mut proc) = data.procs {
+                    local_gpu_pids.append(&mut proc.1);
+                    local_gpu_total_mem += proc.0;
+                }
+            }
+
+            #[cfg(target_os = "linux")]
+            if let Some(data) = amd::get_amd_vecs(
+                &self.temperature_type,
+                &self.filters.temp_filter,
+                &self.widgets_to_harvest,
+                self.last_collection_time,
+            ) {
+                if let Some(mut temp) = data.temperature {
+                    if let Some(sensors) = &mut self.data.temperature_sensors {
+                        sensors.append(&mut temp);
+                    } else {
+                        self.data.temperature_sensors = Some(temp);
+                    }
+                }
+                if let Some(mut mem) = data.memory {
+                    local_gpu.append(&mut mem);
                 }
-                if let Some(proc) = data.procs {
-                    self.gpu_pids = Some(proc.1);
-                    self.gpus_total_mem = Some(proc.0);
+                if let Some(mut proc) = data.procs {
+                    local_gpu_pids.append(&mut proc.1);
+                    local_gpu_total_mem += proc.0;
                 }
             }
+
+            self.data.gpu = (!local_gpu.is_empty()).then_some(local_gpu);
+            self.gpu_pids = (!local_gpu_pids.is_empty()).then_some(local_gpu_pids);
+            self.gpus_total_mem = (local_gpu_total_mem > 0).then_some(local_gpu_total_mem);
         }
     }