Add ggml Ascend CANN backend compile and Ascend GPU detect

Makefile

@@ -34,6 +34,9 @@ ifneq ($(HIP_COMPILER),)
 endif
 endif
 
+ifeq ($(OLLAMA_SKIP_CANN_GENERATE),)
+	RUNNER_TARGETS += cann
+endif
 
 all: runners exe
 

discover/ascend_linux.go

@@ -0,0 +1,101 @@
+//go:build linux
+
+package discover
+
+/*
+#cgo linux LDFLAGS: -lrt -lpthread -ldl -lstdc++ -lm
+#cgo windows LDFLAGS: -lpthread
+
+#include "gpu_info_ascend.h"
+
+*/
+import "C"
+import (
+	"fmt"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"strings"
+	"unsafe"
+)
+
+var AscendLinuxGlobs = []string{
+	"/usr/local/Ascend/latest/aarch64-linux/lib64/libascendcl.so*",
+	"/usr/local/Ascend/ascend-toolkit/latest/aarch64-linux/lib64/libascendcl.so*",
+}
+
+var AscendMgmtName = "libascendcl.so"
+
+var (
+	ascendLibPath string
+)
+
+type ascendHandles struct {
+	ascend      *C.ascend_handle_t
+	deviceCount int
+}
+
+type AscendGPUInfo struct {
+	GpuInfo
+	index int //nolint:unused,nolintlint
+}
+type AscendGPUInfoList []AscendGPUInfo
+
+func initAscendHandles() *ascendHandles {
+	aHandles := &ascendHandles{}
+
+	// Short Circuit if we already know which library to use
+	if ascendLibPath != "" {
+		aHandles.deviceCount, aHandles.ascend, _ = LoadAscendMgmt([]string{ascendLibPath})
+		return aHandles
+	}
+
+	ascendToolkitHome := os.Getenv("ASCEND_TOOLKIT_HOME")
+	if ascendToolkitHome != "" {
+		AscendLinuxGlobs = append(AscendLinuxGlobs, filepath.Join(ascendToolkitHome, "/lib64/libascendcl.so*"))
+	}
+
+	ascendLibPaths := FindGPULibs(AscendMgmtName, AscendLinuxGlobs)
+	if len(ascendLibPaths) > 0 {
+		deviceCount, ascend, libPath := LoadAscendMgmt(ascendLibPaths)
+		if ascend != nil {
+			slog.Debug("detected GPUs", "count", deviceCount, "library", libPath)
+			aHandles.ascend = ascend
+			aHandles.deviceCount = deviceCount
+			ascendLibPath = libPath
+			return aHandles
+		}
+	}
+
+	return aHandles
+}
+
+func LoadAscendMgmt(ascendLibPath []string) (int, *C.ascend_handle_t, string) {
+	var resp C.ascend_init_resp_t
+	resp.ah.verbose = getVerboseState()
+	for _, libPath := range ascendLibPath {
+		lib := C.CString(libPath)
+		defer C.free(unsafe.Pointer(lib))
+		C.ascend_init(lib, &resp)
+		if resp.err != nil {
+			slog.Debug(fmt.Sprintf("Unable to load ascend management library %s: %s", libPath, C.GoString(resp.err)))
+			C.free(unsafe.Pointer(resp.err))
+		} else {
+			return int(resp.num_devices), &resp.ah, libPath
+		}
+	}
+	return 0, nil, ""
+}
+
+func ascendGetVisibleDevicesEnv(gpuInfo []GpuInfo) (string, string) {
+	ids := []string{}
+	for _, info := range gpuInfo {
+		if info.Library != "ascend" {
+			// TODO shouldn't happen if things are wired correctly...
+			slog.Debug("ascendGetVisibleDevicesEnv skipping over non-ascend device", "library", info.Library)
+			continue
+		}
+		ids = append(ids, info.ID)
+	}
+	return "ASCEND_RT_VISIBLE_DEVICES", strings.Join(ids, ",")
+}

discover/gpu.go

@@ -55,6 +55,7 @@ var (
 	nvmlLibPath   string
 	rocmGPUs      []RocmGPUInfo
 	oneapiGPUs    []OneapiGPUInfo
+	ascendGPUs    []AscendGPUInfo
 
 	// If any discovered GPUs are incompatible, report why
 	unsupportedGPUs []UnsupportedGPUInfo
@@ -202,6 +203,7 @@ func GetGPUInfo() GpuInfoList {
 	needRefresh := true
 	var cHandles *cudaHandles
 	var oHandles *oneapiHandles
+	var aHandles *ascendHandles
 	defer func() {
 		if cHandles != nil {
 			if cHandles.cudart != nil {
@@ -220,6 +222,9 @@ func GetGPUInfo() GpuInfoList {
 				C.oneapi_release(*oHandles.oneapi)
 			}
 		}
+		if aHandles.ascend != nil {
+			C.ascend_release(*aHandles.ascend)
+		}
 	}()
 
 	if !bootstrapped {
@@ -387,8 +392,31 @@ func GetGPUInfo() GpuInfoList {
 		if err != nil {
 			bootstrapErrors = append(bootstrapErrors, err)
 		}
+
+		// Then Ascend
+		aHandles = initAscendHandles()
+
+		for i := range aHandles.deviceCount {
+			if aHandles.ascend != nil {
+				gpuInfo := AscendGPUInfo{
+					GpuInfo: GpuInfo{
+						Library: "ascend",
+					},
+					index: i,
+				}
+				C.ascend_bootstrap(*aHandles.ascend, C.int(i), &memInfo)
+				gpuInfo.TotalMemory = uint64(memInfo.total)
+				gpuInfo.FreeMemory = uint64(memInfo.free)
+				gpuInfo.ID = C.GoString(&memInfo.gpu_id[0])
+				gpuInfo.Name = C.GoString(&memInfo.gpu_name[0])
+				ascendGPUs = append(ascendGPUs, gpuInfo)
+
+				slog.Info(fmt.Sprintf("[%d] Name:%s: FreeMemory:%d, TotalMemory: %d", gpuInfo.ID, gpuInfo.Name, gpuInfo.FreeMemory, gpuInfo.TotalMemory))
+			}
+		}
+
 		bootstrapped = true
-		if len(cudaGPUs) == 0 && len(rocmGPUs) == 0 && len(oneapiGPUs) == 0 {
+		if len(cudaGPUs) == 0 && len(rocmGPUs) == 0 && len(oneapiGPUs) == 0 && len(ascendGPUs) == 0 {
 			slog.Info("no compatible GPUs were discovered")
 		}
 
@@ -492,6 +520,19 @@ func GetGPUInfo() GpuInfoList {
 		if err != nil {
 			slog.Debug("problem refreshing ROCm free memory", "error", err)
 		}
+		if aHandles == nil && len(ascendGPUs) > 0 {
+			aHandles = initAscendHandles()
+		}
+
+		for i, gpu := range ascendGPUs {
+			if aHandles.ascend == nil {
+				// shouldn't happen
+				slog.Warn("nil ascend handle with device count", "count", aHandles.deviceCount)
+				continue
+			}
+			C.ascend_bootstrap(*aHandles.ascend, C.int(gpu.index), &memInfo)
+			ascendGPUs[i].FreeMemory = uint64(memInfo.free)
+		}
 	}
 
 	resp := []GpuInfo{}
@@ -504,6 +545,9 @@ func GetGPUInfo() GpuInfoList {
 	for _, gpu := range oneapiGPUs {
 		resp = append(resp, gpu.GpuInfo)
 	}
+	for _, gpu := range ascendGPUs {
+		resp = append(resp, gpu.GpuInfo)
+	}
 	if len(resp) == 0 {
 		resp = append(resp, cpus[0].GpuInfo)
 	}
@@ -558,6 +602,8 @@ func FindGPULibs(baseLibName string, defaultPatterns []string) []string {
 			var err error
 			for ; err == nil; tmp, err = os.Readlink(libPath) {
 				if !filepath.IsAbs(tmp) {
+					// Resolve possible Symlinks in libPath
+					libPath, _ = filepath.EvalSymlinks(libPath)
 					tmp = filepath.Join(filepath.Dir(libPath), tmp)
 				}
 				libPath = tmp
@@ -710,6 +756,8 @@ func (l GpuInfoList) GetVisibleDevicesEnv() (string, string) {
 		return rocmGetVisibleDevicesEnv(l)
 	case "oneapi":
 		return oneapiGetVisibleDevicesEnv(l)
+	case "ascend":
+		return ascendGetVisibleDevicesEnv(l)
 	default:
 		slog.Debug("no filter required for library " + l[0].Library)
 		return "", ""

discover/gpu_info.h

@@ -65,6 +65,7 @@ void cpu_check_ram(mem_info_t *resp);
 #include "gpu_info_nvcuda.h"
 #include "gpu_info_nvml.h"
 #include "gpu_info_oneapi.h"
+#include "gpu_info_ascend.h"
 
 #endif  // __GPU_INFO_H__
 #endif  // __APPLE__