Skip to content

Commit

Permalink
Generate updated code
Browse files Browse the repository at this point in the history 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
  • Loading branch information
@elezar
elezar committed Oct 12, 2023
1 parent cc9e7fb commit e0f6516
Show file tree
Hide file tree
Showing 3 changed files with 277 additions and 170 deletions.
172 changes: 3 additions & 169 deletions pkg/nvml/init.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,43 +16,23 @@ package nvml

import (
"fmt"

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

import "C"

const (
nvmlLibraryName = "libnvidia-ml.so.1"
nvmlLibraryLoadFlags = dl.RTLD_LAZY | dl.RTLD_GLOBAL
)

var nvml *dl.DynamicLibrary

// nvml.Init()
func Init() Return {
lib := dl.New(nvmlLibraryName, nvmlLibraryLoadFlags)
err := lib.Open()
if err != nil {
if err := nvml.Open(); err != nil {
return ERROR_LIBRARY_NOT_FOUND
}

nvml = lib
updateVersionedSymbols()

return nvmlInit()
}

// nvml.InitWithFlags()
func InitWithFlags(Flags uint32) Return {
lib := dl.New(nvmlLibraryName, nvmlLibraryLoadFlags)
err := lib.Open()
if err != nil {
if err := nvml.Open(); err != nil {
return ERROR_LIBRARY_NOT_FOUND
}

nvml = lib

return nvmlInitWithFlags(Flags)
}

Expand All @@ -65,154 +45,8 @@ func Shutdown() Return {

err := nvml.Close()
if err != nil {
panic(fmt.Sprintf("error closing %s: %v", nvmlLibraryName, err))
panic(fmt.Sprintf("error closing %s: %v", nvml, err))
}

return ret
}

// Default all versioned APIs to v1 (to infer the types)
var nvmlInit = nvmlInit_v1
var nvmlDeviceGetPciInfo = nvmlDeviceGetPciInfo_v1
var nvmlDeviceGetCount = nvmlDeviceGetCount_v1
var nvmlDeviceGetHandleByIndex = nvmlDeviceGetHandleByIndex_v1
var nvmlDeviceGetHandleByPciBusId = nvmlDeviceGetHandleByPciBusId_v1
var nvmlDeviceGetNvLinkRemotePciInfo = nvmlDeviceGetNvLinkRemotePciInfo_v1
var nvmlDeviceRemoveGpu = nvmlDeviceRemoveGpu_v1
var nvmlDeviceGetGridLicensableFeatures = nvmlDeviceGetGridLicensableFeatures_v1
var nvmlEventSetWait = nvmlEventSetWait_v1
var nvmlDeviceGetAttributes = nvmlDeviceGetAttributes_v1
var nvmlComputeInstanceGetInfo = nvmlComputeInstanceGetInfo_v1
var DeviceGetComputeRunningProcesses = deviceGetComputeRunningProcesses_v1
var DeviceGetGraphicsRunningProcesses = deviceGetGraphicsRunningProcesses_v1
var DeviceGetMPSComputeRunningProcesses = deviceGetMPSComputeRunningProcesses_v1
var GetBlacklistDeviceCount = GetExcludedDeviceCount
var GetBlacklistDeviceInfoByIndex = GetExcludedDeviceInfoByIndex
var nvmlDeviceGetGpuInstancePossiblePlacements = nvmlDeviceGetGpuInstancePossiblePlacements_v1
var nvmlVgpuInstanceGetLicenseInfo = nvmlVgpuInstanceGetLicenseInfo_v1

type BlacklistDeviceInfo = ExcludedDeviceInfo
type ProcessInfo_v1Slice []ProcessInfo_v1
type ProcessInfo_v2Slice []ProcessInfo_v2

func (pis ProcessInfo_v1Slice) ToProcessInfoSlice() []ProcessInfo {
var newInfos []ProcessInfo
for _, pi := range pis {
info := ProcessInfo{
Pid: pi.Pid,
UsedGpuMemory: pi.UsedGpuMemory,
GpuInstanceId: 0xFFFFFFFF, // GPU instance ID is invalid in v1
ComputeInstanceId: 0xFFFFFFFF, // Compute instance ID is invalid in v1
}
newInfos = append(newInfos, info)
}
return newInfos
}

func (pis ProcessInfo_v2Slice) ToProcessInfoSlice() []ProcessInfo {
var newInfos []ProcessInfo
for _, pi := range pis {
info := ProcessInfo{
Pid: pi.Pid,
UsedGpuMemory: pi.UsedGpuMemory,
GpuInstanceId: pi.GpuInstanceId,
ComputeInstanceId: pi.ComputeInstanceId,
}
newInfos = append(newInfos, info)
}
return newInfos
}

// updateVersionedSymbols()
func updateVersionedSymbols() {
err := nvml.Lookup("nvmlInit_v2")
if err == nil {
nvmlInit = nvmlInit_v2
}
err = nvml.Lookup("nvmlDeviceGetPciInfo_v2")
if err == nil {
nvmlDeviceGetPciInfo = nvmlDeviceGetPciInfo_v2
}
err = nvml.Lookup("nvmlDeviceGetPciInfo_v3")
if err == nil {
nvmlDeviceGetPciInfo = nvmlDeviceGetPciInfo_v3
}
err = nvml.Lookup("nvmlDeviceGetCount_v2")
if err == nil {
nvmlDeviceGetCount = nvmlDeviceGetCount_v2
}
err = nvml.Lookup("nvmlDeviceGetHandleByIndex_v2")
if err == nil {
nvmlDeviceGetHandleByIndex = nvmlDeviceGetHandleByIndex_v2
}
err = nvml.Lookup("nvmlDeviceGetHandleByPciBusId_v2")
if err == nil {
nvmlDeviceGetHandleByPciBusId = nvmlDeviceGetHandleByPciBusId_v2
}
err = nvml.Lookup("nvmlDeviceGetNvLinkRemotePciInfo_v2")
if err == nil {
nvmlDeviceGetNvLinkRemotePciInfo = nvmlDeviceGetNvLinkRemotePciInfo_v2
}
// Unable to overwrite nvmlDeviceRemoveGpu() because the v2 function takes
// a different set of parameters than the v1 function.
//err = nvml.Lookup("nvmlDeviceRemoveGpu_v2")
//if err == nil {
// nvmlDeviceRemoveGpu = nvmlDeviceRemoveGpu_v2
//}
err = nvml.Lookup("nvmlDeviceGetGridLicensableFeatures_v2")
if err == nil {
nvmlDeviceGetGridLicensableFeatures = nvmlDeviceGetGridLicensableFeatures_v2
}
err = nvml.Lookup("nvmlDeviceGetGridLicensableFeatures_v3")
if err == nil {
nvmlDeviceGetGridLicensableFeatures = nvmlDeviceGetGridLicensableFeatures_v3
}
err = nvml.Lookup("nvmlDeviceGetGridLicensableFeatures_v4")
if err == nil {
nvmlDeviceGetGridLicensableFeatures = nvmlDeviceGetGridLicensableFeatures_v4
}
err = nvml.Lookup("nvmlEventSetWait_v2")
if err == nil {
nvmlEventSetWait = nvmlEventSetWait_v2
}
err = nvml.Lookup("nvmlDeviceGetAttributes_v2")
if err == nil {
nvmlDeviceGetAttributes = nvmlDeviceGetAttributes_v2
}
err = nvml.Lookup("nvmlComputeInstanceGetInfo_v2")
if err == nil {
nvmlComputeInstanceGetInfo = nvmlComputeInstanceGetInfo_v2
}
err = nvml.Lookup("nvmlDeviceGetComputeRunningProcesses_v2")
if err == nil {
DeviceGetComputeRunningProcesses = deviceGetComputeRunningProcesses_v2
}
err = nvml.Lookup("nvmlDeviceGetComputeRunningProcesses_v3")
if err == nil {
DeviceGetComputeRunningProcesses = deviceGetComputeRunningProcesses_v3
}
err = nvml.Lookup("nvmlDeviceGetGraphicsRunningProcesses_v2")
if err == nil {
DeviceGetGraphicsRunningProcesses = deviceGetGraphicsRunningProcesses_v2
}
err = nvml.Lookup("nvmlDeviceGetGraphicsRunningProcesses_v3")
if err == nil {
DeviceGetGraphicsRunningProcesses = deviceGetGraphicsRunningProcesses_v3
}
err = nvml.Lookup("nvmlDeviceGetMPSComputeRunningProcesses_v2")
if err == nil {
DeviceGetMPSComputeRunningProcesses = deviceGetMPSComputeRunningProcesses_v2
}
err = nvml.Lookup("nvmlDeviceGetMPSComputeRunningProcesses_v3")
if err == nil {
DeviceGetMPSComputeRunningProcesses = deviceGetMPSComputeRunningProcesses_v3
}
err = nvml.Lookup("nvmlDeviceGetGpuInstancePossiblePlacements_v2")
if err == nil {
nvmlDeviceGetGpuInstancePossiblePlacements = nvmlDeviceGetGpuInstancePossiblePlacements_v2
}
err = nvml.Lookup("nvmlVgpuInstanceGetLicenseInfo_v2")
if err == nil {
nvmlVgpuInstanceGetLicenseInfo = nvmlVgpuInstanceGetLicenseInfo_v2
}
}
Loading

0 comments on commit e0f6516

Please sign in to comment.