compute_caps.py

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "html5lib",
#     "httpx",
# ]
# ///

import lzma

import html5lib
import httpx

NVIDIA_WEB_PAGES = [
    "https://developer.nvidia.com/cuda-gpus",
    "https://developer.nvidia.com/cuda-legacy-gpus",
]
PCI_IDS_XZ = "https://pci-ids.ucw.cz/v2.2/pci.ids.xz"


class ComputeCapsComputer:
    # GPU marketing name to compute capability, e.g.,
    # {"H100": "9.0"}
    gpu_to_cap: dict[str, str]
    # Architecture name to all marketing names using that architecture, e.g.,
    # {"GH100": {"H100", "H200 SXM 141GB", ...}}.
    buckets: dict[str, str]
    # PCI product ID to (possible) architecture name, e.g.,
    # {"2335": "GH100"}
    product_to_bucket: dict[str, str]
    # Architecture name to compute capability, e.g.,
    # {"GH100": "9.0"}
    bucket_to_cap: dict[str, str]
    # PCI product ID to compute capability, e.g.,
    # {"2335": "9.0"}
    gpu_product_to_cap: dict[str, str]

    def __init__(self):
        self.gpu_to_cap = {}
        self.buckets = {}
        self.product_to_bucket = {}

    # Populate gpu_to_cap. Each call is additive.
    def parse_nvidia_webpage(self, url: str, content: bytes) -> None:
        h = html5lib.parse(content, namespaceHTMLElements=False)
        for table in h.findall(".//table"):
            rows = table.findall(".//tr")
            if next(rows[0].itertext()) == "Compute Capability":
                break
        else:
            raise RuntimeError(f"Unable to find table at {url}, did the web page change?")
        for row in rows:
            it = row.itertext()
            cap = next(it)
            if cap.startswith(("1.", "2.", "3.")):
                # these are really old
                continue
            for gpu in it:
                gpu = gpu.removeprefix("NVIDIA ")
                if (oldcap := self.gpu_to_cap.get(gpu, cap)) != cap:
                    raise RuntimeError(f"GPU {gpu} is both {oldcap} and {cap}??")
                self.gpu_to_cap[gpu] = cap

    # Populate buckets and product_to_bucket. Each call is additive (if
    # you somehow have multiple PCI ID files)
    def parse_pci_ids(self, pci_ids: bytes) -> None:
        nvidia = False
        for line in pci_ids.decode().splitlines():
            if not line or line[0] == "#":
                continue
            elif line[0] == "\t":
                if not nvidia:
                    continue
                if line[1] == "\t":
                    continue  # subvendor ID, we don't care
                product_id, product_name = line.split(maxsplit=1)
                if "[" in product_name:
                    bucket_name, specific = product_name.split("[")
                    bucket_name = bucket_name.split(" / ", maxsplit=2)[0].rstrip()
                    specific = specific.removesuffix("]")
                else:
                    bucket_name, specific = product_name, product_name
                bucket = self.buckets.setdefault(bucket_name, {bucket_name})
                bucket.add(specific)
                self.product_to_bucket[product_id] = bucket_name
            else:
                vendor_id, vendor_name = line.split(maxsplit=1)
                if vendor_name != "NVIDIA Corporation":
                    nvidia = False
                    continue
                nvidia = True

    # Populate bucket_to_cap and gpu_product_to_cap from the loaded
    # data. Each call resets these variables.
    def reconcile(self) -> None:
        self.bucket_to_cap = {}
        self.gpu_product_to_cap = {}
        for name, bucket in self.buckets.items():
            caps = {self.gpu_to_cap.get(gpu) for gpu in bucket}
            caps.discard(None)
            if len(caps) == 0:
                continue  # presumably not a GPU
            elif len(caps) > 1:
                raise RuntimeError(f"Multiple compute caps for {bucket}: {caps}")
            self.bucket_to_cap[name] = caps.pop()
        for product, bucket in self.product_to_bucket.items():
            cap = self.bucket_to_cap.get(bucket)
            if cap is None and bucket.endswith("M"):
                # Try the non-mobile GPU's classification.
                bucket = bucket[:-1]
                cap = self.bucket_to_cap.get(cap)
            if cap is None and bucket.endswith("GL"):
                # These seem to be the same except in one case,
                # GM204/GM204GL.
                bucket = bucket[:-2]
                cap = self.bucket_to_cap.get(cap)
            if cap is not None:
                self.gpu_product_to_cap[product] = cap


def main():
    c = ComputeCapsComputer()

    for page in NVIDIA_WEB_PAGES:
        r = httpx.get(page).raise_for_status()
        c.parse_nvidia_webpage(page, r.content)

    r = httpx.get(PCI_IDS_XZ).raise_for_status()
    pci_ids = lzma.decompress(r.content)
    c.parse_pci_ids(pci_ids)

    c.reconcile()

    for product in c.product_to_bucket:
        print(f"0x{product} => {c.gpu_product_to_cap.get(product)}")

#    print("fn product_to_cap(product_id: u16) -> &'static str {")
#    print("    match product_id {")
#    #for product, cap in c.gpu_product_to_cap.items():
#        print(f"        0x{product} => Some(\"{cap}\"),")
#    print("        _ => None,")
#    print("    }")
#    print("}")



if __name__ == "__main__":
    main()