gpu.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670
  1. //go:build linux || windows
  2. package gpu
  3. /*
  4. #cgo linux LDFLAGS: -lrt -lpthread -ldl -lstdc++ -lm
  5. #cgo windows LDFLAGS: -lpthread
  6. #include "gpu_info.h"
  7. */
  8. import "C"
  9. import (
  10. "fmt"
  11. "log/slog"
  12. "os"
  13. "path/filepath"
  14. "runtime"
  15. "strings"
  16. "sync"
  17. "unsafe"
  18. "github.com/ollama/ollama/envconfig"
  19. "github.com/ollama/ollama/format"
  20. )
  21. type cudaHandles struct {
  22. deviceCount int
  23. cudart *C.cudart_handle_t
  24. nvcuda *C.nvcuda_handle_t
  25. nvml *C.nvml_handle_t
  26. }
  27. type oneapiHandles struct {
  28. oneapi *C.oneapi_handle_t
  29. deviceCount int
  30. }
  31. const (
  32. cudaMinimumMemory = 457 * format.MebiByte
  33. rocmMinimumMemory = 457 * format.MebiByte
  34. // TODO OneAPI minimum memory
  35. )
  36. var (
  37. gpuMutex sync.Mutex
  38. bootstrapped bool
  39. cpuCapability CPUCapability
  40. cpus []CPUInfo
  41. cudaGPUs []CudaGPUInfo
  42. nvcudaLibPath string
  43. cudartLibPath string
  44. oneapiLibPath string
  45. nvmlLibPath string
  46. rocmGPUs []RocmGPUInfo
  47. oneapiGPUs []OneapiGPUInfo
  48. )
  49. // With our current CUDA compile flags, older than 5.0 will not work properly
  50. var CudaComputeMin = [2]C.int{5, 0}
  51. var RocmComputeMin = 9
  52. // TODO find a better way to detect iGPU instead of minimum memory
  53. const IGPUMemLimit = 1 * format.GibiByte // 512G is what they typically report, so anything less than 1G must be iGPU
  54. // Note: gpuMutex must already be held
  55. func initCudaHandles() *cudaHandles {
  56. // TODO - if the ollama build is CPU only, don't do these checks as they're irrelevant and confusing
  57. cHandles := &cudaHandles{}
  58. // Short Circuit if we already know which library to use
  59. if nvmlLibPath != "" {
  60. cHandles.nvml, _ = LoadNVMLMgmt([]string{nvmlLibPath})
  61. return cHandles
  62. }
  63. if nvcudaLibPath != "" {
  64. cHandles.deviceCount, cHandles.nvcuda, _ = LoadNVCUDAMgmt([]string{nvcudaLibPath})
  65. return cHandles
  66. }
  67. if cudartLibPath != "" {
  68. cHandles.deviceCount, cHandles.cudart, _ = LoadCUDARTMgmt([]string{cudartLibPath})
  69. return cHandles
  70. }
  71. slog.Debug("searching for GPU discovery libraries for NVIDIA")
  72. var cudartMgmtPatterns []string
  73. // Aligned with driver, we can't carry as payloads
  74. nvcudaMgmtPatterns := NvcudaGlobs
  75. if runtime.GOOS == "windows" {
  76. localAppData := os.Getenv("LOCALAPPDATA")
  77. cudartMgmtPatterns = []string{filepath.Join(localAppData, "Programs", "Ollama", CudartMgmtName)}
  78. }
  79. libDir := LibraryDir()
  80. if libDir != "" {
  81. cudartMgmtPatterns = []string{filepath.Join(libDir, CudartMgmtName)}
  82. }
  83. cudartMgmtPatterns = append(cudartMgmtPatterns, CudartGlobs...)
  84. if len(NvmlGlobs) > 0 {
  85. nvmlLibPaths := FindGPULibs(NvmlMgmtName, NvmlGlobs)
  86. if len(nvmlLibPaths) > 0 {
  87. nvml, libPath := LoadNVMLMgmt(nvmlLibPaths)
  88. if nvml != nil {
  89. slog.Debug("nvidia-ml loaded", "library", libPath)
  90. cHandles.nvml = nvml
  91. nvmlLibPath = libPath
  92. }
  93. }
  94. }
  95. nvcudaLibPaths := FindGPULibs(NvcudaMgmtName, nvcudaMgmtPatterns)
  96. if len(nvcudaLibPaths) > 0 {
  97. deviceCount, nvcuda, libPath := LoadNVCUDAMgmt(nvcudaLibPaths)
  98. if nvcuda != nil {
  99. slog.Debug("detected GPUs", "count", deviceCount, "library", libPath)
  100. cHandles.nvcuda = nvcuda
  101. cHandles.deviceCount = deviceCount
  102. nvcudaLibPath = libPath
  103. return cHandles
  104. }
  105. }
  106. cudartLibPaths := FindGPULibs(CudartMgmtName, cudartMgmtPatterns)
  107. if len(cudartLibPaths) > 0 {
  108. deviceCount, cudart, libPath := LoadCUDARTMgmt(cudartLibPaths)
  109. if cudart != nil {
  110. slog.Debug("detected GPUs", "library", libPath, "count", deviceCount)
  111. cHandles.cudart = cudart
  112. cHandles.deviceCount = deviceCount
  113. cudartLibPath = libPath
  114. return cHandles
  115. }
  116. }
  117. return cHandles
  118. }
  119. // Note: gpuMutex must already be held
  120. func initOneAPIHandles() *oneapiHandles {
  121. oHandles := &oneapiHandles{}
  122. // Short Circuit if we already know which library to use
  123. if oneapiLibPath != "" {
  124. oHandles.deviceCount, oHandles.oneapi, _ = LoadOneapiMgmt([]string{oneapiLibPath})
  125. return oHandles
  126. }
  127. oneapiLibPaths := FindGPULibs(OneapiMgmtName, OneapiGlobs)
  128. if len(oneapiLibPaths) > 0 {
  129. oHandles.deviceCount, oHandles.oneapi, oneapiLibPath = LoadOneapiMgmt(oneapiLibPaths)
  130. }
  131. return oHandles
  132. }
  133. func GetCPUInfo() GpuInfoList {
  134. gpuMutex.Lock()
  135. if !bootstrapped {
  136. gpuMutex.Unlock()
  137. GetGPUInfo()
  138. } else {
  139. gpuMutex.Unlock()
  140. }
  141. return GpuInfoList{cpus[0].GpuInfo}
  142. }
  143. func GetGPUInfo() GpuInfoList {
  144. // TODO - consider exploring lspci (and equivalent on windows) to check for
  145. // GPUs so we can report warnings if we see Nvidia/AMD but fail to load the libraries
  146. gpuMutex.Lock()
  147. defer gpuMutex.Unlock()
  148. needRefresh := true
  149. var cHandles *cudaHandles
  150. var oHandles *oneapiHandles
  151. defer func() {
  152. if cHandles != nil {
  153. if cHandles.cudart != nil {
  154. C.cudart_release(*cHandles.cudart)
  155. }
  156. if cHandles.nvcuda != nil {
  157. C.nvcuda_release(*cHandles.nvcuda)
  158. }
  159. if cHandles.nvml != nil {
  160. C.nvml_release(*cHandles.nvml)
  161. }
  162. }
  163. if oHandles != nil {
  164. if oHandles.oneapi != nil {
  165. // TODO - is this needed?
  166. C.oneapi_release(*oHandles.oneapi)
  167. }
  168. }
  169. }()
  170. if !bootstrapped {
  171. slog.Info("looking for compatible GPUs")
  172. needRefresh = false
  173. cpuCapability = GetCPUCapability()
  174. var memInfo C.mem_info_t
  175. mem, err := GetCPUMem()
  176. if err != nil {
  177. slog.Warn("error looking up system memory", "error", err)
  178. }
  179. cpus = []CPUInfo{
  180. {
  181. GpuInfo: GpuInfo{
  182. memInfo: mem,
  183. Library: "cpu",
  184. Variant: cpuCapability.String(),
  185. ID: "0",
  186. },
  187. },
  188. }
  189. // Fallback to CPU mode if we're lacking required vector extensions on x86
  190. if cpuCapability < GPURunnerCPUCapability && runtime.GOARCH == "amd64" {
  191. slog.Warn("CPU does not have minimum vector extensions, GPU inference disabled", "required", GPURunnerCPUCapability, "detected", cpuCapability)
  192. bootstrapped = true
  193. // No need to do any GPU discovery, since we can't run on them
  194. return GpuInfoList{cpus[0].GpuInfo}
  195. }
  196. depPath := LibraryDir()
  197. // Load ALL libraries
  198. cHandles = initCudaHandles()
  199. // NVIDIA
  200. for i := range cHandles.deviceCount {
  201. if cHandles.cudart != nil || cHandles.nvcuda != nil {
  202. gpuInfo := CudaGPUInfo{
  203. GpuInfo: GpuInfo{
  204. Library: "cuda",
  205. },
  206. index: i,
  207. }
  208. var driverMajor int
  209. var driverMinor int
  210. if cHandles.cudart != nil {
  211. C.cudart_bootstrap(*cHandles.cudart, C.int(i), &memInfo)
  212. } else {
  213. C.nvcuda_bootstrap(*cHandles.nvcuda, C.int(i), &memInfo)
  214. driverMajor = int(cHandles.nvcuda.driver_major)
  215. driverMinor = int(cHandles.nvcuda.driver_minor)
  216. }
  217. if memInfo.err != nil {
  218. slog.Info("error looking up nvidia GPU memory", "error", C.GoString(memInfo.err))
  219. C.free(unsafe.Pointer(memInfo.err))
  220. continue
  221. }
  222. if memInfo.major < CudaComputeMin[0] || (memInfo.major == CudaComputeMin[0] && memInfo.minor < CudaComputeMin[1]) {
  223. slog.Info(fmt.Sprintf("[%d] CUDA GPU is too old. Compute Capability detected: %d.%d", i, memInfo.major, memInfo.minor))
  224. continue
  225. }
  226. gpuInfo.TotalMemory = uint64(memInfo.total)
  227. gpuInfo.FreeMemory = uint64(memInfo.free)
  228. gpuInfo.ID = C.GoString(&memInfo.gpu_id[0])
  229. gpuInfo.Compute = fmt.Sprintf("%d.%d", memInfo.major, memInfo.minor)
  230. gpuInfo.computeMajor = int(memInfo.major)
  231. gpuInfo.computeMinor = int(memInfo.minor)
  232. gpuInfo.MinimumMemory = cudaMinimumMemory
  233. gpuInfo.DriverMajor = driverMajor
  234. gpuInfo.DriverMinor = driverMinor
  235. variant := cudaVariant(gpuInfo)
  236. if depPath != "" {
  237. gpuInfo.DependencyPath = depPath
  238. // Check for variant specific directory
  239. if variant != "" {
  240. if _, err := os.Stat(filepath.Join(depPath, "cuda_"+variant)); err == nil {
  241. gpuInfo.DependencyPath = filepath.Join(depPath, "cuda_"+variant)
  242. }
  243. }
  244. }
  245. gpuInfo.Name = C.GoString(&memInfo.gpu_name[0])
  246. gpuInfo.Variant = variant
  247. // query the management library as well so we can record any skew between the two
  248. // which represents overhead on the GPU we must set aside on subsequent updates
  249. if cHandles.nvml != nil {
  250. C.nvml_get_free(*cHandles.nvml, C.int(gpuInfo.index), &memInfo.free, &memInfo.total, &memInfo.used)
  251. if memInfo.err != nil {
  252. slog.Warn("error looking up nvidia GPU memory", "error", C.GoString(memInfo.err))
  253. C.free(unsafe.Pointer(memInfo.err))
  254. } else {
  255. if memInfo.free != 0 && uint64(memInfo.free) > gpuInfo.FreeMemory {
  256. gpuInfo.OSOverhead = uint64(memInfo.free) - gpuInfo.FreeMemory
  257. slog.Info("detected OS VRAM overhead",
  258. "id", gpuInfo.ID,
  259. "library", gpuInfo.Library,
  260. "compute", gpuInfo.Compute,
  261. "driver", fmt.Sprintf("%d.%d", gpuInfo.DriverMajor, gpuInfo.DriverMinor),
  262. "name", gpuInfo.Name,
  263. "overhead", format.HumanBytes2(gpuInfo.OSOverhead),
  264. )
  265. }
  266. }
  267. }
  268. // TODO potentially sort on our own algorithm instead of what the underlying GPU library does...
  269. cudaGPUs = append(cudaGPUs, gpuInfo)
  270. }
  271. }
  272. // Intel
  273. if envconfig.IntelGPU() {
  274. oHandles = initOneAPIHandles()
  275. if oHandles != nil && oHandles.oneapi != nil {
  276. for d := range oHandles.oneapi.num_drivers {
  277. if oHandles.oneapi == nil {
  278. // shouldn't happen
  279. slog.Warn("nil oneapi handle with driver count", "count", int(oHandles.oneapi.num_drivers))
  280. continue
  281. }
  282. devCount := C.oneapi_get_device_count(*oHandles.oneapi, C.int(d))
  283. for i := range devCount {
  284. gpuInfo := OneapiGPUInfo{
  285. GpuInfo: GpuInfo{
  286. Library: "oneapi",
  287. },
  288. driverIndex: int(d),
  289. gpuIndex: int(i),
  290. }
  291. // TODO - split bootstrapping from updating free memory
  292. C.oneapi_check_vram(*oHandles.oneapi, C.int(d), i, &memInfo)
  293. // TODO - convert this to MinimumMemory based on testing...
  294. var totalFreeMem float64 = float64(memInfo.free) * 0.95 // work-around: leave some reserve vram for mkl lib used in ggml-sycl backend.
  295. memInfo.free = C.uint64_t(totalFreeMem)
  296. gpuInfo.TotalMemory = uint64(memInfo.total)
  297. gpuInfo.FreeMemory = uint64(memInfo.free)
  298. gpuInfo.ID = C.GoString(&memInfo.gpu_id[0])
  299. gpuInfo.Name = C.GoString(&memInfo.gpu_name[0])
  300. gpuInfo.DependencyPath = depPath
  301. oneapiGPUs = append(oneapiGPUs, gpuInfo)
  302. }
  303. }
  304. }
  305. }
  306. rocmGPUs = AMDGetGPUInfo()
  307. bootstrapped = true
  308. if len(cudaGPUs) == 0 && len(rocmGPUs) == 0 && len(oneapiGPUs) == 0 {
  309. slog.Info("no compatible GPUs were discovered")
  310. }
  311. }
  312. // For detected GPUs, load library if not loaded
  313. // Refresh free memory usage
  314. if needRefresh {
  315. mem, err := GetCPUMem()
  316. if err != nil {
  317. slog.Warn("error looking up system memory", "error", err)
  318. } else {
  319. slog.Debug("updating system memory data",
  320. slog.Group(
  321. "before",
  322. "total", format.HumanBytes2(cpus[0].TotalMemory),
  323. "free", format.HumanBytes2(cpus[0].FreeMemory),
  324. "free_swap", format.HumanBytes2(cpus[0].FreeSwap),
  325. ),
  326. slog.Group(
  327. "now",
  328. "total", format.HumanBytes2(mem.TotalMemory),
  329. "free", format.HumanBytes2(mem.FreeMemory),
  330. "free_swap", format.HumanBytes2(mem.FreeSwap),
  331. ),
  332. )
  333. cpus[0].FreeMemory = mem.FreeMemory
  334. cpus[0].FreeSwap = mem.FreeSwap
  335. }
  336. var memInfo C.mem_info_t
  337. if cHandles == nil && len(cudaGPUs) > 0 {
  338. cHandles = initCudaHandles()
  339. }
  340. for i, gpu := range cudaGPUs {
  341. if cHandles.nvml != nil {
  342. C.nvml_get_free(*cHandles.nvml, C.int(gpu.index), &memInfo.free, &memInfo.total, &memInfo.used)
  343. } else if cHandles.cudart != nil {
  344. C.cudart_bootstrap(*cHandles.cudart, C.int(gpu.index), &memInfo)
  345. } else if cHandles.nvcuda != nil {
  346. C.nvcuda_get_free(*cHandles.nvcuda, C.int(gpu.index), &memInfo.free, &memInfo.total)
  347. memInfo.used = memInfo.total - memInfo.free
  348. } else {
  349. // shouldn't happen
  350. slog.Warn("no valid cuda library loaded to refresh vram usage")
  351. break
  352. }
  353. if memInfo.err != nil {
  354. slog.Warn("error looking up nvidia GPU memory", "error", C.GoString(memInfo.err))
  355. C.free(unsafe.Pointer(memInfo.err))
  356. continue
  357. }
  358. if memInfo.free == 0 {
  359. slog.Warn("error looking up nvidia GPU memory")
  360. continue
  361. }
  362. if cHandles.nvml != nil && gpu.OSOverhead > 0 {
  363. // When using the management library update based on recorded overhead
  364. memInfo.free -= C.uint64_t(gpu.OSOverhead)
  365. }
  366. slog.Debug("updating cuda memory data",
  367. "gpu", gpu.ID,
  368. "name", gpu.Name,
  369. "overhead", format.HumanBytes2(gpu.OSOverhead),
  370. slog.Group(
  371. "before",
  372. "total", format.HumanBytes2(gpu.TotalMemory),
  373. "free", format.HumanBytes2(gpu.FreeMemory),
  374. ),
  375. slog.Group(
  376. "now",
  377. "total", format.HumanBytes2(uint64(memInfo.total)),
  378. "free", format.HumanBytes2(uint64(memInfo.free)),
  379. "used", format.HumanBytes2(uint64(memInfo.used)),
  380. ),
  381. )
  382. cudaGPUs[i].FreeMemory = uint64(memInfo.free)
  383. }
  384. if oHandles == nil && len(oneapiGPUs) > 0 {
  385. oHandles = initOneAPIHandles()
  386. }
  387. for i, gpu := range oneapiGPUs {
  388. if oHandles.oneapi == nil {
  389. // shouldn't happen
  390. slog.Warn("nil oneapi handle with device count", "count", oHandles.deviceCount)
  391. continue
  392. }
  393. C.oneapi_check_vram(*oHandles.oneapi, C.int(gpu.driverIndex), C.int(gpu.gpuIndex), &memInfo)
  394. // TODO - convert this to MinimumMemory based on testing...
  395. var totalFreeMem float64 = float64(memInfo.free) * 0.95 // work-around: leave some reserve vram for mkl lib used in ggml-sycl backend.
  396. memInfo.free = C.uint64_t(totalFreeMem)
  397. oneapiGPUs[i].FreeMemory = uint64(memInfo.free)
  398. }
  399. err = RocmGPUInfoList(rocmGPUs).RefreshFreeMemory()
  400. if err != nil {
  401. slog.Debug("problem refreshing ROCm free memory", "error", err)
  402. }
  403. }
  404. resp := []GpuInfo{}
  405. for _, gpu := range cudaGPUs {
  406. resp = append(resp, gpu.GpuInfo)
  407. }
  408. for _, gpu := range rocmGPUs {
  409. resp = append(resp, gpu.GpuInfo)
  410. }
  411. for _, gpu := range oneapiGPUs {
  412. resp = append(resp, gpu.GpuInfo)
  413. }
  414. if len(resp) == 0 {
  415. resp = append(resp, cpus[0].GpuInfo)
  416. }
  417. return resp
  418. }
  419. func FindGPULibs(baseLibName string, defaultPatterns []string) []string {
  420. // Multiple GPU libraries may exist, and some may not work, so keep trying until we exhaust them
  421. var ldPaths []string
  422. gpuLibPaths := []string{}
  423. slog.Debug("Searching for GPU library", "name", baseLibName)
  424. // Start with our bundled libraries
  425. patterns := []string{filepath.Join(LibraryDir(), baseLibName)}
  426. switch runtime.GOOS {
  427. case "windows":
  428. ldPaths = strings.Split(os.Getenv("PATH"), ";")
  429. case "linux":
  430. ldPaths = strings.Split(os.Getenv("LD_LIBRARY_PATH"), ":")
  431. default:
  432. return gpuLibPaths
  433. }
  434. // Then with whatever we find in the PATH/LD_LIBRARY_PATH
  435. for _, ldPath := range ldPaths {
  436. d, err := filepath.Abs(ldPath)
  437. if err != nil {
  438. continue
  439. }
  440. patterns = append(patterns, filepath.Join(d, baseLibName))
  441. }
  442. patterns = append(patterns, defaultPatterns...)
  443. slog.Debug("gpu library search", "globs", patterns)
  444. for _, pattern := range patterns {
  445. // Nvidia PhysX known to return bogus results
  446. if strings.Contains(pattern, "PhysX") {
  447. slog.Debug("skipping PhysX cuda library path", "path", pattern)
  448. continue
  449. }
  450. // Ignore glob discovery errors
  451. matches, _ := filepath.Glob(pattern)
  452. for _, match := range matches {
  453. // Resolve any links so we don't try the same lib multiple times
  454. // and weed out any dups across globs
  455. libPath := match
  456. tmp := match
  457. var err error
  458. for ; err == nil; tmp, err = os.Readlink(libPath) {
  459. if !filepath.IsAbs(tmp) {
  460. tmp = filepath.Join(filepath.Dir(libPath), tmp)
  461. }
  462. libPath = tmp
  463. }
  464. new := true
  465. for _, cmp := range gpuLibPaths {
  466. if cmp == libPath {
  467. new = false
  468. break
  469. }
  470. }
  471. if new {
  472. gpuLibPaths = append(gpuLibPaths, libPath)
  473. }
  474. }
  475. }
  476. slog.Debug("discovered GPU libraries", "paths", gpuLibPaths)
  477. return gpuLibPaths
  478. }
  479. func LoadCUDARTMgmt(cudartLibPaths []string) (int, *C.cudart_handle_t, string) {
  480. var resp C.cudart_init_resp_t
  481. resp.ch.verbose = getVerboseState()
  482. for _, libPath := range cudartLibPaths {
  483. lib := C.CString(libPath)
  484. defer C.free(unsafe.Pointer(lib))
  485. C.cudart_init(lib, &resp)
  486. if resp.err != nil {
  487. slog.Debug("Unable to load cudart", "library", libPath, "error", C.GoString(resp.err))
  488. C.free(unsafe.Pointer(resp.err))
  489. } else {
  490. return int(resp.num_devices), &resp.ch, libPath
  491. }
  492. }
  493. return 0, nil, ""
  494. }
  495. func LoadNVCUDAMgmt(nvcudaLibPaths []string) (int, *C.nvcuda_handle_t, string) {
  496. var resp C.nvcuda_init_resp_t
  497. resp.ch.verbose = getVerboseState()
  498. for _, libPath := range nvcudaLibPaths {
  499. lib := C.CString(libPath)
  500. defer C.free(unsafe.Pointer(lib))
  501. C.nvcuda_init(lib, &resp)
  502. if resp.err != nil {
  503. // Decide what log level based on the type of error message to help users understand why
  504. msg := C.GoString(resp.err)
  505. switch resp.cudaErr {
  506. case C.CUDA_ERROR_INSUFFICIENT_DRIVER, C.CUDA_ERROR_SYSTEM_DRIVER_MISMATCH:
  507. slog.Warn("version mismatch between driver and cuda driver library - reboot or upgrade may be required", "library", libPath, "error", msg)
  508. case C.CUDA_ERROR_NO_DEVICE:
  509. slog.Info("no nvidia devices detected", "library", libPath)
  510. case C.CUDA_ERROR_UNKNOWN:
  511. slog.Warn("unknown error initializing cuda driver library", "library", libPath, "error", msg)
  512. slog.Warn("see https://github.com/ollama/ollama/blob/main/docs/troubleshooting.md for more information")
  513. default:
  514. if strings.Contains(msg, "wrong ELF class") {
  515. slog.Debug("skipping 32bit library", "library", libPath)
  516. } else {
  517. slog.Info("unable to load cuda driver library", "library", libPath, "error", msg)
  518. }
  519. }
  520. C.free(unsafe.Pointer(resp.err))
  521. } else {
  522. return int(resp.num_devices), &resp.ch, libPath
  523. }
  524. }
  525. return 0, nil, ""
  526. }
  527. func LoadNVMLMgmt(nvmlLibPaths []string) (*C.nvml_handle_t, string) {
  528. var resp C.nvml_init_resp_t
  529. resp.ch.verbose = getVerboseState()
  530. for _, libPath := range nvmlLibPaths {
  531. lib := C.CString(libPath)
  532. defer C.free(unsafe.Pointer(lib))
  533. C.nvml_init(lib, &resp)
  534. if resp.err != nil {
  535. slog.Info(fmt.Sprintf("Unable to load NVML management library %s: %s", libPath, C.GoString(resp.err)))
  536. C.free(unsafe.Pointer(resp.err))
  537. } else {
  538. return &resp.ch, libPath
  539. }
  540. }
  541. return nil, ""
  542. }
  543. func LoadOneapiMgmt(oneapiLibPaths []string) (int, *C.oneapi_handle_t, string) {
  544. var resp C.oneapi_init_resp_t
  545. num_devices := 0
  546. resp.oh.verbose = getVerboseState()
  547. for _, libPath := range oneapiLibPaths {
  548. lib := C.CString(libPath)
  549. defer C.free(unsafe.Pointer(lib))
  550. C.oneapi_init(lib, &resp)
  551. if resp.err != nil {
  552. slog.Debug("Unable to load oneAPI management library", "library", libPath, "error", C.GoString(resp.err))
  553. C.free(unsafe.Pointer(resp.err))
  554. } else {
  555. for i := range resp.oh.num_drivers {
  556. num_devices += int(C.oneapi_get_device_count(resp.oh, C.int(i)))
  557. }
  558. return num_devices, &resp.oh, libPath
  559. }
  560. }
  561. return 0, nil, ""
  562. }
  563. func getVerboseState() C.uint16_t {
  564. if envconfig.Debug() {
  565. return C.uint16_t(1)
  566. }
  567. return C.uint16_t(0)
  568. }
  569. // Given the list of GPUs this instantiation is targeted for,
  570. // figure out the visible devices environment variable
  571. //
  572. // If different libraries are detected, the first one is what we use
  573. func (l GpuInfoList) GetVisibleDevicesEnv() (string, string) {
  574. if len(l) == 0 {
  575. return "", ""
  576. }
  577. switch l[0].Library {
  578. case "cuda":
  579. return cudaGetVisibleDevicesEnv(l)
  580. case "rocm":
  581. return rocmGetVisibleDevicesEnv(l)
  582. case "oneapi":
  583. return oneapiGetVisibleDevicesEnv(l)
  584. default:
  585. slog.Debug("no filter required for library " + l[0].Library)
  586. return "", ""
  587. }
  588. }
  589. func LibraryDir() string {
  590. // On Windows/linux we bundle the dependencies at the same level as the executable
  591. appExe, err := os.Executable()
  592. if err != nil {
  593. slog.Warn("failed to lookup executable path", "error", err)
  594. }
  595. cwd, err := os.Getwd()
  596. if err != nil {
  597. slog.Warn("failed to lookup working directory", "error", err)
  598. }
  599. // Scan for any of our dependeices, and pick first match
  600. for _, root := range []string{filepath.Dir(appExe), filepath.Join(filepath.Dir(appExe), envconfig.LibRelativeToExe()), cwd} {
  601. libDep := filepath.Join("lib", "ollama")
  602. if _, err := os.Stat(filepath.Join(root, libDep)); err == nil {
  603. return filepath.Join(root, libDep)
  604. }
  605. // Developer mode, local build
  606. if _, err := os.Stat(filepath.Join(root, runtime.GOOS+"-"+runtime.GOARCH, libDep)); err == nil {
  607. return filepath.Join(root, runtime.GOOS+"-"+runtime.GOARCH, libDep)
  608. }
  609. if _, err := os.Stat(filepath.Join(root, "dist", runtime.GOOS+"-"+runtime.GOARCH, libDep)); err == nil {
  610. return filepath.Join(root, "dist", runtime.GOOS+"-"+runtime.GOARCH, libDep)
  611. }
  612. }
  613. slog.Warn("unable to locate gpu dependency libraries")
  614. return ""
  615. }