From bb62eded7747e12ca083b8214acc9ebaa47559ae Mon Sep 17 00:00:00 2001
From: Mykola Grymalyuk <khronokernel@icloud.com>
Date: Tue, 12 Oct 2021 21:42:33 -0600
Subject: [PATCH] Add Curie and Comet Lake IDs

---
 data/pci_data.py          | 18 ++++++++++++++++++
 resources/device_probe.py |  8 +++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/data/pci_data.py b/data/pci_data.py
index 736cd4b76..67cc3c972 100644
--- a/data/pci_data.py
+++ b/data/pci_data.py
@@ -3,6 +3,18 @@ class nvidia_ids:
     # Courteous of envytools as well as Macrumors:
     # https://envytools.readthedocs.io/en/latest/hw/pciid.html
     # https://forums.macrumors.com/threads/2011-imac-graphics-card-upgrade.1596614/
+    curie_ids = [
+        0x0040,
+        0x00f0,
+        0x0220,
+        0x0140,
+        0x0160,
+        0x0090,
+        0x01d0,
+        0x0390,
+        0x0290,
+    ]
+
     tesla_ids = [
         # G80
         0x0190,  # G80 [GeForce 8800 GTS / 8800 GTX]
@@ -843,6 +855,12 @@ class intel_ids:
         0x3E98,
     ]
 
+    comet_lake_ids = [
+        0x9BC8,
+        0x9BC5,
+        0x9BC4,
+    ]
+
     ice_lake_ids = [
         # AppleIntelICLLPGraphicsFramebuffer IDs
         0xFF05,
diff --git a/resources/device_probe.py b/resources/device_probe.py
index e3d832470..111424c39 100644
--- a/resources/device_probe.py
+++ b/resources/device_probe.py
@@ -160,6 +160,7 @@ class NVIDIA(GPU):
 
     class Archs(enum.Enum):
         # pylint: disable=invalid-name
+        Curie = "Curie"
         Fermi = "Fermi"
         Tesla = "Tesla"
         Kepler = "Kepler"
@@ -169,7 +170,9 @@ class NVIDIA(GPU):
 
     def detect_arch(self):
         # G80/G80GL
-        if self.device_id in pci_data.nvidia_ids.tesla_ids:
+        if self.device_id in pci_data.nvidia_ids.curie_ids:
+            self.arch = NVIDIA.Archs.Curie
+        elif self.device_id in pci_data.nvidia_ids.tesla_ids:
             self.arch = NVIDIA.Archs.Tesla
         elif self.device_id in pci_data.nvidia_ids.fermi_ids:
             self.arch = NVIDIA.Archs.Fermi
@@ -237,6 +240,7 @@ class Intel(GPU):
         Skylake = "Skylake"
         Kaby_Lake = "Kaby Lake"
         Coffee_Lake = "Coffee Lake"
+        Comet_Lake = "Comet Lake"
         Ice_Lake = "Ice Lake"
         Unknown = "Unknown"
 
@@ -263,6 +267,8 @@ class Intel(GPU):
             self.arch = Intel.Archs.Kaby_Lake
         elif self.device_id in pci_data.intel_ids.coffee_lake_ids:
             self.arch = Intel.Archs.Coffee_Lake
+        elif self.device_id in pci_data.intel_ids.comet_lake_ids:
+            self.arch = Intel.Archs.Comet_Lake
         elif self.device_id in pci_data.intel_ids.ice_lake_ids:
             self.arch = Intel.Archs.Ice_Lake
         else: