gpu.go

//go:build linux

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/NVIDIA/go-nvml/pkg/nvml"
)

func detectGPU() GPUInfo {
	info := GPUInfo{}

	ret := nvml.Init()
	if ret != nvml.SUCCESS {
		info.Error = fmt.Sprintf("NVML init failed: %s (is NVIDIA driver installed?)", nvml.ErrorString(ret))
		return info
	}
	defer nvml.Shutdown()

	info.Available = true

	if v, ret := nvml.SystemGetDriverVersion(); ret == nvml.SUCCESS {
		info.DriverVersion = v
	}

	if v, ret := nvml.SystemGetCudaDriverVersion_v2(); ret == nvml.SUCCESS {
		major := v / 1000
		minor := (v % 1000) / 10
		info.CUDAVersion = fmt.Sprintf("%d.%d", major, minor)
	}

	info.CUDAToolkit = detectCUDAToolkit()

	count, ret := nvml.DeviceGetCount()
	if ret != nvml.SUCCESS {
		info.Error = fmt.Sprintf("Failed to get device count: %s", nvml.ErrorString(ret))
		return info
	}
	info.DeviceCount = count

	for i := 0; i < count; i++ {
		device, ret := nvml.DeviceGetHandleByIndex(i)
		if ret != nvml.SUCCESS {
			continue
		}

		dev := GPUDevice{Index: i}

		if name, ret := device.GetName(); ret == nvml.SUCCESS {
			dev.Name = name
		}

		// Memory detection cascade:
		//   1. NVML GetMemoryInfo_v2 (supports unified memory on newer drivers)
		//   2. NVML GetMemoryInfo (standard path)
		//   3. cudaMemGetInfo via dlopen (may work where NVML doesn't, e.g. GB10)
		//   4. /proc/meminfo fallback (unified memory, use system RAM)
		memReported := false
		if mem, ret := device.GetMemoryInfo_v2(); ret == nvml.SUCCESS && mem.Total > 0 {
			dev.MemoryTotalGB = float64(mem.Total) / 1024 / 1024 / 1024
			dev.MemoryFreeGB = float64(mem.Free) / 1024 / 1024 / 1024
			dev.MemoryUsedGB = float64(mem.Used) / 1024 / 1024 / 1024
			dev.MemorySource = "nvml"
			memReported = true
		}
		if !memReported {
			if mem, ret := device.GetMemoryInfo(); ret == nvml.SUCCESS && mem.Total > 0 {
				dev.MemoryTotalGB = float64(mem.Total) / 1024 / 1024 / 1024
				dev.MemoryFreeGB = float64(mem.Free) / 1024 / 1024 / 1024
				dev.MemoryUsedGB = float64(mem.Used) / 1024 / 1024 / 1024
				dev.MemorySource = "nvml"
				memReported = true
			}
		}
		// Try CUDA runtime API — may succeed on unified memory architectures
		// where NVML reports 0 (e.g. GB10/DGX Spark).
		if !memReported {
			if cudart, err := NewCUDARuntime(); err == nil {
				if err := cudart.SetDevice(i); err == nil {
					if free, total, err := cudart.MemGetInfo(); err == nil && total > 0 {
						dev.MemoryTotalGB = float64(total) / 1024 / 1024 / 1024
						dev.MemoryFreeGB = float64(free) / 1024 / 1024 / 1024
						dev.MemoryUsedGB = float64(total-free) / 1024 / 1024 / 1024
						dev.MemorySource = "cuda_runtime"
						memReported = true
					}
				}
				cudart.Close()
			}
		}
		// Unified memory architectures (e.g. GB10/Grace-Blackwell) report 0
		// from both NVML and CUDA. Fall back to system RAM since GPU shares it.
		if !memReported {
			dev.UnifiedMemory = true
			dev.MemorySource = "system_ram"
			dev.MemoryTotalGB = parseMemTotal()
		}

		if major, minor, ret := device.GetCudaComputeCapability(); ret == nvml.SUCCESS {
			dev.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
		}

		if mode, _, ret := device.GetEccMode(); ret == nvml.SUCCESS {
			dev.ECCEnabled = mode == nvml.FEATURE_ENABLED
		}

		if mode, ret := device.GetPersistenceMode(); ret == nvml.SUCCESS && mode == nvml.FEATURE_ENABLED {
			info.PersistenceMode = true
		}

		info.Devices = append(info.Devices, dev)
	}

	return info
}

func detectNVIDIAModules(gpu GPUInfo) NVIDIAModules {
	info := NVIDIAModules{}
	if _, err := os.Stat("/sys/module/nvidia"); err == nil {
		info.NvidiaLoaded = true
	}
	if _, err := os.Stat("/sys/module/nvidia_uvm"); err == nil {
		info.UVMLoaded = true
	}
	if _, err := os.Stat("/sys/module/nvidia_modeset"); err == nil {
		info.ModesetLoaded = true
	}
	if _, err := os.Stat("/sys/module/nvidia_drm"); err == nil {
		info.DRMLoaded = true
	}
	info.PersistenceMode = gpu.PersistenceMode
	return info
}

func detectCUDAToolkit() string {
	if out, err := exec.Command("nvcc", "--version").Output(); err == nil {
		for _, line := range strings.Split(string(out), "\n") {
			if strings.Contains(line, "release") {
				parts := strings.Split(line, "release ")
				if len(parts) >= 2 {
					ver := strings.Split(parts[1], ",")[0]
					return strings.TrimSpace(ver)
				}
			}
		}
	}

	matches, _ := filepath.Glob("/usr/local/cuda*/version.json")
	if _, err := os.Stat("/usr/local/cuda/version.json"); err == nil {
		matches = append([]string{"/usr/local/cuda/version.json"}, matches...)
	}
	for _, path := range matches {
		data, err := os.ReadFile(path)
		if err != nil {
			continue
		}
		var vj struct {
			CUDA struct {
				Version string `json:"version"`
			} `json:"cuda"`
		}
		if json.Unmarshal(data, &vj) == nil && vj.CUDA.Version != "" {
			return vj.CUDA.Version
		}
	}

	return ""
}