Refactor GPU identification

Author: cmdr2

tests/test_gpu_db.py

@@ -0,0 +1,183 @@
+"""
+Tests for the gpu_db module to ensure GPU identification patterns work correctly.
+"""
+
+import pytest
+from torchruntime.gpu_db import (
+    get_nvidia_arch,
+    get_amd_gfx_info,
+    get_gpu_type,
+    is_gpu_vendor,
+    NVIDIA,
+    AMD,
+    INTEL,
+)
+
+
+class TestNvidiaArchDetection:
+    """Test NVIDIA GPU architecture detection."""
+
+    def test_kepler_detection(self):
+        device_names = ["NVIDIA GeForce GK104", "GeForce GTX 780 Ti"]
+        assert get_nvidia_arch(device_names) == 3.7
+
+    def test_maxwell_detection(self):
+        device_names = ["NVIDIA GeForce GM107", "GeForce GTX 950"]
+        assert get_nvidia_arch(device_names) == 5.0
+
+    def test_pascal_detection(self):
+        device_names = ["NVIDIA GeForce GP104", "GeForce GTX 1080"]
+        assert get_nvidia_arch(device_names) == 6.0
+
+    def test_volta_detection(self):
+        device_names = ["NVIDIA Tesla GV100"]
+        assert get_nvidia_arch(device_names) == 7.0
+
+    def test_turing_detection(self):
+        device_names = ["NVIDIA GeForce TU116", "GeForce GTX 1660"]
+        assert get_nvidia_arch(device_names) == 7.5
+
+    def test_ampere_detection(self):
+        device_names = ["NVIDIA GeForce GA102", "GeForce RTX 3080"]
+        assert get_nvidia_arch(device_names) == 8.6
+
+    def test_hopper_detection(self):
+        device_names = ["NVIDIA GH100", "H100 PCIe"]
+        assert get_nvidia_arch(device_names) == 9.0
+
+    def test_ada_lovelace_detection(self):
+        device_names = ["NVIDIA GeForce AD102", "GeForce RTX 4090"]
+        assert get_nvidia_arch(device_names) == 8.9
+
+    def test_blackwell_detection(self):
+        device_names = ["NVIDIA GeForce RTX 5090"]
+        assert get_nvidia_arch(device_names) == 12.0
+
+    def test_blackwell_gb_code_detection(self):
+        device_names = ["NVIDIA GB100"]
+        assert get_nvidia_arch(device_names) == 12.0
+
+    def test_unknown_architecture(self):
+        device_names = ["NVIDIA GeForce Unknown"]
+        assert get_nvidia_arch(device_names) == 0
+
+
+class TestAMDGfxInfo:
+    """Test AMD GPU GFX information retrieval."""
+
+    def test_renoir_gfx(self):
+        name, gfx, hsa = get_amd_gfx_info("1636")
+        assert name == "Renoir"
+        assert gfx == "gfx90c"
+        assert hsa == "9.0.12"
+
+    def test_phoenix_gfx(self):
+        name, gfx, hsa = get_amd_gfx_info("164f")
+        assert name == "Phoenix"
+        assert gfx == "gfx1103"
+        assert hsa == "11.0.1"
+
+    def test_strix_gfx(self):
+        name, gfx, hsa = get_amd_gfx_info("150e")
+        assert name == "Strix"
+        assert gfx == "gfx1150"
+        assert hsa == "11.5.0"
+
+    def test_unknown_device_id(self):
+        name, gfx, hsa = get_amd_gfx_info("9999")
+        assert name == ""
+        assert gfx == ""
+        assert hsa == ""
+
+
+class TestGPUTypeDetection:
+    """Test GPU type detection (discrete, integrated, none)."""
+
+    def test_nvidia_discrete_by_brand(self):
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA GeForce RTX 4080") == "DISCRETE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA Quadro P1000") == "DISCRETE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA Tesla V100") == "DISCRETE"
+
+    def test_nvidia_discrete_by_arch_code(self):
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA GP104") == "DISCRETE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA TU116") == "DISCRETE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA GA102") == "DISCRETE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA AD103") == "DISCRETE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA GH100") == "DISCRETE"
+
+    def test_nvidia_excluded_devices(self):
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA Audio Controller") == "NONE"
+        assert get_gpu_type(NVIDIA, "1234", "NVIDIA USB Controller") == "NONE"
+
+    def test_amd_discrete_by_name(self):
+        assert get_gpu_type(AMD, "1234", "AMD Radeon RX 7900 XT") == "DISCRETE"
+        assert get_gpu_type(AMD, "1234", "AMD Navi 31") == "DISCRETE"
+        assert get_gpu_type(AMD, "1234", "AMD Instinct MI300") == "DISCRETE"
+
+    def test_amd_integrated_by_device_id(self):
+        assert get_gpu_type(AMD, "1636", "AMD Renoir") == "INTEGRATED"
+        assert get_gpu_type(AMD, "164f", "AMD Phoenix") == "INTEGRATED"
+        assert get_gpu_type(AMD, "150e", "AMD Strix") == "INTEGRATED"
+
+    def test_amd_excluded_devices(self):
+        assert get_gpu_type(AMD, "1234", "AMD Audio Device") == "NONE"
+        assert get_gpu_type(AMD, "1234", "AMD USB Controller") == "NONE"
+
+    def test_intel_discrete(self):
+        assert get_gpu_type(INTEL, "1234", "Intel Arc A770") == "DISCRETE"
+
+    def test_intel_integrated(self):
+        assert get_gpu_type(INTEL, "1234", "Intel Iris Xe Graphics") == "INTEGRATED"
+        assert get_gpu_type(INTEL, "1234", "Intel HD Graphics 620") == "INTEGRATED"
+        assert get_gpu_type(INTEL, "1234", "Intel UHD Graphics 630") == "INTEGRATED"
+
+    def test_intel_excluded_devices(self):
+        assert get_gpu_type(INTEL, "1234", "Intel Audio Device") == "NONE"
+
+
+class TestGPUVendorCheck:
+    """Test GPU vendor identification."""
+
+    def test_known_vendors(self):
+        assert is_gpu_vendor(NVIDIA) is True
+        assert is_gpu_vendor(AMD) is True
+        assert is_gpu_vendor(INTEL) is True
+
+    def test_unknown_vendor(self):
+        assert is_gpu_vendor("9999") is False
+
+
+class TestPatternConstruction:
+    """Test that patterns are correctly built without duplication."""
+
+    def test_nvidia_discrete_pattern_includes_all_architectures(self):
+        """Verify NVIDIA_DISCRETE_PATTERN includes all architecture codes."""
+        from torchruntime.gpu_db import NVIDIA_DISCRETE_PATTERN
+
+        # Test that all architecture patterns are included
+        arch_codes = ["gk104", "gm107", "gp104", "gv100", "tu116", "ga102", "gh100", "ad102", "gb100"]
+        for code in arch_codes:
+            assert NVIDIA_DISCRETE_PATTERN.search(code), f"Pattern should match {code}"
+
+        # Test brand names
+        brands = ["GeForce", "Quadro", "Tesla", "RTX", "GTX", "Titan"]
+        for brand in brands:
+            assert NVIDIA_DISCRETE_PATTERN.search(brand), f"Pattern should match {brand}"
+
+        # Test Blackwell model numbers
+        models = ["5060", "5070", "5080", "5090"]
+        for model in models:
+            assert NVIDIA_DISCRETE_PATTERN.search(model), f"Pattern should match {model}"
+
+    def test_amd_discrete_pattern_separate(self):
+        """Verify AMD has a separate discrete pattern."""
+        from torchruntime.gpu_db import AMD_DISCRETE_PATTERN
+
+        # Should match AMD product lines
+        assert AMD_DISCRETE_PATTERN.search("Radeon RX 7900")
+        assert AMD_DISCRETE_PATTERN.search("Instinct MI300")
+        assert AMD_DISCRETE_PATTERN.search("Navi 31")
+
+        # Should not match integrated codenames
+        assert not AMD_DISCRETE_PATTERN.search("Phoenix")
+        assert not AMD_DISCRETE_PATTERN.search("Renoir")

torchruntime/configuration.py

@@ -1,7 +1,8 @@
 import os
 
 from .consts import AMD
-from .device_db import get_gpus, GPU_DEVICES
+from .device_db import get_gpus
+from .gpu_db import get_amd_gfx_info
 from .platform_detection import get_torch_platform, os_name
 
 
@@ -110,9 +111,7 @@ def _set_rocm_vars_for_integrated(gpu_infos):
     gpu = gpu_infos[0]
     env = {}
 
-    integrated_amd_gpus = GPU_DEVICES[AMD]["integrated"]
-
-    family_name, gfx_id, hsa_version = integrated_amd_gpus.get(gpu.device_id, ("", "", ""))
+    family_name, gfx_id, hsa_version = get_amd_gfx_info(gpu.device_id)
     env["HSA_OVERRIDE_GFX_VERSION"] = hsa_version
 
     if gfx_id.startswith("gfx8"):

torchruntime/device_db.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 
 from .consts import NVIDIA, AMD, INTEL
+from .gpu_db import is_gpu_vendor, get_gpu_type
 
 DEVICE_DB_FILE = "gpu_pci_ids.db"  # this file will only include AMD, NVIDIA and Discrete Intel GPUs
 
@@ -20,97 +21,9 @@ class GPU:
     is_discrete: bool
 
 
-GPU_DEVICES = {  # if the value is a regex, it'll be applied to the device_name. if the value is a dict, the pci_id will be looked up
-    AMD: {
-        "discrete": re.compile(r"\b(?:radeon|instinct|fire|rage|polaris|aldebaran|navi)\b", re.IGNORECASE),
-        "integrated": {  # pci_id -> (device_name, gfx_name, hsa_override)
-            "15dd": ("Raven Ridge", "gfx902", "9.1.0"),
-            "15d8": ("Picasso", "gfx903", "9.1.0"),
-            # gfx90c - https://github.com/ROCm/clr/commit/fbbae8055fbad748254bec1675a6970ce1dce594
-            "1636": ("Renoir", "gfx90c", "9.0.12"),
-            "164c": ("Lucienne", "gfx90c", "9.0.12"),
-            "1638": ("Cezanne", "gfx90c", "9.0.12"),
-            "15e7": ("Barcelo", "gfx90c", "9.0.12"),
-            "163f": ("VanGogh", "gfx1033", "10.3.3"),
-            "164d": ("Rembrandt", "gfx1035", "10.3.5"),
-            "1681": ("Rembrandt", "gfx1035", "10.3.5"),
-            "164e": ("Raphael", "gfx1036", "10.3.6"),
-            "1506": ("Mendocino", "gfx1037", "10.3.7"),
-            "13c0": ("Granite Ridge", "gfx1030", "10.3.0"),
-            "164f": ("Phoenix", "gfx1103", "11.0.1"),
-            "15bf": ("Phoenix1", "gfx1103", "11.0.3"),
-            "15c8": ("Phoenix2", "gfx1103", "11.0.3"),
-            "1900": ("Phoenix3", "gfx1103", "11.0.4"),
-            "1901": ("Phoenix4", "gfx1103", "11.0.4"),
-            # gfx1150 - https://github.com/ROCm/clr/commit/4bc515aa62368c6189d19e14b3ec18cb6dd4415e
-            "150e": ("Strix", "gfx1150", "11.5.0"),
-            "1586": ("Strix Halo", "gfx1151", "11.5.1"),
-            "1114": ("Krackan", "gfx1151", "11.5.1"),
-        },
-        "exclude": re.compile(r"\b(?:audio|bridge|arden|oberon|stoney|wani|usb|switch)\b", re.IGNORECASE),
-    },
-    INTEL: {
-        "discrete": re.compile(r"\b(?:arc)\b", re.IGNORECASE),
-        "integrated": re.compile(r"\b(?:iris|hd graphics|uhd graphics)\b", re.IGNORECASE),
-        "exclude": re.compile(r"\b(?:audio|bridge)\b", re.IGNORECASE),
-    },
-    NVIDIA: {
-        "discrete": re.compile(
-            r"\b(?:geforce|riva|quadro|tesla|ion|grid|rtx|tu\d{2,}.+t\d{2,}|gk1\d{2}\w*|gm10\d\w*|gp10\d\w*|gv100\w*|tu1\d{2}\w*|ga10\d\w*|ad10\d\w*)\b",
-            re.IGNORECASE,
-        ),
-        "exclude": re.compile(
-            r"\b(?:audio|switch|pci|memory|smbus|ide|co-processor|bridge|usb|sata|controller)\b", re.IGNORECASE
-        ),
-    },
-}
-
-
 os_name = platform.system()
 
 
-def is_gpu_vendor(vendor_id):
-    return vendor_id in GPU_DEVICES
-
-
-def get_gpu_type(vendor_id, device_id, device_name):
-    """
-    Returns the GPU type of the given PCI device.
-
-    Args:
-        vendor_id (str): PCI Vendor ID.
-        device_id (str): PCI Device ID.
-        device_name (str): PCI Device Name.
-
-    Returns:
-        gpu_type (str): "DISCRETE", "INTEGRATED", or "NONE"
-    """
-
-    def matches(pattern):
-        if isinstance(pattern, re.Pattern):
-            return pattern.search(device_name)
-        if isinstance(pattern, dict):
-            return device_id in pattern
-        return False
-
-    vendor_devices = GPU_DEVICES[vendor_id]
-
-    discrete_devices = vendor_devices.get("discrete")
-    integrated_devices = vendor_devices.get("integrated")
-    exclude_devices = vendor_devices.get("exclude")
-
-    if matches(exclude_devices):
-        return "NONE"
-
-    if matches(integrated_devices):  # check integrated first, to avoid matching "Radeon" and classifying it as discrete
-        return "INTEGRATED"
-
-    if matches(discrete_devices):
-        return "DISCRETE"
-
-    return "NONE"
-
-
 def get_windows_output():
     try:
         command = [