From 5744129b0e65fe502305472e653940e1333e9db3 Mon Sep 17 00:00:00 2001
From: Alexander Andreev <andreevlex.as@gmail.com>
Date: Wed, 14 Jul 2021 21:12:49 +0300
Subject: [PATCH] Impl to_gpu for new index types

IVFFlatIndexImpl
IVFScalarQuantizerIndexImpl
---
 src/index/gpu.rs | 137 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 137 insertions(+)

diff --git a/src/index/gpu.rs b/src/index/gpu.rs
index 3093069..8003b10 100644
--- a/src/index/gpu.rs
+++ b/src/index/gpu.rs
@@ -1,6 +1,8 @@
 //! GPU Index implementation
 
 use super::flat::FlatIndexImpl;
+use super::ivf_flat::IVFFlatIndexImpl;
+use super::scalar_quantizer::IVFScalarQuantizerIndexImpl;
 use super::{
     AssignSearchResult, CpuIndex, FromInnerPtr, Idx, Index, IndexImpl, NativeIndex,
     RangeSearchResult, SearchResult,
@@ -424,6 +426,141 @@ impl FlatIndexImpl {
     }
 }
 
+impl IVFFlatIndexImpl {
+    /// Build a GPU in from the given CPU native index, yielding two
+    /// independent indices. The operation fails if the index does
+    /// not provide GPU support.
+    pub fn to_gpu<'gpu, G>(
+        &self,
+        gpu_res: &'gpu G,
+        device: i32,
+    ) -> Result<GpuIndexImpl<'gpu, IVFFlatIndexImpl>>
+    where
+        G: GpuResourcesProvider,
+    {
+        GpuIndexImpl::from_cpu(self, gpu_res, device)
+    }
+
+    /// Build a GPU index from the given CPU native index, discarding the
+    /// CPU-backed index. The operation fails if the index does not
+    /// provide GPU support.
+    pub fn into_gpu<'gpu, G>(
+        self,
+        gpu_res: &'gpu G,
+        device: i32,
+    ) -> Result<GpuIndexImpl<'gpu, IVFFlatIndexImpl>>
+    where
+        G: GpuResourcesProvider,
+    {
+        self.to_gpu(gpu_res, device)
+    }
+
+    /// Build a GPU index from the given CPU native index.
+    ///
+    /// # Errors
+    ///
+    /// The operation fails if the number of GPU resources and number of
+    /// devices do not match, or the index does not provide GPU support.
+    pub fn to_gpu_multiple<'gpu, G: 'gpu>(
+        &self,
+        gpu_res: &'gpu [G],
+        devices: &[i32],
+    ) -> Result<GpuIndexImpl<'gpu, IVFFlatIndexImpl>>
+    where
+        G: GpuResourcesProvider,
+    {
+        GpuIndexImpl::from_cpu_multiple(&self, gpu_res, devices)
+    }
+
+    /// Build a GPU index from the given CPU native index. The index residing
+    /// in CPU memory is discarded in the process.
+    ///
+    /// # Errors
+    ///
+    /// The operation fails if the number of GPU resources and number of
+    /// devices do not match, or the index does not provide GPU support.
+    pub fn into_gpu_multiple<'gpu, G: 'gpu>(
+        self,
+        gpu_res: &'gpu [G],
+        devices: &[i32],
+    ) -> Result<GpuIndexImpl<'gpu, IVFFlatIndexImpl>>
+    where
+        G: GpuResourcesProvider,
+    {
+        self.to_gpu_multiple(gpu_res, devices)
+        // let the CPU index drop naturally
+    }
+}
+
+impl<Q> IVFScalarQuantizerIndexImpl<Q>
+where
+    Q: NativeIndex + CpuIndex,
+{
+    /// Build a GPU in from the given CPU native index, yielding two
+    /// independent indices. The operation fails if the index does
+    /// not provide GPU support.
+    pub fn to_gpu<'gpu, G>(
+        &self,
+        gpu_res: &'gpu G,
+        device: i32,
+    ) -> Result<GpuIndexImpl<'gpu, IVFScalarQuantizerIndexImpl<Q>>>
+    where
+        G: GpuResourcesProvider,
+    {
+        GpuIndexImpl::from_cpu(self, gpu_res, device)
+    }
+
+    /// Build a GPU index from the given CPU native index, discarding the
+    /// CPU-backed index. The operation fails if the index does not
+    /// provide GPU support.
+    pub fn into_gpu<'gpu, G>(
+        self,
+        gpu_res: &'gpu G,
+        device: i32,
+    ) -> Result<GpuIndexImpl<'gpu, IVFScalarQuantizerIndexImpl<Q>>>
+    where
+        G: GpuResourcesProvider,
+    {
+        self.to_gpu(gpu_res, device)
+    }
+
+    /// Build a GPU index from the given CPU native index.
+    ///
+    /// # Errors
+    ///
+    /// The operation fails if the number of GPU resources and number of
+    /// devices do not match, or the index does not provide GPU support.
+    pub fn to_gpu_multiple<'gpu, G: 'gpu>(
+        &self,
+        gpu_res: &'gpu [G],
+        devices: &[i32],
+    ) -> Result<GpuIndexImpl<'gpu, IVFScalarQuantizerIndexImpl<Q>>>
+    where
+        G: GpuResourcesProvider,
+    {
+        GpuIndexImpl::from_cpu_multiple(&self, gpu_res, devices)
+    }
+
+    /// Build a GPU index from the given CPU native index. The index residing
+    /// in CPU memory is discarded in the process.
+    ///
+    /// # Errors
+    ///
+    /// The operation fails if the number of GPU resources and number of
+    /// devices do not match, or the index does not provide GPU support.
+    pub fn into_gpu_multiple<'gpu, G: 'gpu>(
+        self,
+        gpu_res: &'gpu [G],
+        devices: &[i32],
+    ) -> Result<GpuIndexImpl<'gpu, IVFScalarQuantizerIndexImpl<Q>>>
+    where
+        G: GpuResourcesProvider,
+    {
+        self.to_gpu_multiple(gpu_res, devices)
+        // let the CPU index drop naturally
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::super::{index_factory, CpuIndex, Idx, Index};