rocm-bootstrap 0.1.0__py3-none-any.whl

rocm_bootstrap/__init__.py

@@ -0,0 +1,61 @@
+"""rocm-bootstrap: GPU detection and target groupings for the ROCm Python ecosystem.
+
+This package provides:
+    - Canonical GFX target hierarchy (family/sub-family/target)
+    - Pure-Python GPU detection via Linux sysfs
+    - Python dist-safe and module-safe naming APIs
+    - WheelNext variant plugin for AMD GPU detection
+"""
+
+from rocm_bootstrap.detect import DetectedGpu, detect_gfx_targets, detect_gpus
+from rocm_bootstrap.naming import (
+    PackageNames,
+    bundle_names,
+    device_dist_name,
+    device_module_name,
+    is_valid_dist_name,
+    is_valid_module_name,
+)
+from rocm_bootstrap.targets import (
+    ALL_FAMILIES,
+    ALL_SUB_FAMILIES,
+    ALL_TARGETS,
+    GfxTarget,
+    PackagingLevel,
+    TargetBundle,
+    XnackMode,
+    all_bundles,
+    bundle_for_target,
+    lookup_bundle,
+    lookup_target,
+    packaging_chain,
+    parse_gfx_target_version,
+)
+
+__all__ = [
+    # targets
+    "GfxTarget",
+    "TargetBundle",
+    "PackagingLevel",
+    "XnackMode",
+    "ALL_FAMILIES",
+    "ALL_SUB_FAMILIES",
+    "ALL_TARGETS",
+    "all_bundles",
+    "bundle_for_target",
+    "lookup_bundle",
+    "lookup_target",
+    "packaging_chain",
+    "parse_gfx_target_version",
+    # naming
+    "PackageNames",
+    "bundle_names",
+    "device_dist_name",
+    "device_module_name",
+    "is_valid_dist_name",
+    "is_valid_module_name",
+    # detect
+    "DetectedGpu",
+    "detect_gpus",
+    "detect_gfx_targets",
+]

rocm_bootstrap/_platform.py

@@ -0,0 +1,121 @@
+"""Fakeable I/O layer for GPU detection.
+
+ALL system I/O that detect.py needs lives here. Tests monkeypatch these
+functions to inject fake sysfs content and clinfo output, so detect.py
+never reads files or runs subprocesses directly.
+"""
+
+import os
+import subprocess
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# KFD topology
+# ---------------------------------------------------------------------------
+
+KFD_TOPOLOGY = Path("/sys/class/kfd/kfd/topology/nodes")
+
+
+def list_kfd_nodes() -> list[Path]:
+    """List /sys/class/kfd/kfd/topology/nodes/* directories.
+
+    Returns directories sorted by name (node number). Returns an empty
+    list if the KFD topology path does not exist.
+    """
+    if not KFD_TOPOLOGY.is_dir():
+        return []
+    return sorted(p for p in KFD_TOPOLOGY.iterdir() if p.is_dir() and p.name.isdigit())
+
+
+def read_kfd_properties(node_path: Path) -> str:
+    """Read raw text of a KFD node's ``properties`` file.
+
+    Args:
+        node_path: Path to a KFD topology node directory
+            (e.g., ``/sys/class/kfd/kfd/topology/nodes/1``).
+
+    Returns:
+        Raw file content as a string.
+
+    Raises:
+        FileNotFoundError: If the properties file does not exist.
+    """
+    return (node_path / "properties").read_text()
+
+
+# ---------------------------------------------------------------------------
+# DRM / ip_discovery
+# ---------------------------------------------------------------------------
+
+DRM_CLASS = Path("/sys/class/drm")
+
+
+def list_drm_cards() -> list[Path]:
+    """List /sys/class/drm/card* directories (numbered cards only).
+
+    Returns directories sorted by name. Returns an empty list if the
+    DRM class path does not exist.
+    """
+    if not DRM_CLASS.is_dir():
+        return []
+    return sorted(
+        p
+        for p in DRM_CLASS.iterdir()
+        if p.is_dir() and p.name.startswith("card") and p.name[4:].isdigit()
+    )
+
+
+def read_ip_discovery_version(card_path: Path) -> tuple[int, int, int] | None:
+    """Read GC IP version from ip_discovery sysfs.
+
+    Looks for ``<card_path>/device/ip_discovery/die/0/GC/0/{major,minor,revision}``.
+
+    Args:
+        card_path: Path to a DRM card directory
+            (e.g., ``/sys/class/drm/card1``).
+
+    Returns:
+        ``(major, minor, revision)`` tuple, or ``None`` if the
+        ip_discovery path does not exist.
+    """
+    gc_path = card_path / "device" / "ip_discovery" / "die" / "0" / "GC" / "0"
+    try:
+        major = int((gc_path / "major").read_text().strip())
+        minor = int((gc_path / "minor").read_text().strip())
+        revision = int((gc_path / "revision").read_text().strip())
+    except (FileNotFoundError, ValueError):
+        return None
+    return (major, minor, revision)
+
+
+# ---------------------------------------------------------------------------
+# Windows: clinfo subprocess
+# ---------------------------------------------------------------------------
+
+
+def run_clinfo() -> str:
+    """Run ``clinfo`` and return its stdout.
+
+    Used on Windows where sysfs is not available. Returns empty string
+    if clinfo is not found or fails.
+    """
+    try:
+        result = subprocess.run(
+            ["clinfo"],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        return result.stdout
+    except FileNotFoundError:
+        return ""
+
+
+# ---------------------------------------------------------------------------
+# Environment variables
+# ---------------------------------------------------------------------------
+
+
+def get_env(key: str) -> str | None:
+    """Read an environment variable. Patchable for tests."""
+    return os.environ.get(key)

rocm_bootstrap/detect.py

@@ -0,0 +1,307 @@
+"""GPU detection orchestration.
+
+Detects AMD GPUs on the current system by reading kernel sysfs interfaces.
+All I/O is delegated to the :mod:`rocm_bootstrap._platform` module so that
+tests can monkeypatch it with fake sysfs content.
+
+Detection chain:
+    1. ``ROCM_BOOTSTRAP_DISABLE_DETECTION`` → return ``[]``
+    2. ``ROCM_BOOTSTRAP_FORCE_GFX_ARCH`` → parse forced targets
+    3. KFD topology (``/sys/class/kfd/kfd/topology/nodes/*/properties``)
+    4. ip_discovery (``/sys/class/drm/card*/device/ip_discovery/die/0/GC/0/``)
+    5. Return empty list if all fail
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from dataclasses import dataclass
+
+from rocm_bootstrap import _platform
+from rocm_bootstrap.targets import (
+    GfxTarget,
+    lookup_target,
+    packaging_chain,
+    parse_gfx_target_version,
+)
+
+
+@dataclass(frozen=True)
+class DetectedGpu:
+    """A GPU detected on the system.
+
+    Attributes:
+        target: The :class:`~rocm_bootstrap.targets.GfxTarget` for this GPU.
+        node_id: KFD topology node index, or DRM card number.
+        gpu_id: KFD ``gpu_id`` property. 0 for CPU-only nodes.
+        pci_id: PCI device ID string if available (e.g., ``"0x7550"``).
+    """
+
+    target: GfxTarget
+    node_id: int
+    gpu_id: int = 0
+    pci_id: str | None = None
+
+
+def detect_gpus() -> list[DetectedGpu]:
+    """Detect AMD GPUs on the current system.
+
+    Returns a list of :class:`DetectedGpu` instances, one per GPU found.
+    Returns an empty list if detection is disabled, no GPUs are found,
+    or the system lacks the required sysfs interfaces.
+
+    Environment variables:
+        ``ROCM_BOOTSTRAP_DISABLE_DETECTION``:
+            Set to ``1`` to skip detection entirely.
+        ``ROCM_BOOTSTRAP_FORCE_GFX_ARCH``:
+            Comma-separated list of GFX target names to return instead
+            of detecting (e.g., ``"gfx942,gfx942"`` for two MI300X GPUs).
+    """
+    # 1. Check disable flag
+    if _platform.get_env("ROCM_BOOTSTRAP_DISABLE_DETECTION") == "1":
+        return []
+
+    # 2. Check forced arch override
+    forced = _platform.get_env("ROCM_BOOTSTRAP_FORCE_GFX_ARCH")
+    if forced:
+        return _parse_forced_targets(forced)
+
+    # 3. Try KFD topology
+    gpus = _detect_via_kfd_topology()
+    if gpus:
+        return gpus
+
+    # 4. Try ip_discovery fallback
+    gpus = _detect_via_ip_discovery()
+    if gpus:
+        return gpus
+
+    # 5. Nothing found
+    return []
+
+
+def detect_gfx_targets() -> list[GfxTarget]:
+    """Convenience: detect GPUs and return deduplicated GfxTarget list.
+
+    Returns unique targets in detection order (first occurrence kept).
+    """
+    seen: set[str] = set()
+    targets: list[GfxTarget] = []
+    for gpu in detect_gpus():
+        if gpu.target.name not in seen:
+            seen.add(gpu.target.name)
+            targets.append(gpu.target)
+    return targets
+
+
+# ---------------------------------------------------------------------------
+# Internal detection methods
+# ---------------------------------------------------------------------------
+
+
+def _parse_forced_targets(forced: str) -> list[DetectedGpu]:
+    """Parse ROCM_BOOTSTRAP_FORCE_GFX_ARCH into DetectedGpu list."""
+    gpus: list[DetectedGpu] = []
+    for i, name in enumerate(forced.split(",")):
+        name = name.strip()
+        if not name:
+            continue
+        target = lookup_target(name)
+        gpus.append(DetectedGpu(target=target, node_id=i))
+    return gpus
+
+
+def _parse_kfd_properties(text: str) -> dict[str, int]:
+    """Parse KFD properties file into a dict of name → int value.
+
+    Lines are formatted as ``key value`` pairs, one per line. Non-integer
+    values are skipped.
+    """
+    props: dict[str, int] = {}
+    for line in text.splitlines():
+        parts = line.split()
+        if len(parts) == 2:
+            try:
+                props[parts[0]] = int(parts[1])
+            except ValueError:
+                continue
+    return props
+
+
+def _detect_via_kfd_topology() -> list[DetectedGpu]:
+    """Detect GPUs via KFD topology sysfs nodes."""
+    nodes = _platform.list_kfd_nodes()
+    if not nodes:
+        return []
+
+    gpus: list[DetectedGpu] = []
+    for node_path in nodes:
+        try:
+            text = _platform.read_kfd_properties(node_path)
+        except FileNotFoundError:
+            continue
+
+        props = _parse_kfd_properties(text)
+
+        # Skip CPU-only nodes (simd_count == 0 indicates no GPU CUs)
+        simd_count = props.get("simd_count", 0)
+        if simd_count == 0:
+            continue
+
+        gtv = props.get("gfx_target_version", 0)
+        if gtv == 0:
+            continue
+
+        target = parse_gfx_target_version(gtv)
+
+        node_id = int(node_path.name)
+        gpu_id = props.get("gpu_id", 0)
+        device_id = props.get("device_id")
+        pci_id = f"0x{device_id:x}" if device_id is not None else None
+
+        gpus.append(
+            DetectedGpu(
+                target=target,
+                node_id=node_id,
+                gpu_id=gpu_id,
+                pci_id=pci_id,
+            )
+        )
+
+    return gpus
+
+
+def _detect_via_ip_discovery() -> list[DetectedGpu]:
+    """Detect GPUs via DRM ip_discovery sysfs (fallback)."""
+    cards = _platform.list_drm_cards()
+    if not cards:
+        return []
+
+    gpus: list[DetectedGpu] = []
+    for card_path in cards:
+        version = _platform.read_ip_discovery_version(card_path)
+        if version is None:
+            continue
+
+        major, minor, revision = version
+        gtv = major * 10000 + minor * 100 + revision
+        target = parse_gfx_target_version(gtv)
+
+        card_num = int(card_path.name.removeprefix("card"))
+        gpus.append(
+            DetectedGpu(
+                target=target,
+                node_id=card_num,
+            )
+        )
+
+    return gpus
+
+
+# ---------------------------------------------------------------------------
+# CLI entry point
+# ---------------------------------------------------------------------------
+
+
+def main(argv: list[str] | None = None) -> None:
+    """CLI entry point for ``rocm-bootstrap-detect``."""
+    parser = argparse.ArgumentParser(
+        prog="rocm-bootstrap-detect",
+        description="Detect AMD GPUs and print target information.",
+    )
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument(
+        "--unique",
+        "-u",
+        action="store_const",
+        dest="mode",
+        const="unique",
+        help="Print unique target names, one per line.",
+    )
+    group.add_argument(
+        "--verbose",
+        "-v",
+        action="store_const",
+        dest="mode",
+        const="verbose",
+        help="Human-readable output with hierarchy and generic ISA details.",
+    )
+    group.add_argument(
+        "--hierarchy",
+        action="store_const",
+        dest="mode",
+        const="hierarchy",
+        help="Print unique bundle hierarchies (target sub-family family).",
+    )
+    parser.set_defaults(mode="unique")
+    args = parser.parse_args(argv)
+
+    gpus = detect_gpus()
+
+    if args.mode == "verbose":
+        _print_verbose(gpus)
+    elif args.mode == "hierarchy":
+        _print_hierarchy(gpus)
+    else:
+        _print_unique(gpus)
+
+
+def _print_unique(gpus: list[DetectedGpu]) -> None:
+    """Machine-consumable: one unique target name per line."""
+    for target in _unique_targets(gpus):
+        print(target.name)
+
+
+def _print_hierarchy(gpus: list[DetectedGpu]) -> None:
+    """Unique packaging chains, space-separated: target sub_family family."""
+    seen: set[str] = set()
+    for target in _unique_targets(gpus):
+        if target.name in seen:
+            continue
+        seen.add(target.name)
+        chain = packaging_chain(target)
+        print(" ".join(b.key for b in chain))
+
+
+def _print_verbose(gpus: list[DetectedGpu]) -> None:
+    """Human-readable output with full details."""
+    if not gpus:
+        print("No AMD GPUs detected.")
+        return
+
+    print(f"Detected {len(gpus)} AMD GPU(s):\n")
+    for gpu in gpus:
+        t = gpu.target
+        parts = [f"  Node {gpu.node_id}: {t.name}"]
+        parts.append(f"major={t.major} minor={t.minor} stepping={t.stepping}")
+        if gpu.pci_id:
+            parts.append(f"PCI={gpu.pci_id}")
+        if gpu.gpu_id:
+            parts.append(f"gpu_id={gpu.gpu_id}")
+        print("  ".join(parts))
+
+        chain = packaging_chain(t)
+        # chain is (target_bundle, sub_family_bundle, family_bundle)
+        for bundle in chain[1:]:  # skip target-level (already printed)
+            label = bundle.level.value.replace("_", "-")
+            line = f"    {label}: {bundle.key} ({bundle.display_name})"
+            if bundle.llvm_generic:
+                line += f"  generic: {bundle.llvm_generic}"
+            print(line)
+        print()
+
+
+def _unique_targets(gpus: list[DetectedGpu]) -> list[GfxTarget]:
+    """Deduplicate targets preserving detection order."""
+    seen: set[str] = set()
+    targets: list[GfxTarget] = []
+    for gpu in gpus:
+        if gpu.target.name not in seen:
+            seen.add(gpu.target.name)
+            targets.append(gpu.target)
+    return targets
+
+
+if __name__ == "__main__":
+    main()

rocm_bootstrap/naming.py

@@ -0,0 +1,131 @@
+"""Python dist-safe and module-safe name generation for GFX targets.
+
+This module provides deterministic, validated name generation for all
+levels of the GFX packaging hierarchy. Every generated name is guaranteed
+to be valid for its intended context (pip/wheel dist names or Python
+module identifiers).
+
+Naming convention:
+    - Bundle ``key`` uses underscores → directly a valid Python identifier.
+    - ``dist_name`` replaces ``_`` with ``-`` → valid wheel/package name.
+    - ``module_name`` IS the ``key`` → valid Python identifier.
+    - All names start with ``gfx`` → no leading digit issues.
+"""
+
+import re
+from dataclasses import dataclass
+
+from rocm_bootstrap.targets import TargetBundle
+
+# PEP 625 / PyPA: distribution names are lowercase, alphanumeric + hyphens +
+# underscores + periods. Must start with alphanumeric.
+_DIST_NAME_RE = re.compile(r"^[a-z0-9]([a-z0-9._-]*[a-z0-9])?$")
+
+
+@dataclass(frozen=True)
+class PackageNames:
+    """Dist-safe and module-safe names for a packaging level.
+
+    Attributes:
+        dist_name: Name safe for Python distribution/wheel usage.
+            Lowercase, uses hyphens as separators.
+            E.g., ``"gfx11-5"``, ``"gfx942"``.
+        module_name: Name safe for Python identifiers/modules.
+            Uses underscores, no hyphens, no leading digits.
+            E.g., ``"gfx11_5"``, ``"gfx942"``.
+    """
+
+    dist_name: str
+    module_name: str
+
+
+def bundle_names(bundle: TargetBundle) -> PackageNames:
+    """Get Python-safe package names for a bundle.
+
+    The bundle's ``key`` is the ``module_name`` (already uses underscores).
+    The ``dist_name`` is derived by replacing ``_`` with ``-``.
+
+    Args:
+        bundle: Any :class:`~rocm_bootstrap.targets.TargetBundle`.
+
+    Returns:
+        :class:`PackageNames` with validated dist and module names.
+    """
+    module_name = bundle.key
+    dist_name = module_name.replace("_", "-")
+    return PackageNames(dist_name=dist_name, module_name=module_name)
+
+
+def device_dist_name(prefix: str, bundle: TargetBundle) -> str:
+    """Generate a device package distribution name.
+
+    Combines a dist-name prefix with the bundle's dist name.
+
+    Examples::
+
+        device_dist_name("rocm-sdk-device", bundle_gfx11_5)
+        # -> "rocm-sdk-device-gfx11-5"
+
+        device_dist_name("amd-torch-device", bundle_gfx942)
+        # -> "amd-torch-device-gfx942"
+
+    Args:
+        prefix: Distribution name prefix (e.g., ``"rocm-sdk-device"``).
+        bundle: The :class:`~rocm_bootstrap.targets.TargetBundle`.
+
+    Returns:
+        Combined dist name string.
+    """
+    names = bundle_names(bundle)
+    return f"{prefix}-{names.dist_name}"
+
+
+def device_module_name(prefix: str, bundle: TargetBundle) -> str:
+    """Generate a device package module name.
+
+    Combines a module-name prefix with the bundle's module name.
+
+    Examples::
+
+        device_module_name("rocm_sdk_device", bundle_gfx11_5)
+        # -> "rocm_sdk_device_gfx11_5"
+
+    Args:
+        prefix: Module name prefix (e.g., ``"rocm_sdk_device"``).
+        bundle: The :class:`~rocm_bootstrap.targets.TargetBundle`.
+
+    Returns:
+        Combined module name string.
+    """
+    names = bundle_names(bundle)
+    return f"{prefix}_{names.module_name}"
+
+
+def is_valid_dist_name(name: str) -> bool:
+    """Check if a string is valid as a Python distribution name.
+
+    Per PEP 625 / PyPA naming specification: lowercase alphanumeric,
+    hyphens, underscores, periods. Must start and end with alphanumeric.
+
+    Args:
+        name: String to validate.
+
+    Returns:
+        ``True`` if valid.
+    """
+    return bool(_DIST_NAME_RE.match(name))
+
+
+def is_valid_module_name(name: str) -> bool:
+    """Check if a string is valid as a Python module/identifier name.
+
+    Must be a valid Python identifier: starts with letter or underscore,
+    contains only letters, digits, underscores. Must not be a keyword.
+
+    Args:
+        name: String to validate.
+
+    Returns:
+        ``True`` if valid.
+    """
+    return name.isidentifier()