baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

gpu: add env var for detecting Intel oneapi gpus (#5076)

Browse source 
* gpu: add env var for detecting intel oneapi gpus

* fix build error
This commit is contained in:
Jeffrey Morgan 2024-06-16 20:09:05 -04:00 committed by GitHub
parent 4c2c8f93dd
commit 163cd3e77c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 35 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
envconfig/config.go
View file

@ -57,6 +57,8 @@ var (
SchedSpread bool
// Set via OLLAMA_TMPDIR in the environment
TmpDir string
// Set via OLLAMA_INTEL_GPU in the environment
IntelGpu bool
// Set via CUDA_VISIBLE_DEVICES in the environment
CudaVisibleDevices string
@ -101,6 +103,7 @@ func AsMap() map[string]EnvVar {
ret["ROCR_VISIBLE_DEVICES"] = EnvVar{"ROCR_VISIBLE_DEVICES", RocrVisibleDevices, "Set which AMD devices are visible"}
ret["GPU_DEVICE_ORDINAL"] = EnvVar{"GPU_DEVICE_ORDINAL", GpuDeviceOrdinal, "Set which AMD devices are visible"}
ret["HSA_OVERRIDE_GFX_VERSION"] = EnvVar{"HSA_OVERRIDE_GFX_VERSION", HsaOverrideGfxVersion, "Override the gfx used for all detected AMD GPUs"}
ret["OLLAMA_INTEL_GPU"] = EnvVar{"OLLAMA_INTEL_GPU", IntelGpu, "Enable experimental Intel GPU detection"}
}
return ret
}
@ -276,6 +279,10 @@ func LoadConfig() {
slog.Error("invalid setting", "OLLAMA_HOST", Host, "error", err, "using default port", Host.Port)
}
if set, err := strconv.ParseBool(clean("OLLAMA_INTEL_GPU")); err == nil {
IntelGpu = set
}
CudaVisibleDevices = clean("CUDA_VISIBLE_DEVICES")
HipVisibleDevices = clean("HIP_VISIBLE_DEVICES")
RocrVisibleDevices = clean("ROCR_VISIBLE_DEVICES")

6
gpu/gpu.go
View file

@ -280,8 +280,9 @@ func GetGPUInfo() GpuInfoList {
}
// Intel
if envconfig.IntelGpu {
oHandles = initOneAPIHandles()
for d := 0; oHandles.oneapi != nil && d < int(oHandles.oneapi.num_drivers); d++ {
for d := range oHandles.oneapi.num_drivers {
if oHandles.oneapi == nil {
// shouldn't happen
slog.Warn("nil oneapi handle with driver count", "count", int(oHandles.oneapi.num_drivers))
@ -293,7 +294,7 @@ func GetGPUInfo() GpuInfoList {
GpuInfo: GpuInfo{
Library: "oneapi",
},
driverIndex: d,
driverIndex: int(d),
gpuIndex: int(i),
}
// TODO - split bootstrapping from updating free memory
@ -309,6 +310,7 @@ func GetGPUInfo() GpuInfoList {
oneapiGPUs = append(oneapiGPUs, gpuInfo)
}
}
}
rocmGPUs = AMDGetGPUInfo()
bootstrapped = true