feat: Enhanced platform detection (ARM64, RISC-V, device models)

This PR enhances platform detection for better edge device support.

Changes:
- Add ARM64 detection (Raspberry Pi, Apple Silicon, etc.)
- Add RISC-V detection (emerging architecture)
- Add specific device model detection (Broadcom, Apple Silicon)
- Enhanced logging with detailed platform info

Inspired by Q-Lite project for ESP32/Pico integration:
https://github.com/RalphBigBear/q-lite

Benefits:
- Better support for ARM64 edge devices (Pi, etc.)
- Prepares for RISC-V support (emerging)
- Improved debugging with detailed platform info
This commit is contained in:
Ralph 2026-02-12 00:07:36 +08:00
parent 2f09686724
commit 4f2f78a4a5

View File

@ -140,14 +140,66 @@ def get_dist_info():
return False, 0, 0, 1
def autodetect_device_type():
# prefer to use CUDA if available, otherwise use MPS, otherwise fallback on CPU
"""
Enhanced platform detection for ARM64, RISC-V, and edge devices.
Detects:
- CUDA/ROCm/MPS/CPU (original)
- ARM64 (Raspberry Pi, Apple Silicon, etc.)
- RISC-V (emerging architecture)
- Specific device models (Broadcom, Apple Silicon)
Inspired by Q-Lite for edge device integration.
"""
import platform
import torch
# Original detection (CUDA/MPS/CPU)
if torch.cuda.is_available():
device_type = "cuda"
elif torch.backends.mps.is_available():
device_type = "mps"
else:
device_type = "cpu"
print0(f"Autodetected device type: {device_type}")
# Enhanced detection (ARM64/RISC-V)
machine = platform.machine()
system = platform.system()
if machine == "aarch64":
device_type = f"{device_type}:arm64"
# Detect specific ARM device
try:
with open("/proc/cpuinfo", "r") as f:
cpuinfo = f.read()
if "Raspberry Pi" in cpuinfo:
# Extract Pi version
for line in cpuinfo.split("\n"):
if "Hardware" in line:
device_model = line.split(":")[1].strip()
device_type = f"{device_type}:rpi-{device_model}"
break
elif "broadcom" in cpuinfo.lower():
device_type = f"{device_type}:broadcom-arm"
except (FileNotFoundError, PermissionError):
# Non-Linux ARM64 (e.g., macOS ARM64)
if system == "Darwin":
device_type = f"{device_type}:apple-silicon"
else:
device_type = f"{device_type}:generic-arm64"
elif machine == "riscv64":
device_type = f"{device_type}:riscv64"
# Add platform info logging
print0(f"Autodetected platform:")
print0(f" System: {system}")
print0(f" Machine: {machine}")
print0(f" Device: {device_type}")
return device_type
return device_type
def compute_init(device_type="cuda"): # cuda|cpu|mps