55 lines
1.5 KiB
Go
55 lines
1.5 KiB
Go
|
package llm
|
||
|
|
|
||
|
|
import (
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/jmorganca/ollama/gpu"
|
||
|
|
"github.com/stretchr/testify/assert"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestGetDynLibs(t *testing.T) {
|
||
|
|
availableDynLibs = map[string]string{
|
||
|
|
"cpu": "X_cpu",
|
||
|
|
}
|
||
|
|
assert.Equal(t, false, rocmDynLibPresent())
|
||
|
|
res := getDynLibs(gpu.GpuInfo{Library: "cpu"})
|
||
|
|
assert.Len(t, res, 1)
|
||
|
|
assert.Equal(t, availableDynLibs["cpu"], res[0])
|
||
|
|
|
||
|
|
availableDynLibs = map[string]string{
|
||
|
|
"rocm_v5": "X_rocm_v5",
|
||
|
|
"rocm_v6": "X_rocm_v6",
|
||
|
|
"cpu": "X_cpu",
|
||
|
|
}
|
||
|
|
assert.Equal(t, true, rocmDynLibPresent())
|
||
|
|
res = getDynLibs(gpu.GpuInfo{Library: "rocm"})
|
||
|
|
assert.Len(t, res, 3)
|
||
|
|
assert.Equal(t, availableDynLibs["rocm_v5"], res[0])
|
||
|
|
assert.Equal(t, availableDynLibs["rocm_v6"], res[1])
|
||
|
|
assert.Equal(t, availableDynLibs["cpu"], res[2])
|
||
|
|
|
||
|
|
res = getDynLibs(gpu.GpuInfo{Library: "rocm", Variant: "v6"})
|
||
|
|
assert.Len(t, res, 3)
|
||
|
|
assert.Equal(t, availableDynLibs["rocm_v6"], res[0])
|
||
|
|
assert.Equal(t, availableDynLibs["rocm_v5"], res[1])
|
||
|
|
assert.Equal(t, availableDynLibs["cpu"], res[2])
|
||
|
|
|
||
|
|
res = getDynLibs(gpu.GpuInfo{Library: "cuda"})
|
||
|
|
assert.Len(t, res, 1)
|
||
|
|
assert.Equal(t, availableDynLibs["cpu"], res[0])
|
||
|
|
|
||
|
|
res = getDynLibs(gpu.GpuInfo{Library: "default"})
|
||
|
|
assert.Len(t, res, 1)
|
||
|
|
assert.Equal(t, "default", res[0])
|
||
|
|
|
||
|
|
availableDynLibs = map[string]string{
|
||
|
|
"rocm": "X_rocm_v5",
|
||
|
|
"cpu": "X_cpu",
|
||
|
|
}
|
||
|
|
assert.Equal(t, true, rocmDynLibPresent())
|
||
|
|
res = getDynLibs(gpu.GpuInfo{Library: "rocm", Variant: "v6"})
|
||
|
|
assert.Len(t, res, 2)
|
||
|
|
assert.Equal(t, availableDynLibs["rocm"], res[0])
|
||
|
|
assert.Equal(t, availableDynLibs["cpu"], res[1])
|
||
|
|
}
|