Python API for IVF-Flat serialization

cpp/CMakeLists.txt

@@ -364,6 +364,7 @@ if(RAFT_COMPILE_LIBRARY)
     src/raft_runtime/neighbors/brute_force_knn_int64_t_float.cu
     src/raft_runtime/neighbors/ivf_flat_build.cu
     src/raft_runtime/neighbors/ivf_flat_search.cu
+    src/raft_runtime/neighbors/ivf_flat_serialize.cu
     src/raft_runtime/neighbors/ivfpq_build.cu
     src/raft_runtime/neighbors/ivfpq_deserialize.cu
     src/raft_runtime/neighbors/ivfpq_search_float_int64_t.cu

cpp/include/raft_runtime/neighbors/ivf_flat.hpp

@@ -43,7 +43,15 @@ namespace raft::runtime::neighbors::ivf_flat {
   void extend(raft::device_resources const& handle,                                  \
               raft::device_matrix_view<const T, IdxT, row_major> new_vectors,        \
               std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
-              raft::neighbors::ivf_flat::index<T, IdxT>* idx);
+              raft::neighbors::ivf_flat::index<T, IdxT>* idx);                       \
+                                                                                     \
+  void serialize(raft::device_resources const& handle,                               \
+                 const std::string& filename,                                        \
+                 const raft::neighbors::ivf_flat::index<T, IdxT>& index);            \
+                                                                                     \
+  void deserialize(raft::device_resources const& handle,                             \
+                   const std::string& filename,                                      \
+                   raft::neighbors::ivf_flat::index<T, IdxT>* index);
 
 RAFT_INST_BUILD_EXTEND(float, int64_t)
 RAFT_INST_BUILD_EXTEND(int8_t, int64_t)

cpp/src/raft_runtime/neighbors/ivf_flat_serialize.cu

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <string>
+
+#include <raft/core/device_resources.hpp>
+#include <raft/neighbors/ivf_flat_serialize.cuh>
+#include <raft/neighbors/ivf_flat_types.hpp>
+#include <raft_runtime/neighbors/ivf_flat.hpp>
+
+namespace raft::runtime::neighbors::ivf_flat {
+
+#define RAFT_IVF_FLAT_SERIALIZE_INST(DTYPE)                                            \
+  void serialize(raft::device_resources const& handle,                                 \
+                 const std::string& filename,                                          \
+                 const raft::neighbors::ivf_flat::index<DTYPE, int64_t>& index)        \
+  {                                                                                    \
+    raft::neighbors::ivf_flat::serialize(handle, filename, index);                     \
+  };                                                                                   \
+                                                                                       \
+  void deserialize(raft::device_resources const& handle,                               \
+                   const std::string& filename,                                        \
+                   raft::neighbors::ivf_flat::index<DTYPE, int64_t>* index)            \
+  {                                                                                    \
+    if (!index) { RAFT_FAIL("Invalid index pointer"); }                                \
+    *index = raft::neighbors::ivf_flat::deserialize<DTYPE, int64_t>(handle, filename); \
+  };
+
+RAFT_IVF_FLAT_SERIALIZE_INST(float);
+RAFT_IVF_FLAT_SERIALIZE_INST(int8_t);
+RAFT_IVF_FLAT_SERIALIZE_INST(uint8_t);
+
+#undef RAFT_IVF_FLAT_SERIALIZE_INST
+}  // namespace raft::runtime::neighbors::ivf_flat

python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 
-from .ivf_flat import Index, IndexParams, SearchParams, build, extend, search
+from .ivf_flat import Index, IndexParams, SearchParams, build, extend, search, save, load
 
 __all__ = [
     "Index",
@@ -22,4 +22,6 @@
     "build",
     "extend",
     "search",
+    "save",
+    "load"
 ]

python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd

@@ -133,3 +133,28 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \
         device_matrix_view[uint8_t, int64_t, row_major] queries,
         device_matrix_view[int64_t, int64_t, row_major] neighbors,
         device_matrix_view[float, int64_t, row_major] distances) except +
+
+
+    cdef void serialize(const device_resources& handle,
+                        const string& filename,
+                        const index[float, int64_t]& index) except +
+
+    cdef void deserialize(const device_resources& handle,
+                          const string& filename,
+                          index[float, int64_t]* index) except +
+
+    cdef void serialize(const device_resources& handle,
+                        const string& filename,
+                        const index[uint8_t, int64_t]& index) except +
+
+    cdef void deserialize(const device_resources& handle,
+                          const string& filename,
+                          index[uint8_t, int64_t]* index) except +
+
+    cdef void serialize(const device_resources& handle,
+                        const string& filename,
+                        const index[int8_t, int64_t]& index) except +
+
+    cdef void deserialize(const device_resources& handle,
+                          const string& filename,
+                          index[int8_t, int64_t]* index) except +

python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx

@@ -708,3 +708,141 @@ def search(SearchParams search_params,
         raise ValueError("query dtype %s not supported" % queries_dt)
 
     return (distances, neighbors)
+
+@auto_sync_handle
+def save(filename, Index index, handle=None):
+    """
+    Saves the index to file.
+
+    Saving / loading the index is experimental. The serialization format is
+    subject to change.
+
+    Parameters
+    ----------
+    filename : string
+        Name of the file.
+    index : Index
+        Trained IVF-Flat index.
+    {handle_docstring}
+
+    Examples
+    --------
+    >>> import cupy as cp
+
+    >>> from pylibraft.common import DeviceResources
+    >>> from pylibraft.neighbors import ivf_flat
+
+    >>> n_samples = 50000
+    >>> n_features = 50
+    >>> dataset = cp.random.random_sample((n_samples, n_features),
+    ...                                   dtype=cp.float32)
+
+    >>> # Build index
+    >>> handle = DeviceResources()
+    >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, handle=handle)
+    >>> ivf_flat.save("my_index.bin", index, handle=handle)
+    """
+    if not index.trained:
+        raise ValueError("Index need to be built before saving it.")
+
+    if handle is None:
+        handle = DeviceResources()
+    cdef device_resources* handle_ = \
+        <device_resources*><size_t>handle.getHandle()
+
+    cdef string c_filename = filename.encode('utf-8')
+
+    cdef IndexFloat idx_float
+    cdef IndexInt8 idx_int8
+    cdef IndexUint8 idx_uint8
+
+    if index.active_index_type == "float32":
+        idx_float = index
+        c_ivf_flat.serialize(deref(handle_), c_filename, deref(idx_float.index))
+    elif index.active_index_type == "byte":
+        idx_int8 = index
+        c_ivf_flat.serialize(deref(handle_), c_filename, deref(idx_int8.index))
+    elif index.active_index_type == "ubyte":
+        idx_uint8 = index
+        c_ivf_flat.serialize(deref(handle_), c_filename, deref(idx_uint8.index))
+    else:
+        raise ValueError("Index dtype %s not supported" % index.active_index_type)
+
+
+@auto_sync_handle
+def load(filename, dtype, handle=None):
+    """
+    Loads index from file.
+
+    Saving / loading the index is experimental. The serialization format is
+    subject to change, therefore loading an index saved with a previous
+    version of raft is not guaranteed to work.
+
+    Parameters
+    ----------
+    filename : string
+        Name of the file.
+    dtype : data type object
+        dataset type, supported values [np.float32, np.byte, np.ubyte]
+    {handle_docstring}
+
+    Returns
+    -------
+    index : Index
+
+    Examples
+    --------
+    >>> import cupy as cp
+
+    >>> from pylibraft.common import DeviceResources
+    >>> from pylibraft.neighbors import ivf_flat
+
+    >>> n_samples = 50000
+    >>> n_features = 50
+    >>> dataset = cp.random.random_sample((n_samples, n_features),
+    ...                                   dtype=cp.float32)
+
+    >>> # Build and save index
+    >>> handle = DeviceResources()
+    >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, handle=handle)
+    >>> ivf_flat.save("my_index.bin", index, handle=handle)
+    >>> del index
+
+    >>> n_queries = 100
+    >>> queries = cp.random.random_sample((n_queries, n_features),
+    ...                                   dtype=cp.float32)
+    >>> handle = DeviceResources()
+    >>> index = ivf_flat.load("my_index.bin", dtype=cp.float32, handle=handle)
+
+    >>> distances, neighbors = ivf_flat.search(ivf_pq.SearchParams(), index,
+    ...                                      queries, k=10, handle=handle)
+    """
+    if handle is None:
+        handle = DeviceResources()
+    cdef device_resources* handle_ = \
+        <device_resources*><size_t>handle.getHandle()
+
+    cdef string c_filename = filename.encode('utf-8')
+    cdef IndexFloat idx_float
+    cdef IndexInt8 idx_int8
+    cdef IndexUint8 idx_uint8
+
+    dataset_dt = np.dtype(dtype)
+    if dataset_dt == np.float32:
+        idx_float = IndexFloat(handle)
+        c_ivf_flat.deserialize(deref(handle_), c_filename, idx_float.index)
+        idx_float.trained = True
+        return idx_float
+    elif dataset_dt == np.byte:
+        idx_int8 = IndexInt8(handle)
+        c_ivf_flat.deserialize(deref(handle_), c_filename, idx_int8.index)
+        idx_int8.trained = True
+        return idx_int8
+    elif dataset_dt == np.ubyte:
+        idx_uint8 = IndexUint8(handle)
+        c_ivf_flat.deserialize(deref(handle_), c_filename, idx_uint8.index)
+        idx_uint8.trained = True
+        return idx_uint8
+    else:
+        raise ValueError("Index dtype %s not supported" % dtype)
+

python/pylibraft/pylibraft/test/test_ivf_flat.py

@@ -165,9 +165,7 @@ def run_ivf_flat_build_search_test(
         "inner_product": "cosine",
         "euclidean": "euclidean",
     }[metric]
-    nn_skl = NearestNeighbors(
-        n_neighbors=k, algorithm="brute", metric=skl_metric
-    )
+    nn_skl = NearestNeighbors(n_neighbors=k, algorithm="brute", metric=skl_metric)
     nn_skl.fit(dataset)
     skl_idx = nn_skl.kneighbors(queries, return_distance=False)
 
@@ -184,9 +182,7 @@ def run_ivf_flat_build_search_test(
 @pytest.mark.parametrize("n_lists", [100])
 @pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8])
 @pytest.mark.parametrize("array_type", ["device"])
-def test_ivf_pq_dtypes(
-    n_rows, n_cols, n_queries, n_lists, dtype, inplace, array_type
-):
+def test_ivf_pq_dtypes(n_rows, n_cols, n_queries, n_lists, dtype, inplace, array_type):
     # Note that inner_product tests use normalized input which we cannot
     # represent in int8, therefore we test only sqeuclidean metric here.
     run_ivf_flat_build_search_test(
@@ -236,9 +232,7 @@ def test_ivf_flat_n(params):
     )
 
 
-@pytest.mark.parametrize(
-    "metric", ["sqeuclidean", "inner_product", "euclidean"]
-)
+@pytest.mark.parametrize("metric", ["sqeuclidean", "inner_product", "euclidean"])
 @pytest.mark.parametrize("dtype", [np.float32])
 def test_ivf_flat_build_params(metric, dtype):
     run_ivf_flat_build_search_test(
@@ -419,9 +413,9 @@ def test_search_inputs(params):
 
     q_dt = params.get("q_dt", np.float32)
     q_order = params.get("q_order", "C")
-    queries = generate_data(
-        (n_queries, params.get("q_cols", n_cols)), q_dt
-    ).astype(q_dt, order=q_order)
+    queries = generate_data((n_queries, params.get("q_cols", n_cols)), q_dt).astype(
+        q_dt, order=q_order
+    )
     queries_device = device_ndarray(queries)
 
     idx_dt = params.get("idx_dt", np.int64)
@@ -461,3 +455,51 @@ def test_search_inputs(params):
             out_idx_device,
             out_dist_device,
         )
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.ubyte])
+def test_save_load(dtype):
+    n_rows = 10000
+    n_cols = 50
+    n_queries = 1000
+
+    dataset = generate_data((n_rows, n_cols), dtype)
+    dataset_device = device_ndarray(dataset)
+
+    build_params = ivf_flat.IndexParams(n_lists=100, metric="sqeuclidean")
+    index = ivf_flat.build(build_params, dataset_device)
+
+    assert index.trained
+    filename = "my_index.bin"
+    ivf_flat.save(filename, index)
+    loaded_index = ivf_flat.load(filename, dtype)
+
+    assert index.metric == loaded_index.metric
+    assert index.n_lists == loaded_index.n_lists
+    assert index.dim == loaded_index.dim
+    assert index.adaptive_centers == loaded_index.adaptive_centers
+    # assert index.size == loaded_index.size
+
+    queries = generate_data((n_queries, n_cols), dtype)
+
+    queries_device = device_ndarray(queries)
+    search_params = ivf_flat.SearchParams(n_probes=100)
+    k = 10
+
+    distance_dev, neighbors_dev = ivf_flat.search(
+        search_params, index, queries_device, k
+    )
+
+    neighbors = neighbors_dev.copy_to_host()
+    dist = distance_dev.copy_to_host()
+    del index
+
+    distance_dev, neighbors_dev = ivf_flat.search(
+        search_params, loaded_index, queries_device, k
+    )
+
+    neighbors2 = neighbors_dev.copy_to_host()
+    dist2 = distance_dev.copy_to_host()
+
+    assert np.all(neighbors == neighbors2)
+    assert np.allclose(dist, dist2, rtol=1e-6)