Merge pull request #1 from AnyVisionltd/support-multiple-gpus

Support multiple gpus

Author: NitzanZAnyvision

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

@@ -1,31 +1,69 @@
-//+build gpu
+//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>
 */
 import "C"
 import (
 	"errors"
+	"unsafe"
 )
 
-
 func TransferToGpu(index Index) (Index, error) {
-	var gpuResources *C.FaissStandardGpuResources
+	var gpuResource *C.FaissStandardGpuResources
 	var gpuIndex *C.FaissGpuIndex
-	c := C.faiss_StandardGpuResources_new(&gpuResources)
+	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(gpuResources, 0, index.cPtr(), &gpuIndex)
+	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
+}
+
+func TransferToAllGPUs(index Index, gpuIndexes []int) (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")
+		}
+	}
+
+	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, resource: gpuResources}, nil
+	return &faissIndex{idx: gpuIndex, resources: gpuResources}, nil
 }
 
 func TransferToCpu(gpuIndex Index) (Index, error) {
@@ -37,15 +75,17 @@ func TransferToCpu(gpuIndex Index) (Index, error) {
 	}
 
 	Free(gpuIndex)
-	
+
 	return &faissIndex{idx: cpuIndex}, nil
 }
 
 func Free(index Index) {
-	var gpuResource *C.FaissStandardGpuResources
-	gpuResource = index.cGpuResource()
-	C.faiss_StandardGpuResources_free(gpuResource)
+	gpuResources := index.cGpuResource()
+	for _, gpuResource := range gpuResources {
+		C.faiss_StandardGpuResources_free(gpuResource)
+	}
 	index.Delete()
+
 }
 
 func CreateGpuIndex() (Index, error) {
@@ -56,5 +96,8 @@ func CreateGpuIndex() (Index, error) {
 		return nil, errors.New("error on init gpu %v")
 	}
 
-	return &faissIndex{idx: gpuIndex, resource: gpuResource}, nil
+	var gpuResources []*C.FaissStandardGpuResources
+	gpuResources = append(gpuResources, gpuResource)
+
+	return &faissIndex{idx: gpuIndex, resources: gpuResources}, nil
 }

gpu_bindings_cpu.go

@@ -2,6 +2,7 @@
 
 package faiss
 
+import "C"
 import "errors"
 
 func TransferToGpu(index Index) (Index, error) {
@@ -19,3 +20,11 @@ func Free(gpuIndex Index) error {
 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

@@ -117,14 +117,15 @@ func TestTransferToGpuAndBack(t *testing.T) {
 
 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 := TransferToGpu(flatIndex)
+		flatIndexGpu, err := TransferToAllGPUs(flatIndex, gpus)
 		require.Nil(t, err)
 
 		t.Log("created indexes, freeing..")
-		err = Free(flatIndexGpu)
+		Free(flatIndexGpu)
 		require.Nil(t, err)
 		t.Log("freed, memory should be freed..")
 		time.Sleep(1 * time.Second)

index.go

@@ -59,16 +59,23 @@ type Index interface {
 
 	cPtr() *C.FaissIndex
 
-	cGpuResource() *C.FaissStandardGpuResources
+	cGpuResource() []*C.FaissStandardGpuResources
+
+	cGpuMultipleClonerOptions() *C.FaissGpuMultipleClonerOptions
 }
 
 type faissIndex struct {
 	idx *C.FaissIndex
-	resource *C.FaissStandardGpuResources
+	resources []*C.FaissStandardGpuResources
+	options  *C.FaissGpuMultipleClonerOptions
+}
+
+func (idx *faissIndex) cGpuResource() []*C.FaissStandardGpuResources {
+	return idx.resources
 }
 
-func (idx *faissIndex) cGpuResource() *C.FaissStandardGpuResources {
-	return idx.resource
+func (idx *faissIndex) cGpuMultipleClonerOptions() *C.FaissGpuMultipleClonerOptions {
+	return idx.options
 }
 
 func (idx *faissIndex) cPtr() *C.FaissIndex {