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/gpu_bindings.go b/gpu_bindings.go new file mode 100644 index 0000000..5a8096b --- /dev/null +++ b/gpu_bindings.go @@ -0,0 +1,60 @@ +//+build gpu + +package faiss + +/* +#include +#include +*/ +import "C" +import ( + "errors" +) + + +func TransferToGpu(index Index) (Index, error) { + var gpuResources *C.FaissStandardGpuResources + var gpuIndex *C.FaissGpuIndex + c := C.faiss_StandardGpuResources_new(&gpuResources) + if c != 0 { + return nil, errors.New("error on init gpu %v") + } + + exitCode := C.faiss_index_cpu_to_gpu(gpuResources, 0, index.cPtr(), &gpuIndex) + if exitCode != 0 { + return nil, errors.New("error transferring to gpu") + } + + return &faissIndex{idx: gpuIndex, resource: 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) { + var gpuResource *C.FaissStandardGpuResources + gpuResource = index.cGpuResource() + 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") + } + + return &faissIndex{idx: gpuIndex, resource: gpuResource}, nil +} diff --git a/gpu_bindings_cpu.go b/gpu_bindings_cpu.go new file mode 100644 index 0000000..f24b073 --- /dev/null +++ b/gpu_bindings_cpu.go @@ -0,0 +1,21 @@ +//+build cpu + +package faiss + +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..") +} diff --git a/gpu_bindings_test.go b/gpu_bindings_test.go new file mode 100644 index 0000000..b65cf42 --- /dev/null +++ b/gpu_bindings_test.go @@ -0,0 +1,133 @@ +//+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++ { + t.Logf("creating index %v", i) + flatIndex, err := NewIndexFlatIP(256) + require.Nil(t, err) + flatIndexGpu, err := TransferToGpu(flatIndex) + require.Nil(t, err) + + t.Log("created indexes, freeing..") + err = Free(flatIndexGpu) + require.Nil(t, err) + t.Log("freed, memory should be freed..") + time.Sleep(1 * time.Second) + } + +} \ No newline at end of file diff --git a/index.go b/index.go index 40282f0..29ff19c 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,17 @@ type Index interface { Delete() cPtr() *C.FaissIndex + + cGpuResource() *C.FaissStandardGpuResources } type faissIndex struct { idx *C.FaissIndex + resource *C.FaissStandardGpuResources +} + +func (idx *faissIndex) cGpuResource() *C.FaissStandardGpuResources { + return idx.resource } 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]) + +}