Source of nvml_windows.go - NVIDIA GPU

Raw file
Source viewDiff to previous
18.63 KB
xxxxxxxxxx
//   - `pendingEnabled` (bool): `true` if ECC will be enabled on the next reboot, `false` if it will be disabled.
 
package nvml
​
/*
#cgo CFLAGS: -I${SRCDIR}/nvml-sdk/include
#cgo CFLAGS: -DNVML_NO_UNVERSIONED_FUNC_DEFS=1
​
#include "nvml.h"
*/
import "C"
import (
    "errors"
    "sync"
    "syscall"
    "unsafe"
​
    "golang.org/x/sys/windows"
    "golang.zabbix.com/sdk/errs"
)
​
var (
    _ Device = (*NVMLDevice)(nil)
    _ Runner = (*NVMLRunner)(nil)
)
​
// NVMLRunner manages the loading and retrieval of functions from a Windows DLL.
// It is responsible for ensuring thread-safe access to the list of loaded
// procedures and facilitates dynamic function calls from the DLL.
type NVMLRunner struct {
    dll         *windows.DLL
    procListMux *sync.Mutex
    procList    map[string]*windows.Proc
}
​
// NVMLDevice represents an NVML GPU device, identified by a unique handle.
type NVMLDevice struct {
    handle uintptr
    runner *NVMLRunner // Reference to the Runner (formerly NVMLRunner)
}
​
// NewNVMLRunner creates a new NVML Runner instance, loading the NVML library.
func NewNVMLRunner() (*NVMLRunner, error) {
    dll, err := windows.LoadDLL("nvml.dll")
    if err != nil {
        return nil, errs.WrapConst(err, ErrLibraryNotFound) //nolint:wrapcheck
    }
​
    runner := &NVMLRunner{
        dll:         dll,
        procList:    make(map[string]*windows.Proc),
        procListMux: &sync.Mutex{},
    }
​
    return runner, nil
}
​
// Init initializes the NVML library using the older NVML interface.
func (runner *NVMLRunner) Init() error {
    err := runner.callProc("nvmlInit")
    if err != nil {
        return errs.Wrap(err, "failed while calling procedure")
    }
​
    return nil
}
​
// InitV2 initializes the NVML library using the NVML v2 interface.
func (runner *NVMLRunner) InitV2() error {
    err := runner.callProc("nvmlInit_v2")
    if err != nil {
        return errs.Wrap(err, "failed while calling procedure")
    }
​
    return nil
}
​
// GetNVMLVersion retrieves the version of the NVML library currently in use.
func (runner *NVMLRunner) GetNVMLVersion() (string, error) {
    var version [systemNVMLVersionBufferSize]byte
​
    err := runner.callProc("nvmlSystemGetNVMLVersion",
        uintptr(unsafe.Pointer(&version[0])),
        uintptr(systemNVMLVersionBufferSize),
    )
    if err != nil {
        return "", errs.Wrap(err, "failed while calling procedure")
    }
​
    return windows.ByteSliceToString(version[:]), nil
}
​
// GetDriverVersion retrieves the version of the NVIDIA driver currently in use.