From a16a86c31a2fc515bfdf8e0569c597457f10e968 Mon Sep 17 00:00:00 2001 From: Kevin Klues Date: Tue, 22 Dec 2020 07:07:48 +0000 Subject: [PATCH] Only check for nvmlEventSetWait_v2 symbol once at initizilation Signed-off-by: Kevin Klues --- bindings/go/nvml/bindings.go | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/bindings/go/nvml/bindings.go b/bindings/go/nvml/bindings.go index 6898918..a650c13 100644 --- a/bindings/go/nvml/bindings.go +++ b/bindings/go/nvml/bindings.go @@ -48,6 +48,8 @@ const ( XidCriticalError = C.nvmlEventTypeXidCriticalError ) +var nvmlEventSetWait = nvmlEventSetWait_v1 + type handle struct{ dev C.nvmlDevice_t } type EventSet struct{ set C.nvmlEventSet_t } type Event struct { @@ -58,6 +60,14 @@ type Event struct { Edata uint64 } +func nvmlEventSetWait_v1(Set C.nvmlEventSet_t, Data *C.nvmlEventData_t, Timeoutms C.uint) C.nvmlReturn_t { + return C.nvmlEventSetWait(Set, Data, Timeoutms) +} + +func nvmlEventSetWait_v2(Set C.nvmlEventSet_t, Data *C.nvmlEventData_t, Timeoutms C.uint) C.nvmlReturn_t { + return C.nvmlEventSetWait_v2(Set, Data, Timeoutms) +} + func uintPtr(c C.uint) *uint { i := uint(c) return &i @@ -86,6 +96,12 @@ func init_() error { if r == C.NVML_ERROR_LIBRARY_NOT_FOUND { return errors.New("could not load NVML library") } + + found := dl.lookupSymbol("nvmlEventSetWait_v2") + if found == C.NVML_SUCCESS { + nvmlEventSetWait = nvmlEventSetWait_v2 + } + return errorString(r) } @@ -157,15 +173,10 @@ func DeleteEventSet(es EventSet) { func WaitForEvent(es EventSet, timeout uint) (Event, error) { var data C.nvmlEventData_t + data.gpuInstanceId = 0xFFFFFFFF + data.computeInstanceId = 0xFFFFFFFF - r := dl.lookupSymbol("nvmlEventSetWait_v2") - if r == C.NVML_SUCCESS { - r = C.nvmlEventSetWait_v2(es.set, &data, C.uint(timeout)) - } else { - r = C.nvmlEventSetWait(es.set, &data, C.uint(timeout)) - data.gpuInstanceId = 0xFFFFFFFF - data.computeInstanceId = 0xFFFFFFFF - } + r := nvmlEventSetWait(es.set, &data, C.uint(timeout)) if r != C.NVML_SUCCESS { return Event{}, errorString(r) }