d88c527be3
This reduces the built-in linux version to not use any vector extensions which enables the resulting builds to run under Rosetta on MacOS in Docker. Then at runtime it checks for the actual CPU vector extensions and loads the best CPU library available
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 TestGetShims(t *testing.T) {
|
|
availableShims = map[string]string{
|
|
"cpu": "X_cpu",
|
|
}
|
|
assert.Equal(t, false, rocmShimPresent())
|
|
res := getShims(gpu.GpuInfo{Library: "cpu"})
|
|
assert.Len(t, res, 1)
|
|
assert.Equal(t, availableShims["cpu"], res[0])
|
|
|
|
availableShims = map[string]string{
|
|
"rocm_v5": "X_rocm_v5",
|
|
"rocm_v6": "X_rocm_v6",
|
|
"cpu": "X_cpu",
|
|
}
|
|
assert.Equal(t, true, rocmShimPresent())
|
|
res = getShims(gpu.GpuInfo{Library: "rocm"})
|
|
assert.Len(t, res, 3)
|
|
assert.Equal(t, availableShims["rocm_v5"], res[0])
|
|
assert.Equal(t, availableShims["rocm_v6"], res[1])
|
|
assert.Equal(t, availableShims["cpu"], res[2])
|
|
|
|
res = getShims(gpu.GpuInfo{Library: "rocm", Variant: "v6"})
|
|
assert.Len(t, res, 3)
|
|
assert.Equal(t, availableShims["rocm_v6"], res[0])
|
|
assert.Equal(t, availableShims["rocm_v5"], res[1])
|
|
assert.Equal(t, availableShims["cpu"], res[2])
|
|
|
|
res = getShims(gpu.GpuInfo{Library: "cuda"})
|
|
assert.Len(t, res, 1)
|
|
assert.Equal(t, availableShims["cpu"], res[0])
|
|
|
|
res = getShims(gpu.GpuInfo{Library: "default"})
|
|
assert.Len(t, res, 1)
|
|
assert.Equal(t, "default", res[0])
|
|
|
|
availableShims = map[string]string{
|
|
"rocm": "X_rocm_v5",
|
|
"cpu": "X_cpu",
|
|
}
|
|
assert.Equal(t, true, rocmShimPresent())
|
|
res = getShims(gpu.GpuInfo{Library: "rocm", Variant: "v6"})
|
|
assert.Len(t, res, 2)
|
|
assert.Equal(t, availableShims["rocm"], res[0])
|
|
assert.Equal(t, availableShims["cpu"], res[1])
|
|
}
|