pciclient: add AMD GPU detection via PCI enumeration

Signed-off-by: Gaurav Sharma <[email protected]>

Author: mgsharm

sources/pciclient/src/lib.rs

@@ -8,8 +8,8 @@ pciclient provides util functions that can:
 mod private;
 
 use private::{
-    call_list_devices, check_efa_attachment, check_inf1_attachment, check_inf2_attachment,
-    check_neuron_attachment, PciClient,
+    call_list_devices, check_amd_gpu_attachment, check_efa_attachment, check_inf1_attachment,
+    check_inf2_attachment, check_neuron_attachment, PciClient,
 };
 
 use bon::Builder;
@@ -123,6 +123,11 @@ pub fn is_neuron_attached() -> Result<bool> {
     check_neuron_attachment(PciClient {})
 }
 
+/// Call `lspci` and check if there is any AMD GPU device attached.
+pub fn is_amd_gpu_attached() -> Result<bool> {
+    check_amd_gpu_attachment(PciClient {})
+}
+
 /// Call `lspci` and check if there are inf1 devices attached
 pub fn is_inf1_instance() -> Result<bool> {
     check_inf1_attachment(PciClient {})

sources/pciclient/src/private.rs

@@ -14,6 +14,7 @@ use crate::{
 
 const AMAZON_VENDOR_CODE: &str = "1d0f";
 const EFA_KEYWORD: &str = "efa";
+const AMD_VENDOR_CODE: &str = "1002";
 
 // Neuron devices specific to inf1 instance types
 lazy_static! {
@@ -34,6 +35,12 @@ lazy_static! {
     };
 }
 
+lazy_static! {
+    static ref AMD_GPU_DEVICES: HashSet<&'static str> = hashset! {
+        "75a3" // AMD Instinct MI355X
+    };
+}
+
 const LSPCI_PATH: &str = "/usr/bin/lspci";
 
 pub(crate) trait CommandExecutor {
@@ -183,6 +190,19 @@ pub(crate) fn check_neuron_attachment<T: CommandExecutor>(command_executor: T) -
     }))
 }
 
+/// Call `lspci` and check if there is any AMD GPU device attached.
+/// Internal usage, adding command_executor as parameter allows us to better unit test.
+/// For external usage, check [`crate::is_amd_gpu_attached`].
+pub(crate) fn check_amd_gpu_attachment<T: CommandExecutor>(command_executor: T) -> Result<bool> {
+    let list_devices_param = ListDevicesParam::builder()
+        .vendor(AMD_VENDOR_CODE.to_string())
+        .build();
+    let list_device_results = call_list_devices(command_executor, list_devices_param)?;
+    Ok(list_device_results
+        .iter()
+        .any(|device_info| AMD_GPU_DEVICES.contains(&device_info.device.as_str())))
+}
+
 /// Call `lspci` and check if there is any inf1-specific Neuron device attached.
 /// Internal usage, adding command_executor as parameter allows us to better unit test.
 /// For external usage, check [`crate::is_inf1_instance`].
@@ -228,8 +248,8 @@ mod test {
     use crate::{ListDevicesOutput, ListDevicesParam, Result};
 
     use super::{
-        call_list_devices, check_efa_attachment, check_inf1_attachment, check_neuron_attachment,
-        parse_list_devices_output, CommandExecutor,
+        call_list_devices, check_amd_gpu_attachment, check_efa_attachment, check_inf1_attachment,
+        check_neuron_attachment, parse_list_devices_output, CommandExecutor,
     };
 
     struct MockPciClient {
@@ -436,6 +456,31 @@ mod test {
         assert!(check_inf1_attachment_result.unwrap());
     }
 
+    #[test]
+    fn test_is_amd_gpu_attached() {
+        let mock_pci_client = MockPciClient {
+            // AMD Instinct MI355X device has vendor 1002 for AMD, device code 75a3.
+            output: vec![
+                r#"00:1e.0 "0300" "1002" "75a3" -p00 "1002" "0123""#.to_string(),
+                r#"00:1f.0 "0302" "10de" "1eb8" -ra1 -p00 "10de" "12a2""#.to_string(),
+            ],
+        };
+        let check_amd_gpu_attachment_result = check_amd_gpu_attachment(mock_pci_client);
+        assert!(check_amd_gpu_attachment_result.is_ok());
+        assert!(check_amd_gpu_attachment_result.unwrap());
+    }
+
+    #[test]
+    fn test_is_amd_gpu_attached_negative_case() {
+        let mock_pci_client = MockPciClient {
+            // Below is an actual output from lspci for ena device (not AMD GPU).
+            output: vec![r#"00:06.0 "0200" "1d0f" "ec20" -p00 "1d0f" "ec20""#.to_string()],
+        };
+        let check_amd_gpu_attachment_result = check_amd_gpu_attachment(mock_pci_client);
+        assert!(check_amd_gpu_attachment_result.is_ok());
+        assert!(!check_amd_gpu_attachment_result.unwrap());
+    }
+
     #[test]
     fn test_is_inf1_attached_negative_case() {
         let mock_pci_client = MockPciClient {