Gpu sharding

_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

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=

gpu_bindings.go

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

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

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

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,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 {

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