[AIE-151] native: support dynamic detection of OpenCL

Signed-off-by: Jacob Howard <jacob.howard@docker.com>
This commit is contained in:
Jacob Howard 2025-04-28 12:13:37 -06:00
parent 4c3ffbfa53
commit 3d8c73c355
No known key found for this signature in database
GPG Key ID: 3E8B8F7FEB46FC66
2 changed files with 33 additions and 6 deletions

View File

@ -14,7 +14,7 @@ func (l *llamaCpp) ensureLatestLlamaCpp(ctx context.Context, log logging.Logger,
llamaCppPath, vendoredServerStoragePath string,
) error {
nvGPUInfoBin := filepath.Join(vendoredServerStoragePath, "com.docker.nv-gpu-info.exe")
var canUseCUDA11 bool
var canUseCUDA11, canUseOpenCL bool
var err error
ShouldUseGPUVariantLock.Lock()
defer ShouldUseGPUVariantLock.Unlock()
@ -25,15 +25,19 @@ func (l *llamaCpp) ensureLatestLlamaCpp(ctx context.Context, log logging.Logger,
l.status = fmt.Sprintf("failed to check CUDA 11 capability: %v", err)
return fmt.Errorf("failed to check CUDA 11 capability: %w", err)
}
} else if runtime.GOARCH == "arm64" {
canUseOpenCL, err = hasOpenCL()
if err != nil {
l.status = fmt.Sprintf("failed to check OpenCL capability: %v", err)
return fmt.Errorf("failed to check OpenCL capability: %w", err)
}
}
}
desiredVersion := "latest"
desiredVariant := "cpu"
if canUseCUDA11 {
desiredVariant = "cuda"
}
// TODO(p1-0tr): we should auto-detect if we can use opencl, but for now assume that we can
if runtime.GOARCH == "arm64" {
} else if canUseOpenCL {
desiredVariant = "opencl"
}
l.status = fmt.Sprintf("looking for updates for %s variant", desiredVariant)

View File

@ -3,10 +3,12 @@ package llamacpp
import (
"bufio"
"context"
"errors"
"fmt"
"os/exec"
"strconv"
"strings"
"syscall"
"github.com/jaypipes/ghw"
)
@ -56,6 +58,27 @@ func hasCUDA11CapableGPU(ctx context.Context, nvGPUInfoBin string) (bool, error)
return false, nil
}
func CanUseGPU(ctx context.Context, nvGPUInfoBin string) (bool, error) {
return hasCUDA11CapableGPU(ctx, nvGPUInfoBin)
func hasOpenCL() (bool, error) {
opencl, err := syscall.LoadLibrary("OpenCL.dll")
if err != nil {
if errors.Is(err, syscall.ERROR_MOD_NOT_FOUND) {
return false, nil
}
return false, fmt.Errorf("unable to load OpenCL DLL: %w", err)
}
// We could perform additional platform and device version checks here (if
// we scaffold out the relevant OpenCL API datatypes in Go), but since users
// can opt-out of GPU support, we can probably skip that and just let users
// disable it if things don't work. Alternatively, we could inspect the GPUs
// found by the ghw package, if it supports (e.g.) Adreno GPUs.
syscall.FreeLibrary(opencl)
return true, nil
}
func CanUseGPU(ctx context.Context, nvGPUInfoBin string) (bool, error) {
haveCUDA11GPU, err := hasCUDA11CapableGPU(ctx, nvGPUInfoBin)
if haveCUDA11GPU || err != nil {
return haveCUDA11GPU, err
}
return hasOpenCL()
}