scripts/check_gpu_compatibility.py

"""Report Anima MixBit runtime-mode eligibility for the local NVIDIA GPU."""

from __future__ import annotations

import argparse
import json
import sys
from typing import Any


STANDARD_VALIDATED_NAMES = (
    "nvidia geforce rtx 3090",
    "nvidia geforce rtx 5080 laptop gpu",
)
LOW_VRAM_VALIDATED_NAMES = ("nvidia geforce rtx 5080 laptop gpu",)


def native_nvfp4_eligible(capability: tuple[int, int]) -> bool:
    """Match the native scaled_mm capability gate used by the runtime."""

    major, minor = capability
    return (major, minor) == (8, 9) or major >= 9


def classify(name: str, capability: tuple[int, int]) -> dict[str, Any]:
    normalized = name.strip().lower()
    native_eligible = native_nvfp4_eligible(capability)
    standard_validated = any(item in normalized for item in STANDARD_VALIDATED_NAMES)
    low_vram_validated = any(item in normalized for item in LOW_VRAM_VALIDATED_NAMES)

    if low_vram_validated:
        low_vram_status = "validated"
    elif native_eligible:
        low_vram_status = "capability_eligible_not_package_validated"
    else:
        low_vram_status = "unsupported_by_native_nvfp4_gate"

    return {
        "name": name,
        "compute_capability": f"{capability[0]}.{capability[1]}",
        "standard_mode": "validated" if standard_validated else "not_package_validated",
        "low_vram_mode": low_vram_status,
        "native_nvfp4_eligible": native_eligible,
        "recommended_mode": (
            "low_vram_or_standard" if native_eligible else "standard"
        ),
    }


def self_test() -> int:
    expectations = {
        (8, 0): False,
        (8, 6): False,
        (8, 9): True,
        (9, 0): True,
        (12, 0): True,
    }
    for capability, expected in expectations.items():
        actual = native_nvfp4_eligible(capability)
        if actual != expected:
            raise AssertionError(f"Capability {capability}: {actual} != {expected}")
    print("PASS: GPU compatibility rules are internally consistent.")
    return 0


def inspect_device(index: int) -> dict[str, Any]:
    try:
        import torch
    except ImportError as error:
        return {
            "ready": False,
            "error": f"PyTorch is not installed: {error}",
        }

    if not torch.cuda.is_available():
        return {
            "ready": False,
            "error": "CUDA is not available. This package has no validated CPU runtime.",
        }
    if index < 0 or index >= torch.cuda.device_count():
        return {
            "ready": False,
            "error": f"CUDA device index {index} is out of range.",
        }

    name = torch.cuda.get_device_name(index)
    capability = tuple(int(value) for value in torch.cuda.get_device_capability(index))
    result = classify(name, capability)
    result.update(
        {
            "ready": True,
            "device_index": index,
            "pytorch": torch.__version__,
            "cuda_runtime": torch.version.cuda,
        }
    )
    return result


def print_human(result: dict[str, Any]) -> None:
    if not result.get("ready"):
        print(f"判定不能: {result['error']}")
        return

    standard = {
        "validated": "実機検証済み",
        "not_package_validated": "このパッケージでは未検証",
    }[result["standard_mode"]]
    low_vram = {
        "validated": "実機検証済み",
        "capability_eligible_not_package_validated": "演算条件適合・パッケージ未検証",
        "unsupported_by_native_nvfp4_gate": "非対応（native NVFP4条件外）",
    }[result["low_vram_mode"]]
    recommendation = (
        "VRAMを優先するなら省VRAM、速度と互換性を優先するなら通常"
        if result["native_nvfp4_eligible"]
        else "通常"
    )
    print(f"GPU: {result['name']}")
    print(f"Compute Capability: {result['compute_capability']}")
    print(f"通常モード: {standard}")
    print(f"省VRAMモード: {low_vram}")
    print(f"推奨: {recommendation}")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", type=int, default=0)
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        return self_test()

    result = inspect_device(args.device)
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print_human(result)
    return 0 if result.get("ready") else 2


if __name__ == "__main__":
    sys.exit(main())