diff --git a/_example/flat_index_map/flat_index_map.go b/_example/flat_index_map/flat_index_map.go new file mode 100644 index 0000000..3ba6f8b --- /dev/null +++ b/_example/flat_index_map/flat_index_map.go @@ -0,0 +1,36 @@ +package main + +import ( + "fmt" + "github.com/DataIntelligenceCrew/go-faiss" +) + +func main() { + dimension := 1 + dbSize := 5 + + index, err := faiss.NewIndexFlat(dimension, faiss.MetricL2) + if err != nil { + fmt.Println(err.Error()) + } + indexMap, err := faiss.NewIndexIDMap(index) + if err != nil { + fmt.Println(err.Error()) + } + xb := []float32{1,2,3,4,5} + ids := make([]int64, dbSize) + for i := 0; i < dbSize; i++ { + ids[i] = int64(i) + } + + err = indexMap.AddWithIDs(xb, ids) + if err != nil { + fmt.Println(err.Error()) + } + toFind := xb[dimension:2*dimension] + distances1, resultIds, err := indexMap.Search(toFind, 5) + fmt.Println(distances1, resultIds, err) + fmt.Println(resultIds[0] == ids[1]) + fmt.Println(distances1[0] == 0) + +} diff --git a/go.mod b/go.mod index 9eed116..ea8f2e9 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,5 @@ -module github.com/DataIntelligenceCrew/go-faiss +module github.com/Anyvisionltd/go-faiss go 1.14 + +require github.com/stretchr/testify v1.7.0 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..b380ae4 --- /dev/null +++ b/go.sum @@ -0,0 +1,10 @@ +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/gpu_bindings.go b/gpu_bindings.go new file mode 100644 index 0000000..2e17ab7 --- /dev/null +++ b/gpu_bindings.go @@ -0,0 +1,115 @@ +//go:build gpu +// +build gpu + +package faiss + +/* + +#include +#include +#include +#include +*/ +import "C" +import ( + "errors" + "unsafe" +) + +func TransferToGpu(index Index) (Index, error) { + var gpuResource *C.FaissStandardGpuResources + var gpuIndex *C.FaissGpuIndex + c := C.faiss_StandardGpuResources_new(&gpuResource) + if c != 0 { + return nil, errors.New("error on init gpu %v") + } + + exitCode := C.faiss_index_cpu_to_gpu(gpuResource, 0, index.cPtr(), &gpuIndex) + + if exitCode != 0 { + return nil, errors.New("error transferring to gpu") + } + + var gpuResources []*C.FaissStandardGpuResources + gpuResources = append(gpuResources, gpuResource) + + return &faissIndex{idx: gpuIndex, resources: gpuResources}, nil +} + +// to use sharding use this version of faiss - git clone -b fix-cpi-gpu-faiss_index_cpu_to_gpu_multiple_with_options https://github.com/AviadHAv/faiss.git +// TransferToAllGPUs - gpuIndexes - which gpus to use [2,4,5] , index - flat or idmap , shard - should shard accross all gpus if not will put same data on all gpus +func TransferToAllGPUs(index Index, gpuIndexes []int, sharding bool) (Index, error) { + amountOfGPUs := len(gpuIndexes) + gpuResources := make([]*C.FaissStandardGpuResources, len(gpuIndexes)) + for i := 0; i < len(gpuIndexes); i++ { + var resourceIndex *C.FaissStandardGpuResources + gpuResources[i] = resourceIndex + } + + var gpuIndex *C.FaissGpuIndex + for i := 0; i < amountOfGPUs; i++ { + c := C.faiss_StandardGpuResources_new(&gpuResources[i]) + if c != 0 { + return nil, errors.New("error on init gpu %v") + } + } + var exitCode C.int + if sharding { + exitCode = C.faiss_index_cpu_to_gpu_multiple_sharding( + (**C.FaissStandardGpuResources)(unsafe.Pointer(&gpuResources[0])), + C.size_t(len(gpuResources)), + (*C.int)(unsafe.Pointer(&gpuIndexes[0])), + C.size_t(len(gpuIndexes)), + index.cPtr(), + &gpuIndex) + } else { + exitCode = C.faiss_index_cpu_to_gpu_multiple( + (**C.FaissStandardGpuResources)(unsafe.Pointer(&gpuResources[0])), + (*C.int)(unsafe.Pointer(&gpuIndexes[0])), + C.size_t(len(gpuIndexes)), + index.cPtr(), + &gpuIndex) + } + + if exitCode != 0 { + return nil, errors.New("error transferring to gpu") + } + + return &faissIndex{idx: gpuIndex, resources: gpuResources}, nil +} + +func TransferToCpu(gpuIndex Index) (Index, error) { + var cpuIndex *C.FaissIndex + + exitCode := C.faiss_index_gpu_to_cpu(gpuIndex.cPtr(), &cpuIndex) + if exitCode != 0 { + return nil, errors.New("error transferring to gpu") + } + + Free(gpuIndex) + + return &faissIndex{idx: cpuIndex}, nil +} + +func Free(index Index) { + gpuResources := index.cGpuResource() + for _, gpuResource := range gpuResources { + C.faiss_StandardGpuResources_free(gpuResource) + } + index.Delete() + +} + +func CreateGpuIndex() (Index, error) { + var gpuResource *C.FaissStandardGpuResources + var gpuIndex *C.FaissGpuIndex + c := C.faiss_StandardGpuResources_new(&gpuResource) + if c != 0 { + return nil, errors.New("error on init gpu %v") + } + + var gpuResources []*C.FaissStandardGpuResources + gpuResources = append(gpuResources, gpuResource) + + return &faissIndex{idx: gpuIndex, resources: gpuResources}, nil +} diff --git a/gpu_bindings_cpu.go b/gpu_bindings_cpu.go new file mode 100644 index 0000000..dfcb902 --- /dev/null +++ b/gpu_bindings_cpu.go @@ -0,0 +1,30 @@ +//+build cpu + +package faiss + +import "C" +import "errors" + +func TransferToGpu(index Index) (Index, error) { + return nil, errors.New("Not supported when running in CPU mode..") +} + +func TransferToCpu(index Index) (Index, error) { + return nil, errors.New("Not supported when running in CPU mode..") +} + +func Free(gpuIndex Index) error { + return errors.New("Not supported when running in CPU mode..") +} + +func CreateGpuIndex() (Index, error) { + return nil, errors.New("Not supported when running in CPU mode..") +} + +func TransferToAllGPUs(index Index,gpuIndexes []int) (Index, error) { + return nil, errors.New("Not supported when running in CPU mode..") +} + +func TransferToAllGPUsWithOptions(index Index,gpuIndexes []int) (Index, error) { + return nil, errors.New("Not supported when running in CPU mode..") +} \ No newline at end of file diff --git a/gpu_bindings_test.go b/gpu_bindings_test.go new file mode 100644 index 0000000..14f362f --- /dev/null +++ b/gpu_bindings_test.go @@ -0,0 +1,134 @@ +//go:build gpu +// +build gpu + +package faiss + +import ( + "github.com/stretchr/testify/require" + "testing" + "time" +) + +func TestFlatIndexOnGpuFunctionality(t *testing.T) { + index, err := NewIndexFlatL2(1) + require.Nil(t, err) + + gpuIdx, err := TransferToGpu(index) + require.Nil(t, err) + + vectorsToAdd := []float32{1, 2, 3, 4, 5} + err = gpuIdx.Add(vectorsToAdd) + require.Nil(t, err) + + distances, resultIds, err := gpuIdx.Search(vectorsToAdd, 5) + require.Nil(t, err) + require.Equal(t, int64(len(vectorsToAdd)), gpuIdx.Ntotal()) + + t.Log(distances, resultIds, err) + for i := range vectorsToAdd { + require.Equal(t, int64(i), resultIds[len(vectorsToAdd)*i]) + require.Zero(t, distances[len(vectorsToAdd)*i]) + } + //This is necessary bc RemoveIDs isn't implemented for GPUIndexs + cpuIdx, err := TransferToCpu(gpuIdx) + require.Nil(t, err) + idsSelector, err := NewIDSelectorBatch([]int64{0}) + cpuIdx.RemoveIDs(idsSelector) + gpuIdx, err = TransferToGpu(cpuIdx) + require.Nil(t, err) + require.Equal(t, int64(len(vectorsToAdd)-1), gpuIdx.Ntotal()) + +} + +func TestIndexIDMapOnGPU(t *testing.T) { + index, err := NewIndexFlatL2(1) + require.Nil(t, err) + + indexMap, err := NewIndexIDMap(index) + require.Nil(t, err) + + gpuIndex, err := TransferToGpu(indexMap) + require.Nil(t, err) + + vectorsToAdd := []float32{1, 2, 3, 4, 5} + ids := make([]int64, len(vectorsToAdd)) + for i := 0; i < len(vectorsToAdd); i++ { + ids[i] = int64(i) + } + + err = gpuIndex.AddWithIDs(vectorsToAdd, ids) + require.Nil(t, err) + + distances, resultIds, err := gpuIndex.Search(vectorsToAdd, 5) + require.Nil(t, err) + t.Log(gpuIndex.D(), gpuIndex.Ntotal()) + t.Log(distances, resultIds, err) + for i := range vectorsToAdd { + require.Equal(t, ids[i], resultIds[len(vectorsToAdd)*i]) + require.Zero(t, distances[len(vectorsToAdd)*i]) + } +} + +func TestTransferToGpuAndBack(t *testing.T) { + index, err := NewIndexFlatL2(1) + require.Nil(t, err) + + indexMap, err := NewIndexIDMap(index) + require.Nil(t, err) + + gpuIndex, err := TransferToGpu(indexMap) + require.Nil(t, err) + + vectorsToAdd := []float32{1, 2, 4, 7, 11} + ids := make([]int64, len(vectorsToAdd)) + for i := 0; i < len(vectorsToAdd); i++ { + ids[i] = int64(i) + } + + err = gpuIndex.AddWithIDs(vectorsToAdd, ids) + require.Nil(t, err) + + //This is necessary bc RemoveIDs isn't implemented for GPUIndexs + cpuIdx, err := TransferToCpu(gpuIndex) + require.Nil(t, err) + idsSelector, err := NewIDSelectorBatch([]int64{0}) + cpuIdx.RemoveIDs(idsSelector) + gpuIndex, err = TransferToGpu(cpuIdx) + require.Nil(t, err) + + require.Equal(t, int64(4), gpuIndex.Ntotal()) + distances2, resultIds2, err := gpuIndex.Search([]float32{1}, 5) + t.Log(distances2, resultIds2, gpuIndex.Ntotal()) + require.Nil(t, err) + require.Equal(t, float32(1), distances2[0]) + + cpuIndex, err := TransferToCpu(gpuIndex) + require.Nil(t, err) + require.Equal(t, int64(4), cpuIndex.Ntotal()) + + idsSelector, err = NewIDSelectorBatch([]int64{0}) + cpuIndex.RemoveIDs(idsSelector) + distances2, resultIds2, err = cpuIndex.Search([]float32{1}, 5) + t.Log(distances2, resultIds2, cpuIndex.Ntotal()) + require.Nil(t, err) + require.Equal(t, float32(1), distances2[0]) + +} + +func TestFreeGPUResource(t *testing.T) { + for i := 0; i < 20; i++ { + gpus := []int{0} + t.Logf("creating index %v", i) + flatIndex, err := NewIndexFlatIP(256) + require.Nil(t, err) + flatIndexGpu, err := TransferToAllGPUs(flatIndex, gpus, false) + require.Nil(t, err) + + t.Log("created indexes, freeing..") + Free(flatIndexGpu) + require.Nil(t, err) + t.Log("freed, memory should be freed..") + time.Sleep(1 * time.Second) + } + +} diff --git a/index.go b/index.go index 40282f0..c7d2790 100644 --- a/index.go +++ b/index.go @@ -5,6 +5,8 @@ package faiss #include #include #include +#include +#include */ import "C" import "unsafe" @@ -56,10 +58,24 @@ type Index interface { Delete() cPtr() *C.FaissIndex + + cGpuResource() []*C.FaissStandardGpuResources + + cGpuMultipleClonerOptions() *C.FaissGpuMultipleClonerOptions } type faissIndex struct { idx *C.FaissIndex + resources []*C.FaissStandardGpuResources + options *C.FaissGpuMultipleClonerOptions +} + +func (idx *faissIndex) cGpuResource() []*C.FaissStandardGpuResources { + return idx.resources +} + +func (idx *faissIndex) cGpuMultipleClonerOptions() *C.FaissGpuMultipleClonerOptions { + return idx.options } func (idx *faissIndex) cPtr() *C.FaissIndex { diff --git a/index_flat.go b/index_flat.go index b8a3c03..a97d6f8 100644 --- a/index_flat.go +++ b/index_flat.go @@ -52,5 +52,5 @@ func (idx *IndexImpl) AsFlat() *IndexFlat { if ptr == nil { panic("index is not a flat index") } - return &IndexFlat{&faissIndex{ptr}} + return &IndexFlat{&faissIndex{idx: ptr}} } diff --git a/index_idmap.go b/index_idmap.go new file mode 100644 index 0000000..4551152 --- /dev/null +++ b/index_idmap.go @@ -0,0 +1,19 @@ +package faiss + +/* +#include +*/ +import "C" +import ( + "errors" +) + +func NewIndexIDMap(index Index) (Index, error) { + var indexMapPointer *C.FaissIndexIDMap + var pointerToIndexMapPointer **C.FaissIndexIDMap + pointerToIndexMapPointer = &indexMapPointer + if C.faiss_IndexIDMap_new(pointerToIndexMapPointer, index.cPtr()) != 0 { + return nil, errors.New("Error occurred while initializing IndexIDMapWrapper") + } + return &faissIndex{idx: indexMapPointer}, nil +} diff --git a/index_idmap_test.go b/index_idmap_test.go new file mode 100644 index 0000000..b9d97e1 --- /dev/null +++ b/index_idmap_test.go @@ -0,0 +1,36 @@ +package faiss + +import ( + "fmt" + "github.com/stretchr/testify/require" + "testing" +) + +func TestNewIndexIDMap(t *testing.T) { + dimension := 1 + dbSize := 5 + + index, err := NewIndexFlat(dimension, MetricL2) + if err != nil { + fmt.Println(err.Error()) + } + indexMap, err := NewIndexIDMap(index) + if err != nil { + fmt.Println(err.Error()) + } + xb := []float32{1,2,3,4,5} + ids := make([]int64, dbSize) + for i := 10; i < dbSize; i++ { + ids[i] = int64(i) + } + + err = indexMap.AddWithIDs(xb, ids) + if err != nil { + fmt.Println(err.Error()) + } + toFind := xb[dimension:2*dimension] + distances1, resultIds, err := indexMap.Search(toFind, 5) + require.Equal(t, resultIds[0], ids[1]) + require.Zero(t, distances1[0]) + +}