C bindings for IDIndexMap

_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)
+
+}

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

gpu_bindings.go

@@ -0,0 +1,60 @@
+//+build gpu
+
+package faiss
+
+/*
+#include <faiss/c_api/gpu/StandardGpuResources_c.h>
+#include <faiss/c_api/gpu/GpuAutoTune_c.h>
+*/
+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
+}

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..")
+}

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)
+	}
+
+}

index.go

@@ -5,6 +5,8 @@ package faiss
 #include <faiss/c_api/Index_c.h>
 #include <faiss/c_api/impl/AuxIndexStructures_c.h>
 #include <faiss/c_api/index_factory_c.h>
+#include <faiss/c_api/gpu/StandardGpuResources_c.h>
+#include <faiss/c_api/gpu/GpuAutoTune_c.h>
 */
 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 {

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}}
 }

index_idmap.go

@@ -0,0 +1,19 @@
+package faiss
+
+/*
+#include <faiss/c_api/MetaIndexes_c.h>
+*/
+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
+}

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])
+
+}