From 55a806e677aed6ea108f52ea176a36f6605034f0 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Thu, 1 Jun 2023 16:14:07 -0400
Subject: [PATCH] Add IterGpuInfo and HasNvidiaGpu methods

---
 StabilityMatrix/Helper/HardwareHelper.cs | 76 +++++++++++++++++++++++-
 1 file changed, 73 insertions(+), 3 deletions(-)

diff --git a/StabilityMatrix/Helper/HardwareHelper.cs b/StabilityMatrix/Helper/HardwareHelper.cs
index 4239475b..16eeec2f 100644
--- a/StabilityMatrix/Helper/HardwareHelper.cs
+++ b/StabilityMatrix/Helper/HardwareHelper.cs
@@ -1,14 +1,21 @@
 using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using Microsoft.Win32;
 
 namespace StabilityMatrix.Helper;
 
-public class HardwareHelper
+public static class HardwareHelper
 {
+    private const string GpuRegistryKeyPath =
+        @"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}";
+
     public static ulong GetGpuMemoryBytes()
     {
         var registry = Registry.LocalMachine;
-        var key = registry.OpenSubKey(@"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}\0000", false);
+        var key = registry.OpenSubKey(
+            @"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}\0000", false);
         if (key == null)
         {
             return 0;
@@ -22,7 +29,8 @@ public class HardwareHelper
     public static string GetGpuChipName()
     {
         var registry = Registry.LocalMachine;
-        var key = registry.OpenSubKey(@"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}\0000", false);
+        var key = registry.OpenSubKey(
+            @"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}\0000", false);
         if (key == null)
         {
             return "Unknown";
@@ -31,4 +39,66 @@ public class HardwareHelper
         var gpuName = key.GetValue("HardwareInformation.ChipType");
         return gpuName?.ToString() ?? "Unknown";
     }
+
+    /// <summary>
+    /// Yields GpuInfo for each GPU in the system.
+    /// </summary>
+    public static IEnumerable<GpuInfo> IterGpuInfo()
+    {
+        using var baseKey = Registry.LocalMachine.OpenSubKey(GpuRegistryKeyPath);
+        if (baseKey == null)
+        {
+            yield break;
+        }
+
+        foreach (var subKeyName in baseKey.GetSubKeyNames().Where(k => k.StartsWith("0")))
+        {
+            using var subKey = baseKey.OpenSubKey(subKeyName);
+            if (subKey != null)
+            {
+                yield return new GpuInfo
+                {
+                    Name = subKey.GetValue("DriverDesc")?.ToString(),
+                    MemoryBytes = Convert.ToUInt64(subKey.GetValue("HardwareInformation.qwMemorySize")),
+                };
+            }
+        }
+    }
+    
+    /// <summary>
+    /// Return true if the system has at least one Nvidia GPU.
+    /// </summary>
+    public static bool HasNvidiaGpu()
+    {
+        return IterGpuInfo().Any(gpu => gpu.Name?.ToLowerInvariant().Contains("nvidia") ?? false);
+    }
+}
+
+[SuppressMessage("ReSharper", "MemberCanBePrivate.Global")]
+public static class Size
+{
+    public const ulong KiB = 1024;
+    public const ulong MiB = KiB * 1024;
+    public const ulong GiB = MiB * 1024;
+}
+
+public enum Level
+{
+    Unknown,
+    Low,
+    Medium,
+    High
+}
+
+public record GpuInfo
+{
+    public string? Name { get; init; } = string.Empty;
+    public ulong? MemoryBytes { get; init; }
+    public Level? MemoryLevel => MemoryBytes switch
+    {
+        <= 0 => Level.Unknown,
+        < 4 * Size.GiB => Level.Low,
+        < 8 * Size.GiB => Level.Medium,
+        _ => Level.High
+    };
 }