Surface Laptop Ultra：微软首搭 NVIDIA Grace，Arm 版 Windows 终于有了旗舰硬件

# check_cuda_env.ps1 — 在 Arm Windows 上检测 CUDA 工具链
# 运行方式：powershell -ExecutionPolicy Bypass -File check_cuda_env.ps1

Write-Host "=== CUDA Environment Check ===" -ForegroundColor Cyan

# 1. 检查 NVIDIA 驱动
$driver = Get-WmiObject Win32_VideoController | Where-Object { $_.Name -match "NVIDIA" }
if ($driver) {
    Write-Host "[OK] NVIDIA GPU: $($driver.Name)" -ForegroundColor Green
    Write-Host "     Driver Version: $($driver.DriverVersion)" -ForegroundColor Green
} else {
    Write-Host "[FAIL] No NVIDIA GPU detected" -ForegroundColor Red
    exit 1
}

# 2. 检查 nvcc
$nvccPath = Get-Command nvcc -ErrorAction SilentlyContinue
if ($nvccPath) {
    $version = nvcc --version | Select-String "release"
    Write-Host "[OK] nvcc found: $version" -ForegroundColor Green
} else {
    Write-Host "[WARN] nvcc not in PATH — install CUDA Toolkit for Arm64" -ForegroundColor Yellow
    Write-Host "       Download: https://developer.nvidia.com/cuda-toolkit-arm64" -ForegroundColor Yellow
}

# 3. 检查 Python + PyTorch (Arm64)
$python = Get-Command python -ErrorAction SilentlyContinue
if ($python) {
    Write-Host "[OK] Python: $(python --version)" -ForegroundColor Green
    $torchCheck = python -c "import torch; print(torch.cuda.is_available())" 2>$null
    if ($torchCheck -eq "True") {
        Write-Host "[OK] PyTorch CUDA available" -ForegroundColor Green
        python -c "import torch; print(f'GPU: {torch.cuda.get_device_name(0)}')"
    } else {
        Write-Host "[WARN] PyTorch CUDA not available — need Arm64 + CUDA build" -ForegroundColor Yellow
    }
} else {
    Write-Host "[WARN] Python not found" -ForegroundColor Yellow
}

# 4. 检查 WSL2 (可选，用于 Linux 开发流程)
$wsl = Get-Command wsl -ErrorAction SilentlyContinue
if ($wsl) {
    Write-Host "[OK] WSL2 available" -ForegroundColor Green
} else {
    Write-Host "[INFO] WSL2 not installed — optional for Linux-based workflows" -ForegroundColor Gray
}

Write-Host ""
Write-Host "=== Quick GPU Compute Test ===" -ForegroundColor Cyan

# 最小 PyTorch GPU 测试
if ($python -and $torchCheck -eq "True") {
    python -c @"
import torch
import time

device = torch.device('cuda')
# 矩阵乘法基准：4096x4096
a = torch.randn(4096, 4096, device=device)
b = torch.randn(4096, 4096, device=device)

# Warmup
for _ in range(3):
    c = a @ b
torch.cuda.synchronize()

start = time.time()
for _ in range(20):
    c = a @ b
torch.cuda.synchronize()
elapsed = time.time() - start

print(f'4096x4096 matmul: {elapsed/20*1000:.2f} ms per op')
print(f'Estimated throughput: {2*4096**3/elapsed/1e12:.2f} TFLOPS (FP32)')
"@
} else {
    Write-Host "Skipping GPU compute test (PyTorch CUDA not ready)" -ForegroundColor Yellow
}

# 安装 PyTorch Arm64 + CUDA（版本号随发布更新，请查阅 pytorch.org）
pip install torch --index-url https://download.pytorch.org/whl/cu126_arm64

# 在 WSL2 (Ubuntu) 内安装
wsl -d Ubuntu
pip install torch --index-url https://download.pytorch.org/whl/cu126
python -c "import torch; print(torch.cuda.is_available())"
# WSL2 内 CUDA 通过 GPU-PV 虚拟化直通，性能损失约 5-10%