gpu/darwin: pure go implemention (#57)

This commit is contained in:
UUBulb
2024-08-26 22:57:07 +08:00
committed by GitHub
parent f03671686f
commit 42187f2bf2
11 changed files with 230 additions and 340 deletions
+1 -1
View File
@@ -26,7 +26,7 @@ func (rsmi *ROCmSMI) Start() error {
if _, err := os.Stat(rsmi.BinPath); os.IsNotExist(err) {
binPath, err := exec.LookPath("rocm-smi")
if err != nil {
return errors.New("Didn't find the adequate tool to query GPU utilization")
return errors.New("didn't find the adequate tool to query GPU utilization")
}
rsmi.BinPath = binPath
}
-144
View File
@@ -1,144 +0,0 @@
#include "gpu_darwin.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#define IOSERVICE_GPU "IOAccelerator"
#define IOSERVICE_PCI "IOPCIDevice"
void *find_properties(io_registry_entry_t service, int depth, CFStringRef key,
CFStringRef dict_key) {
CFTypeRef properties = IORegistryEntrySearchCFProperty(
service, kIOServicePlane, key, kCFAllocatorDefault,
kIORegistryIterateRecursively);
if (properties) {
if (CFGetTypeID(properties) == CFStringGetTypeID()) {
CFStringRef cfStr = (CFStringRef)properties;
char buffer[1024];
CFStringGetCString(cfStr, buffer, sizeof(buffer), kCFStringEncodingUTF8);
CFRelease(properties);
return strdup(buffer);
} else if (CFGetTypeID(properties) == CFDictionaryGetTypeID()) {
CFDictionaryRef cfDict = (CFDictionaryRef)properties;
CFNumberRef cfValue = (CFNumberRef)CFDictionaryGetValue(cfDict, dict_key);
if (cfValue == NULL) {
return NULL;
}
int value;
if (!CFNumberGetValue(cfValue, kCFNumberIntType, &value)) {
return NULL;
}
return (void *)(intptr_t)value;
}
}
return NULL;
}
char **find_devices(char *key) {
io_service_t io_reg_err;
io_iterator_t iterator;
int capacity = 10;
char **cards = malloc(capacity * sizeof(char *));
if (!cards) {
fprintf(stderr, "Memory allocation failed\n");
return NULL;
}
io_reg_err = IOServiceGetMatchingServices(
kIOMainPortDefault, IOServiceMatching(IOSERVICE_GPU), &iterator);
if (io_reg_err != KERN_SUCCESS) {
printf("Error getting GPU entry\n");
return NULL;
}
io_object_t service;
int index = 0;
while ((service = IOIteratorNext(iterator)) != MACH_PORT_NULL) {
CFStringRef cfStr = CFStringCreateWithCString(kCFAllocatorDefault, key,
kCFStringEncodingUTF8);
char *result = find_properties(service, 0, cfStr, CFSTR(""));
CFRelease(cfStr);
IOObjectRelease(service);
if (result != NULL) {
if (index >= capacity) {
capacity += 1;
char **new_cards = (char **)realloc(cards, capacity * sizeof(char *));
if (!new_cards) {
fprintf(stderr, "Memory reallocation failed\n");
for (int i = 0; i < index; i++) {
free(cards[i]);
}
free(cards);
free(result);
return NULL;
}
cards = new_cards;
}
cards[index] = result;
index++;
}
if (result == NULL && strcmp(key, "model") == 0) {
IOObjectRelease(iterator);
io_reg_err = IOServiceGetMatchingServices(
kIOMainPortDefault, IOServiceMatching(IOSERVICE_PCI), &iterator);
if (io_reg_err != KERN_SUCCESS) {
printf("Error getting PCI entry\n");
return NULL;
}
}
}
IOObjectRelease(iterator);
char **result_cards = (char **)realloc(cards, sizeof(char *) * (index + 1));
if (!result_cards) {
fprintf(stderr, "Memory reallocation failed\n");
for (int i = 0; i < index; i++) {
free(cards[i]);
}
free(cards);
return NULL;
}
result_cards[index] = NULL;
return result_cards;
}
int find_utilization(char *key, char *dict_key) {
void *result_ptr;
io_service_t io_reg_err;
io_iterator_t iterator;
io_reg_err = IOServiceGetMatchingServices(
kIOMainPortDefault, IOServiceMatching(IOSERVICE_GPU), &iterator);
if (io_reg_err != KERN_SUCCESS) {
printf("Error getting GPU entry\n");
return 0;
}
io_object_t service = IOIteratorNext(iterator);
if (service != MACH_PORT_NULL) {
CFStringRef cfStr = CFStringCreateWithCString(kCFAllocatorDefault, key,
kCFStringEncodingUTF8);
CFStringRef cfDictStr = CFStringCreateWithCString(
kCFAllocatorDefault, dict_key, kCFStringEncodingUTF8);
result_ptr = find_properties(service, 0, cfStr, cfDictStr);
CFRelease(cfStr);
CFRelease(cfDictStr);
}
IOObjectRelease(service);
IOObjectRelease(iterator);
if (result_ptr == NULL) {
return 0;
}
return (int)(intptr_t)result_ptr;
}
-15
View File
@@ -1,15 +0,0 @@
#ifndef __SMC_H__
#define __SMC_H__ 1
#include <IOKit/IOKitLib.h>
#include <CoreFoundation/CoreFoundation.h>
#if (defined __MAC_OS_X_VERSION_MIN_REQUIRED) && (__MAC_OS_X_VERSION_MIN_REQUIRED < 120000)
#define kIOMainPortDefault kIOMasterPortDefault
#endif
void *find_properties(io_registry_entry_t, int, CFStringRef, CFStringRef);
char **find_devices(char *);
int find_utilization(char *, char *);
#endif
+1 -1
View File
@@ -26,7 +26,7 @@ func (smi *NvidiaSMI) Start() error {
if _, err := os.Stat(smi.BinPath); os.IsNotExist(err) {
binPath, err := exec.LookPath("nvidia-smi")
if err != nil {
return errors.New("Didn't find the adequate tool to query GPU utilization")
return errors.New("didn't find the adequate tool to query GPU utilization")
}
smi.BinPath = binPath
}
+4 -28
View File
@@ -1,36 +1,12 @@
//go:build darwin && !cgo
//go:build darwin
package stat
import (
"os/exec"
"regexp"
"strconv"
"github.com/nezhahq/agent/pkg/gpu"
)
func extractGPUStat(cmd *exec.Cmd) ([]float64, error) {
gs, err := cmd.CombinedOutput()
if err != nil {
return nil, err
}
re := regexp.MustCompile(`"Device Utilization %"\s*=\s*(\d+)`)
matches := re.FindAllSubmatch(gs, -1)
var u []float64
for _, match := range matches {
if len(match) > 1 {
p, _ := strconv.ParseFloat(string(match[1]), 64)
u = append(u, p)
}
}
return u, nil
}
func GetGPUStat() (float64, error) {
ioreg := exec.Command("ioreg", "-rd1", "-c", "IOAccelerator")
gs, err := extractGPUStat(ioreg)
if err != nil || len(gs) == 0 {
return 0, err
}
return gs[0], nil
usage, err := gpu.FindUtilization("PerformanceStatistics", "Device Utilization %")
return float64(usage), err
}
-25
View File
@@ -1,25 +0,0 @@
//go:build darwin && cgo
package stat
// #cgo LDFLAGS: -framework IOKit -framework CoreFoundation
// #include "gpu_darwin.h"
import "C"
import (
"unsafe"
)
func extractGPUStat(key *C.char, dict_key *C.char) (int, error) {
utilization := C.find_utilization(key, dict_key)
return int(utilization), nil
}
func GetGPUStat() (float64, error) {
key := C.CString("PerformanceStatistics")
dict_key := C.CString("Device Utilization %")
defer C.free(unsafe.Pointer(key))
defer C.free(unsafe.Pointer(dict_key))
gs, _ := extractGPUStat(key, dict_key)
return float64(gs), nil
}