fetch_ml/internal/worker/native_bridge_libs.go
Jeremie Fraeys be39b37aec
feat: native GPU detection and NVML bridge for macOS and Linux
- Add dynamic NVML loading for Linux GPU detection
- Add macOS GPU detection via IOKit framework
- Add Zig NVML wrapper for cross-platform GPU queries
- Update native bridge to support platform-specific GPU libs
- Add CMake support for NVML dynamic library
2026-02-21 17:59:59 -05:00

78 lines
1.8 KiB
Go

//go:build cgo && native_libs
// +build cgo,native_libs
package worker
// #cgo darwin LDFLAGS: -L${SRCDIR}/../../native/build -Wl,-rpath,${SRCDIR}/../../native/build -ldataset_hash
// #cgo linux LDFLAGS: -L${SRCDIR}/../../native/build -Wl,-rpath,${SRCDIR}/../../native/build -ldataset_hash -lnvml_gpu -lnvidia-ml
// #include "../../native/dataset_hash/dataset_hash.h"
// #include <stdlib.h>
import "C"
import (
"errors"
"log"
"runtime"
"sync"
"time"
"unsafe"
"github.com/jfraeys/fetch_ml/internal/manifest"
)
var (
hashCtx *C.fh_context_t
hashCtxOnce sync.Once
ctxInitTime time.Time
)
func getHashContext() *C.fh_context_t {
hashCtxOnce.Do(func() {
start := time.Now()
hashCtx = C.fh_init(C.uint32_t(runtime.NumCPU()))
ctxInitTime = time.Now()
log.Printf("[native] hash context initialized: %v (threads: %d)",
time.Since(start), runtime.NumCPU())
})
return hashCtx
}
func dirOverallSHA256HexNative(root string) (string, error) {
ctx := getHashContext()
croot := C.CString(root)
defer C.free(unsafe.Pointer(croot))
result := C.fh_hash_directory_combined(ctx, croot)
if result == nil {
err := C.fh_last_error(ctx)
if err != nil {
return "", errors.New(C.GoString(err))
}
return "", errors.New("native hash failed")
}
defer C.fh_free_string(result)
return C.GoString(result), nil
}
func GetSIMDImplName() string {
return C.GoString(C.fh_get_simd_impl_name())
}
func HasSIMDSHA256() bool {
return C.fh_has_simd_sha256() == 1
}
func ScanArtifactsNative(runDir string) (*manifest.Artifacts, error) {
return ScanArtifacts(runDir)
}
func ExtractTarGzNative(archivePath, dstDir string) error {
return ExtractTarGz(archivePath, dstDir)
}
// DirOverallSHA256HexNative exports the native hash implementation for benchmarks.
func DirOverallSHA256HexNative(root string) (string, error) {
return dirOverallSHA256HexNative(root)
}