fix: add vendor in gpu info (#1952)

Co-authored-by: vansangpfiev <[email protected]>
janhq · Feb 11, 2025 · cdb2446 · cdb2446
1 parent 4be3a74
commit cdb2446
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 24 deletions.
diff --git a/engine/common/hardware_common.h b/engine/common/hardware_common.h
@@ -79,6 +79,7 @@ struct GPU {
   int64_t total_vram;
   std::string uuid;
   bool is_activated = true;
+  std::string vendor;
 };
 
 inline Json::Value ToJson(const std::vector<GPU>& gpus) {
@@ -100,6 +101,7 @@ inline Json::Value ToJson(const std::vector<GPU>& gpus) {
     gpu["total_vram"] = gpus[i].total_vram;
     gpu["uuid"] = gpus[i].uuid;
     gpu["activated"] = gpus[i].is_activated;
+    gpu["vendor"] = gpus[i].vendor;
     if (gpus[i].total_vram > 0) {
       res.append(gpu);
     }

diff --git a/engine/utils/hardware/gpu/vulkan/vulkan_gpu.h b/engine/utils/hardware/gpu/vulkan/vulkan_gpu.h
@@ -24,6 +24,21 @@
 #endif
 
 namespace cortex::hw {
+inline std::string GetVendorStr(uint32_t vendor_id) {
+  switch (vendor_id) {
+    case 0x1002:
+      return "AMD";
+    case 0x10DE:
+      return "NVIDIA";
+    case 0x8086:
+      return "INTEL";
+    case 0x13B5:
+      return "ARM";
+    default:
+      return std::to_string(vendor_id);
+  }
+}
+
 #if defined(_WIN32)
 // Definitions of the used function pointers. Add more if you use other ADL APIs
 typedef int (*ADL_MAIN_CONTROL_CREATE)(ADL_MAIN_MALLOC_CALLBACK, int);
@@ -335,7 +350,8 @@ inline cpp::result<std::vector<cortex::hw::GPU>, std::string> GetGpuInfoList() {
 
   // Get the physical devices
   uint32_t physical_device_count = 0;
-  result = vkEnumeratePhysicalDevices(instance, &physical_device_count, nullptr);
+  result =
+      vkEnumeratePhysicalDevices(instance, &physical_device_count, nullptr);
   if (result != VK_SUCCESS) {
     vkDestroyInstance(instance, nullptr);
     FreeLibrary(vulkan_library);
@@ -376,7 +392,8 @@ inline cpp::result<std::vector<cortex::hw::GPU>, std::string> GetGpuInfoList() {
     VkPhysicalDeviceProperties2 device_properties2 = {};
     device_properties2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;
     device_properties2.pNext = &device_id_properties;
-    device_id_properties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ID_PROPERTIES;
+    device_id_properties.sType =
+        VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ID_PROPERTIES;
 
     vkGetPhysicalDeviceProperties2(physical_device, &device_properties2);
 
@@ -404,15 +421,18 @@ inline cpp::result<std::vector<cortex::hw::GPU>, std::string> GetGpuInfoList() {
 #endif
     int free_vram_MiB =
         total_vram_MiB > used_vram_MiB ? total_vram_MiB - used_vram_MiB : 0;
-    gpus.emplace_back(cortex::hw::GPU{
-        .id = std::to_string(id),
-        .device_id = device_properties.deviceID,
-        .name = device_properties.deviceName,
-        .version = std::to_string(device_properties.driverVersion),
-        .add_info = cortex::hw::AmdAddInfo{},
-        .free_vram = free_vram_MiB,
-        .total_vram = total_vram_MiB,
-        .uuid = uuid_to_string(device_id_properties.deviceUUID)});
+    if (total_vram_MiB > 0) {
+      gpus.emplace_back(cortex::hw::GPU{
+          .id = std::to_string(id),
+          .device_id = device_properties.deviceID,
+          .name = device_properties.deviceName,
+          .version = std::to_string(device_properties.driverVersion),
+          .add_info = cortex::hw::AmdAddInfo{},
+          .free_vram = free_vram_MiB,
+          .total_vram = total_vram_MiB,
+          .uuid = uuid_to_string(device_id_properties.deviceUUID),
+          .vendor = GetVendorStr(device_properties.vendorID)});
+    }
     id++;
   }
 

diff --git a/engine/utils/hardware/gpu_info.h b/engine/utils/hardware/gpu_info.h
@@ -24,10 +24,11 @@ inline std::vector<GPU> GetGPUInfo() {
             .compute_cap = nvidia_gpus[i].compute_cap.value_or("unknown")};
         vulkan_gpus[j].free_vram = std::stoll(nvidia_gpus[i].vram_free);
         vulkan_gpus[j].total_vram = std::stoll(nvidia_gpus[i].vram_total);
+        vulkan_gpus[j].vendor = nvidia_gpus[i].vendor;
       }
     }
   }
-  
+
   if (use_vulkan_info) {
     return vulkan_gpus;
   } else {
@@ -43,7 +44,8 @@ inline std::vector<GPU> GetGPUInfo() {
                       .compute_cap = n.compute_cap.value_or("unknown")},
               .free_vram = std::stoi(n.vram_free),
               .total_vram = std::stoi(n.vram_total),
-              .uuid = n.uuid});
+              .uuid = n.uuid,
+              .vendor = n.vendor});
     }
     return res;
   }

diff --git a/engine/utils/system_info_utils.cc b/engine/utils/system_info_utils.cc
@@ -124,17 +124,16 @@ std::vector<GpuInfo> GetGpuInfoList() {
 
     while (
         std::regex_search(search_start, output.cend(), match, gpu_info_reg)) {
-      GpuInfo gpuInfo = {
-          match[1].str(),                        // id
-          match[2].str(),                        // vram_total
-          match[3].str(),                        // vram_free
-          match[4].str(),                        // name
-          GetGpuArch(match[4].str()),            // arch
-          driver_version,                        // driver_version
-          cuda_version,                          // cuda_driver_version
-          need_fallback ? "0" : match[5].str(),  // compute_cap
-          match[rg_count].str()                  // uuid
-      };
+      GpuInfo gpuInfo = {match[1].str(),              // id
+                         match[2].str(),              // vram_total
+                         match[3].str(),              // vram_free
+                         match[4].str(),              // name
+                         GetGpuArch(match[4].str()),  // arch
+                         driver_version,              // driver_version
+                         cuda_version,                // cuda_driver_version
+                         need_fallback ? "0" : match[5].str(),  // compute_cap
+                         match[rg_count].str(),                 // uuid
+                         "NVIDIA"};
       gpuInfoList.push_back(gpuInfo);
       search_start = match.suffix().first;
     }

diff --git a/engine/utils/system_info_utils.h b/engine/utils/system_info_utils.h
@@ -121,6 +121,7 @@ struct GpuInfo {
   std::optional<std::string> cuda_driver_version;
   std::optional<std::string> compute_cap;
   std::string uuid;
+  std::string vendor;
 };
 
 std::vector<GpuInfo> GetGpuInfoListVulkan();