gpu_bindings.go: added TransferToCpu method

This is potentially necessary for various things, but specifically is
critical for us NOW because RemoveIDs isn't implemented for GPUIndexs,
so until we implement that functionality in faiss, the only solution is
to transfer the index back and forth to the CPU and remove ids which the
index is on CPU.

Author: orihabAnyvision

gpu_bindings.go

@@ -24,3 +24,15 @@ func TransferToGpu(index Index) (Index, error) {
 
 	return &faissIndex{idx: gpuIndex}, 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")
+	}
+
+	return &faissIndex{idx: cpuIndex}, nil
+}
+

gpu_bindings_test.go

@@ -6,7 +6,7 @@ import (
 	"testing"
 )
 
-func TestFlatIndexOnGpu(t *testing.T) {
+func TestFlatIndexOnGpuFunctionality(t *testing.T) {
 	index, err := NewIndexFlatL2(1)
 	require.Nil(t, err)
 
@@ -19,12 +19,22 @@ func TestFlatIndexOnGpu(t *testing.T) {
 
 	distances, resultIds, err := gpuIdx.Search(vectorsToAdd, 5)
 	require.Nil(t, err)
+	require.Equal(t, int64(len(vectorsToAdd)), gpuIdx.Ntotal())
 
 	fmt.Println(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) {
@@ -55,3 +65,51 @@ func TestIndexIDMapOnGPU(t *testing.T) {
 		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)
+	fmt.Println(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)
+	fmt.Println(distances2, resultIds2, cpuIndex.Ntotal())
+	require.Nil(t, err)
+	require.Equal(t, float32(1), distances2[0])
+
+}
+