From deb208709035888a1854ccb458dc262c5c23b175 Mon Sep 17 00:00:00 2001
From: "aviadh@anyvision.co" <aviadh@anyvision.co>
Date: Tue, 21 Jun 2022 17:32:13 +0300
Subject: [PATCH] support gpu sharding

currently there is bug on faiss https://github.com/facebookresearch/faiss/issues/2357
to use sharding you have to use this version of faiss for now git clone -b fix-cpi-gpu-faiss_index_cpu_to_gpu_multiple_with_options https://github.com/AviadHAv/faiss.git
---
 gpu_bindings.go      | 30 +++++++++++++++++++++---------
 gpu_bindings_test.go | 16 ++++++++--------
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/gpu_bindings.go b/gpu_bindings.go
index 9c08289..2e17ab7 100644
--- a/gpu_bindings.go
+++ b/gpu_bindings.go
@@ -8,6 +8,7 @@ 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 (
@@ -35,13 +36,14 @@ func TransferToGpu(index Index) (Index, error) {
 	return &faissIndex{idx: gpuIndex, resources: gpuResources}, nil
 }
 
-func TransferToAllGPUs(index Index, gpuIndexes []int) (Index, error) {
+// 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
@@ -51,13 +53,23 @@ func TransferToAllGPUs(index Index, gpuIndexes []int) (Index, error) {
 			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)
+	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")
diff --git a/gpu_bindings_test.go b/gpu_bindings_test.go
index a86b7ba..14f362f 100644
--- a/gpu_bindings_test.go
+++ b/gpu_bindings_test.go
@@ -1,4 +1,5 @@
-//+build gpu
+//go:build gpu
+// +build gpu
 
 package faiss
 
@@ -15,7 +16,7 @@ func TestFlatIndexOnGpuFunctionality(t *testing.T) {
 	gpuIdx, err := TransferToGpu(index)
 	require.Nil(t, err)
 
-	vectorsToAdd := []float32{1,2,3,4,5}
+	vectorsToAdd := []float32{1, 2, 3, 4, 5}
 	err = gpuIdx.Add(vectorsToAdd)
 	require.Nil(t, err)
 
@@ -49,7 +50,7 @@ func TestIndexIDMapOnGPU(t *testing.T) {
 	gpuIndex, err := TransferToGpu(indexMap)
 	require.Nil(t, err)
 
-	vectorsToAdd := []float32{1,2,3,4,5}
+	vectorsToAdd := []float32{1, 2, 3, 4, 5}
 	ids := make([]int64, len(vectorsToAdd))
 	for i := 0; i < len(vectorsToAdd); i++ {
 		ids[i] = int64(i)
@@ -78,7 +79,7 @@ func TestTransferToGpuAndBack(t *testing.T) {
 	gpuIndex, err := TransferToGpu(indexMap)
 	require.Nil(t, err)
 
-	vectorsToAdd := []float32{1,2,4,7,11}
+	vectorsToAdd := []float32{1, 2, 4, 7, 11}
 	ids := make([]int64, len(vectorsToAdd))
 	for i := 0; i < len(vectorsToAdd); i++ {
 		ids[i] = int64(i)
@@ -101,7 +102,6 @@ func TestTransferToGpuAndBack(t *testing.T) {
 	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())
@@ -117,11 +117,11 @@ func TestTransferToGpuAndBack(t *testing.T) {
 
 func TestFreeGPUResource(t *testing.T) {
 	for i := 0; i < 20; i++ {
-		gpus:= []int{0}
+		gpus := []int{0}
 		t.Logf("creating index %v", i)
 		flatIndex, err := NewIndexFlatIP(256)
 		require.Nil(t, err)
-		flatIndexGpu, err := TransferToAllGPUs(flatIndex, gpus)
+		flatIndexGpu, err := TransferToAllGPUs(flatIndex, gpus, false)
 		require.Nil(t, err)
 
 		t.Log("created indexes, freeing..")
@@ -131,4 +131,4 @@ func TestFreeGPUResource(t *testing.T) {
 		time.Sleep(1 * time.Second)
 	}
 
-}
\ No newline at end of file
+}