chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,93 @@
+"""Automatic detection of the device available on the local machine."""
+
+import os
+import subprocess
+import sys
+from typing import Dict, Optional  # noqa: UP035
+
+import tvm
+from tvm.runtime import Device
+from tvm_ffi import DLDeviceType
+
+from . import logging
+from .style import bold, green, red
+
+FOUND = green("Found")
+NOT_FOUND = red("Not found")
+AUTO_DETECT_DEVICES = ["cuda", "rocm", "metal", "vulkan", "opencl", "cpu"]
+_RESULT_CACHE: Dict[str, bool] = {}  # noqa: UP006
+
+
+logger = logging.getLogger(__name__)
+
+
+def detect_device(device_hint: str) -> Optional[Device]:
+    """Detect locally available device from string hint."""
+    if device_hint == "auto":
+        device = None
+        for device_type in AUTO_DETECT_DEVICES:
+            cur_device = tvm.device(device_type=device_type, index=0)
+            if _device_exists(cur_device):
+                if device is None:
+                    device = cur_device
+        if device is None:
+            logger.info("%s: No available device detected", NOT_FOUND)
+            return None
+        logger.info("Using device: %s", bold(device2str(device)))
+        return device
+    try:
+        device = tvm.device(device_hint)
+    except Exception as err:
+        raise ValueError(f"Invalid device name: {device_hint}") from err
+    if not _device_exists(device):
+        raise ValueError(f"Device is not found on your local environment: {device_hint}")
+    return device
+
+
+def device2str(device: Device) -> str:
+    """Convert a TVM device object to string."""
+    return f"{tvm.runtime.Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}"
+
+
+def _device_exists(device: Device) -> bool:
+    device_type = tvm.runtime.Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
+    device_str = device2str(device)
+    if device_str in _RESULT_CACHE:
+        return _RESULT_CACHE[device_str]
+    cmd = [
+        sys.executable,
+        "-m",
+        "mlc_llm.cli.check_device",
+        device_type,
+    ]
+    prefix = "check_device:"
+    subproc_outputs = [
+        line[len(prefix) :].strip()
+        for line in subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=False,
+            env=os.environ,
+        )
+        .stdout.strip()
+        .splitlines()
+        if line.startswith(prefix)
+    ]
+    if subproc_outputs:
+        if subproc_outputs[0]:
+            for i in subproc_outputs[0].split(","):
+                logger.info("%s device: %s:%s", FOUND, device_type, i)
+                _RESULT_CACHE[f"{device_type}:{i}"] = True
+                if device.dlpack_device_type() == DLDeviceType.kDLCPU:
+                    break
+    else:
+        logger.error(
+            "GPU device detection failed. Please report this issue with the output of command: %s",
+            " ".join(cmd),
+        )
+    if device_str in _RESULT_CACHE:
+        return _RESULT_CACHE[device_str]
+    logger.info("%s device: %s", NOT_FOUND, device_str)
+    _RESULT_CACHE[device_str] = False
+    return False