1. 增加配置持久化
2. 完善细节
This commit is contained in:
sansen
+288
-126
@@ -5,141 +5,231 @@ package gpu
|
||||
import (
|
||||
"log"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"syscall"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
var hideWindow = &syscall.SysProcAttr{HideWindow: true}
|
||||
|
||||
type gpuName struct {
|
||||
Name string
|
||||
Vendor string
|
||||
}
|
||||
|
||||
// --- Method 1: DXGI ---
|
||||
var (
|
||||
dxgi = syscall.NewLazyDLL("dxgi.dll")
|
||||
createDXGIFactory1 = dxgi.NewProc("CreateDXGIFactory1")
|
||||
IID_IDXGIFactory1 = &syscall.GUID{0x770aae78, 0xf26f, 0x4dba, [8]byte{0xa8, 0x29, 0x25, 0x3c, 0x83, 0xd1, 0xb3, 0x87}}
|
||||
)
|
||||
|
||||
type dxgiDesc struct {
|
||||
Description [128]uint16
|
||||
VendorId uint32
|
||||
DeviceId uint32
|
||||
SubSysId uint32
|
||||
Revision uint32
|
||||
DedicatedVideoMemory uint64
|
||||
_ uint64
|
||||
_ uint64
|
||||
_ [8]byte
|
||||
}
|
||||
|
||||
func dxgiGPUList() []gpuName {
|
||||
var factory uintptr
|
||||
r, _, _ := createDXGIFactory1.Call(uintptr(unsafe.Pointer(IID_IDXGIFactory1)), uintptr(unsafe.Pointer(&factory)))
|
||||
if r != 0 {
|
||||
return nil
|
||||
}
|
||||
defer release(factory)
|
||||
|
||||
var gpus []gpuName
|
||||
for i := uint32(0); ; i++ {
|
||||
var adapter uintptr
|
||||
fn := callVtable(factory, 12) // EnumAdapters1
|
||||
r, _, _ = syscall.SyscallN(fn, factory, uintptr(i), uintptr(unsafe.Pointer(&adapter)))
|
||||
if r != 0 {
|
||||
break
|
||||
}
|
||||
var desc dxgiDesc
|
||||
syscall.SyscallN(callVtable(adapter, 8), adapter, uintptr(unsafe.Pointer(&desc))) // GetDesc
|
||||
release(adapter)
|
||||
|
||||
name := syscall.UTF16ToString(desc.Description[:])
|
||||
vendor := classifyByVID(desc.VendorId)
|
||||
log.Printf("[gpu][dxgi] %d: %s (VID=0x%04x)", i, name, desc.VendorId)
|
||||
if name == "" || isVirtualGPU(name) || vendor == "unknown" {
|
||||
continue
|
||||
}
|
||||
gpus = append(gpus, gpuName{Name: name, Vendor: vendor})
|
||||
}
|
||||
return gpus
|
||||
}
|
||||
|
||||
func callVtable(obj uintptr, idx int) uintptr {
|
||||
vtablePtr := *(**uintptr)(unsafe.Pointer(obj))
|
||||
arr := (*[64]uintptr)(unsafe.Pointer(vtablePtr))
|
||||
return arr[idx]
|
||||
}
|
||||
|
||||
func release(obj uintptr) { syscall.SyscallN(callVtable(obj, 2), obj) }
|
||||
|
||||
// --- Method 2: Registry ---
|
||||
func registryGPUList() []gpuName {
|
||||
cmd := exec.Command("reg", "query",
|
||||
`HKLM\SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}`,
|
||||
"/s", "/v", "DriverDesc")
|
||||
cmd.SysProcAttr = hideWindow
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return parseRegOutput(string(out))
|
||||
}
|
||||
|
||||
func parseRegOutput(out string) []gpuName {
|
||||
var gpus []gpuName
|
||||
seen := map[string]bool{}
|
||||
for _, line := range splitLines(out) {
|
||||
line = trim(line)
|
||||
if !startsWith(line, "DriverDesc") {
|
||||
continue
|
||||
}
|
||||
// Format: " DriverDesc REG_SZ NVIDIA GeForce RTX 4090"
|
||||
parts := splitN(line, "REG_SZ", 2)
|
||||
if len(parts) < 2 {
|
||||
continue
|
||||
}
|
||||
name := trim(parts[1])
|
||||
if name == "" || seen[name] || isVirtualGPU(name) {
|
||||
continue
|
||||
}
|
||||
vendor := classifyByName(name)
|
||||
if vendor == "unknown" {
|
||||
log.Printf("[gpu][reg] unknown vendor: %s", name)
|
||||
continue
|
||||
}
|
||||
seen[name] = true
|
||||
gpus = append(gpus, gpuName{Name: name, Vendor: vendor})
|
||||
}
|
||||
return gpus
|
||||
}
|
||||
|
||||
// --- Method 3: PowerShell ---
|
||||
func psGPUList() []gpuName {
|
||||
cmd := exec.Command("powershell", "-NoProfile", "-Command",
|
||||
`Get-CimInstance Win32_VideoController | ForEach-Object { $_.Name + "|" + $_.AdapterCompatibility }`,
|
||||
)
|
||||
cmd.SysProcAttr = hideWindow
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var gpus []gpuName
|
||||
seen := map[string]bool{}
|
||||
for _, line := range splitLines(string(out)) {
|
||||
line = trim(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
parts := splitN(line, "|", 2)
|
||||
name := trim(parts[0])
|
||||
vendorStr := ""
|
||||
if len(parts) > 1 {
|
||||
vendorStr = trim(parts[1])
|
||||
}
|
||||
if name == "" || seen[name] || isVirtualGPU(name) {
|
||||
continue
|
||||
}
|
||||
vendor := classifyByName(vendorStr + " " + name)
|
||||
if vendor == "unknown" {
|
||||
continue
|
||||
}
|
||||
seen[name] = true
|
||||
gpus = append(gpus, gpuName{Name: name, Vendor: vendor})
|
||||
}
|
||||
return gpus
|
||||
}
|
||||
|
||||
// --- CPU Detection ---
|
||||
func detectCPUName() string {
|
||||
cmd := exec.Command("reg", "query",
|
||||
`HKLM\HARDWARE\DESCRIPTION\System\CentralProcessor\0`,
|
||||
"/v", "ProcessorNameString")
|
||||
cmd.SysProcAttr = hideWindow
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
return detectCPUByEnv()
|
||||
}
|
||||
for _, line := range splitLines(string(out)) {
|
||||
line = trim(line)
|
||||
if !startsWith(line, "ProcessorNameString") {
|
||||
continue
|
||||
}
|
||||
parts := splitN(line, "REG_SZ", 2)
|
||||
if len(parts) >= 2 {
|
||||
return trim(parts[1])
|
||||
}
|
||||
}
|
||||
return detectCPUByEnv()
|
||||
}
|
||||
|
||||
func detectCPUByEnv() string {
|
||||
cmd := exec.Command("cmd", "/c", "echo %PROCESSOR_IDENTIFIER%")
|
||||
cmd.SysProcAttr = hideWindow
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
return "CPU"
|
||||
}
|
||||
name := trim(string(out))
|
||||
if name == "" {
|
||||
return "CPU"
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
// --- Orchestrator ---
|
||||
func detectGPUNames() []gpuName {
|
||||
// Try PowerShell first, then wmic
|
||||
gpus := psGPUList()
|
||||
if len(gpus) > 0 {
|
||||
log.Printf("[gpu] PowerShell found %d GPUs", len(gpus))
|
||||
for _, g := range gpus {
|
||||
log.Printf("[gpu] %s -> vendor=%s", g.Name, g.Vendor)
|
||||
}
|
||||
return gpus
|
||||
methods := []struct {
|
||||
name string
|
||||
fn func() []gpuName
|
||||
}{
|
||||
{"DXGI", dxgiGPUList},
|
||||
{"Registry", registryGPUList},
|
||||
{"PowerShell", psGPUList},
|
||||
}
|
||||
|
||||
gpus = wmicGPUList()
|
||||
if len(gpus) > 0 {
|
||||
log.Printf("[gpu] WMIC found %d GPUs", len(gpus))
|
||||
for _, g := range gpus {
|
||||
log.Printf("[gpu] %s -> vendor=%s", g.Name, g.Vendor)
|
||||
for _, m := range methods {
|
||||
gpus := m.fn()
|
||||
if len(gpus) > 0 {
|
||||
log.Printf("[gpu] %s detected %d GPU(s)", m.name, len(gpus))
|
||||
return gpus
|
||||
}
|
||||
return gpus
|
||||
log.Printf("[gpu] %s returned empty", m.name)
|
||||
}
|
||||
|
||||
log.Printf("[gpu] No GPUs detected via PowerShell or WMIC")
|
||||
return nil
|
||||
}
|
||||
|
||||
func psGPUList() []gpuName {
|
||||
out, err := exec.Command("powershell", "-NoProfile", "-Command",
|
||||
"Get-CimInstance Win32_VideoController | Select-Object Name,AdapterCompatibility,DriverVersion | ConvertTo-Csv -NoTypeInformation",
|
||||
).CombinedOutput()
|
||||
log.Printf("[gpu] PS output:\n%s", string(out))
|
||||
if err != nil {
|
||||
log.Printf("[gpu] PS error: %v", err)
|
||||
return nil
|
||||
}
|
||||
return parseCSV(string(out))
|
||||
}
|
||||
|
||||
func wmicGPUList() []gpuName {
|
||||
out, err := exec.Command("wmic", "path", "win32_VideoController",
|
||||
"get", "Name,AdapterCompatibility", "/format:csv").CombinedOutput()
|
||||
log.Printf("[gpu] WMIC output:\n%s", string(out))
|
||||
if err != nil {
|
||||
log.Printf("[gpu] WMIC error: %v", err)
|
||||
return nil
|
||||
}
|
||||
return parseWMIC(string(out))
|
||||
}
|
||||
|
||||
func parseCSV(out string) []gpuName {
|
||||
var gpus []gpuName
|
||||
seen := map[string]bool{}
|
||||
lines := strings.Split(out, "\n")
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "\"Name\"") || strings.HasPrefix(line, "Name") {
|
||||
continue
|
||||
}
|
||||
// Remove quotes: "Name","AdapterCompatibility"
|
||||
line = strings.ReplaceAll(line, "\"", "")
|
||||
parts := strings.SplitN(line, ",", 2)
|
||||
if len(parts) < 1 {
|
||||
continue
|
||||
}
|
||||
name := strings.TrimSpace(parts[0])
|
||||
vendor := ""
|
||||
if len(parts) >= 2 {
|
||||
vendor = strings.TrimSpace(parts[1])
|
||||
}
|
||||
if name == "" || seen[name] {
|
||||
continue
|
||||
}
|
||||
if isVirtualGPU(name) {
|
||||
log.Printf("[gpu] SKIP virtual: %s", name)
|
||||
continue
|
||||
}
|
||||
vendor = classifyVendor(vendor, name)
|
||||
if vendor == "unknown" {
|
||||
log.Printf("[gpu] SKIP unknown vendor (vendor=%q name=%q)", vendor, name)
|
||||
continue
|
||||
}
|
||||
seen[name] = true
|
||||
gpus = append(gpus, gpuName{Name: name, Vendor: vendor})
|
||||
}
|
||||
return gpus
|
||||
}
|
||||
|
||||
func parseWMIC(out string) []gpuName {
|
||||
var gpus []gpuName
|
||||
seen := map[string]bool{}
|
||||
lines := strings.Split(out, "\n")
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "Node,") {
|
||||
continue
|
||||
}
|
||||
parts := strings.SplitN(line, ",", 3)
|
||||
if len(parts) < 3 {
|
||||
continue
|
||||
}
|
||||
name := strings.TrimSpace(parts[2])
|
||||
vendor := strings.TrimSpace(parts[1])
|
||||
if name == "" || seen[name] {
|
||||
continue
|
||||
}
|
||||
if isVirtualGPU(name) {
|
||||
log.Printf("[gpu] SKIP virtual: %s", name)
|
||||
continue
|
||||
}
|
||||
vendor = classifyVendor(vendor, name)
|
||||
if vendor == "unknown" {
|
||||
log.Printf("[gpu] SKIP unknown vendor (vendor=%q name=%q)", vendor, name)
|
||||
continue
|
||||
}
|
||||
seen[name] = true
|
||||
gpus = append(gpus, gpuName{Name: name, Vendor: vendor})
|
||||
}
|
||||
return gpus
|
||||
}
|
||||
|
||||
func classifyVendor(vendor, name string) string {
|
||||
lower := strings.ToLower(vendor + " " + name)
|
||||
switch {
|
||||
case strings.Contains(lower, "nvidia"):
|
||||
func classifyByVID(vid uint32) string {
|
||||
switch vid {
|
||||
case 0x10DE:
|
||||
return "nvidia"
|
||||
case strings.Contains(lower, "intel") || strings.Contains(lower, "uhd graphics") ||
|
||||
strings.Contains(lower, "iris") || strings.Contains(lower, "hd graphics"):
|
||||
case 0x8086, 0x1414:
|
||||
return "intel"
|
||||
case strings.Contains(lower, "amd") || strings.Contains(lower, "radeon") || strings.Contains(lower, "ati"):
|
||||
case 0x1002:
|
||||
return "amd"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
func classifyByName(name string) string {
|
||||
l := toLower(name)
|
||||
switch {
|
||||
case strContains(l, "nvidia"):
|
||||
return "nvidia"
|
||||
case strContains(l, "intel") || strContains(l, "uhd") || strContains(l, "iris") || strContains(l, "hd graphics") || strContains(l, "arc"):
|
||||
return "intel"
|
||||
case strContains(l, "amd") || strContains(l, "radeon") || strContains(l, "ati") || strContains(l, "firepro"):
|
||||
return "amd"
|
||||
default:
|
||||
return "unknown"
|
||||
@@ -147,18 +237,90 @@ func classifyVendor(vendor, name string) string {
|
||||
}
|
||||
|
||||
func isVirtualGPU(name string) bool {
|
||||
lower := strings.ToLower(name)
|
||||
// Filter out virtual/remote/display-only adapters that don't have encoders
|
||||
virtual := []string{
|
||||
l := toLower(name)
|
||||
for _, v := range []string{
|
||||
"virtual", "remote", "rdp", "citrix", "vmware", "hyper-v",
|
||||
"mirror", "indirect", "parsec", "splashtop", "idm",
|
||||
"mirror", "indirect", "parsec", "splashtop",
|
||||
"mirage", "vnc", "displayonly", "basicdisplay",
|
||||
"microsoft basic", "microsoft remote",
|
||||
}
|
||||
for _, v := range virtual {
|
||||
if strings.Contains(lower, v) {
|
||||
} {
|
||||
if strContains(l, v) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func toLower(s string) string {
|
||||
b := make([]byte, len(s))
|
||||
for i := 0; i < len(s); i++ {
|
||||
c := s[i]
|
||||
if c >= 'A' && c <= 'Z' {
|
||||
c += 32
|
||||
}
|
||||
b[i] = c
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func strContains(s, sub string) bool {
|
||||
return len(sub) <= len(s) && indexOf(s, sub) >= 0
|
||||
}
|
||||
|
||||
func indexOf(s, sub string) int {
|
||||
for i := 0; i <= len(s)-len(sub); i++ {
|
||||
if s[i:i+len(sub)] == sub {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func splitLines(s string) []string {
|
||||
var lines []string
|
||||
for {
|
||||
i := indexOf(s, "\n")
|
||||
if i < 0 {
|
||||
if s != "" {
|
||||
lines = append(lines, s)
|
||||
}
|
||||
break
|
||||
}
|
||||
if i > 0 && s[i-1] == '\r' {
|
||||
lines = append(lines, s[:i-1])
|
||||
} else {
|
||||
lines = append(lines, s[:i])
|
||||
}
|
||||
s = s[i+1:]
|
||||
}
|
||||
return lines
|
||||
}
|
||||
|
||||
func trim(s string) string {
|
||||
for len(s) > 0 && (s[0] == ' ' || s[0] == '\t' || s[0] == '\r' || s[0] == '\n') {
|
||||
s = s[1:]
|
||||
}
|
||||
for len(s) > 0 && (s[len(s)-1] == ' ' || s[len(s)-1] == '\t' || s[len(s)-1] == '\r' || s[len(s)-1] == '\n') {
|
||||
s = s[:len(s)-1]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func startsWith(s, prefix string) bool {
|
||||
return len(s) >= len(prefix) && s[:len(prefix)] == prefix
|
||||
}
|
||||
|
||||
func splitN(s, sep string, n int) []string {
|
||||
var r []string
|
||||
for n > 1 {
|
||||
i := indexOf(s, sep)
|
||||
if i < 0 {
|
||||
break
|
||||
}
|
||||
r = append(r, s[:i])
|
||||
s = s[i+len(sep):]
|
||||
n--
|
||||
}
|
||||
r = append(r, s)
|
||||
return r
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user