59 lines
1.4 KiB
Go
59 lines
1.4 KiB
Go
|
//go:build linux || windows
|
||
|
|
|
||
|
|
package gpu
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"log/slog"
|
||
|
|
"os"
|
||
|
|
"path/filepath"
|
||
|
|
"strconv"
|
||
|
|
"strings"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Determine if the given ROCm lib directory is usable by checking for existence of some glob patterns
|
||
|
|
func rocmLibUsable(libDir string) bool {
|
||
|
|
slog.Debug("evaluating potential rocm lib dir " + libDir)
|
||
|
|
for _, g := range ROCmLibGlobs {
|
||
|
|
res, _ := filepath.Glob(filepath.Join(libDir, g))
|
||
|
|
if len(res) == 0 {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
func GetSupportedGFX(libDir string) ([]string, error) {
|
||
|
|
var ret []string
|
||
|
|
files, err := filepath.Glob(filepath.Join(libDir, "rocblas", "library", "TensileLibrary_lazy_gfx*.dat"))
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
for _, file := range files {
|
||
|
|
ret = append(ret, strings.TrimSuffix(strings.TrimPrefix(filepath.Base(file), "TensileLibrary_lazy_"), ".dat"))
|
||
|
|
}
|
||
|
|
return ret, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func amdSetVisibleDevices(ids []int, skip map[int]interface{}) {
|
||
|
|
// Set the visible devices if not already set
|
||
|
|
// TODO - does sort order matter?
|
||
|
|
devices := []string{}
|
||
|
|
for i := range ids {
|
||
|
|
slog.Debug(fmt.Sprintf("i=%d", i))
|
||
|
|
if _, skipped := skip[i]; skipped {
|
||
|
|
slog.Debug("skipped")
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
devices = append(devices, strconv.Itoa(i))
|
||
|
|
}
|
||
|
|
slog.Debug(fmt.Sprintf("devices=%v", devices))
|
||
|
|
|
||
|
|
val := strings.Join(devices, ",")
|
||
|
|
err := os.Setenv("HIP_VISIBLE_DEVICES", val)
|
||
|
|
if err != nil {
|
||
|
|
slog.Warn(fmt.Sprintf("failed to set env: %s", err))
|
||
|
|
}
|
||
|
|
slog.Debug("HIP_VISIBLE_DEVICES=" + val)
|
||
|
|
}
|