IVF-Flat Python wrappers

cpp/CMakeLists.txt

@@ -370,6 +370,8 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/neighbors/ivfpq_search_float_uint64_t.cu
     src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu
     src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu
+    src/distance/neighbors/ivf_flat_search.cu
+    src/distance/neighbors/ivf_flat_build.cu
     src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu
     src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu
     src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu

cpp/include/raft_runtime/neighbors/ivf_flat.hpp

@@ -0,0 +1,84 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <raft/neighbors/ivf_flat_types.hpp>
+
+namespace raft::runtime::neighbors::ivf_flat {
+
+// We define overloads for build and extend with void return type. This is used in the Cython
+// wrappers, where exception handling is not compatible with return type that has nontrivial
+// constructor.
+#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                      \
+  auto build(raft::device_resources const& handle,                           \
+             raft::device_matrix_view<const T, uint64_t, row_major> dataset, \
+             const raft::neighbors::ivf_flat::index_params& params)          \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>;                             \
+                                                                             \
+  auto extend(raft::device_resources const& handle,                          \
+              const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index,   \
+              const T* new_vectors,                                          \
+              const IdxT* new_indices,                                       \
+              IdxT n_rows)                                                   \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>;                             \
+                                                                             \
+  void build(raft::device_resources const& handle,                           \
+             raft::device_matrix_view<const T, uint64_t, row_major> dataset, \
+             const raft::neighbors::ivf_flat::index_params& params,          \
+             raft::neighbors::ivf_flat::index<T, IdxT>* idx);                \
+                                                                             \
+  void extend(raft::device_resources const& handle,                          \
+              raft::neighbors::ivf_flat::index<T, IdxT>* idx,                \
+              const T* new_vectors,                                          \
+              const IdxT* new_indices,                                       \
+              IdxT n_rows);
+
+RAFT_INST_BUILD_EXTEND(float, uint64_t)
+RAFT_INST_BUILD_EXTEND(int8_t, uint64_t)
+RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t)
+
+#undef RAFT_INST_BUILD_EXTEND
+
+// /**
+//  * Save the index to file.
+//  *
+//  * Experimental, both the API and the serialization format are subject to change.
+//  *
+//  * @param[in] handle the raft handle
+//  * @param[in] filename the filename for saving the index
+//  * @param[in] index IVF-PQ index
+//  *
+//  */
+// void save(raft::device_resources const& handle,
+//           const std::string& filename,
+//           const raft::neighbors::ivf_flat::index<uint64_t>& index);
+
+// /**
+//  * Load index from file.
+//  *
+//  * Experimental, both the API and the serialization format are subject to change.
+//  *
+//  * @param[in] handle the raft handle
+//  * @param[in] filename the name of the file that stores the index
+//  * @param[in] index IVF-PQ index
+//  *
+//  */
+// void load(raft::device_resources const& handle,
+//           const std::string& filename,
+//           raft::neighbors::ivf_flat::index<uint64_t>* index);
+
+}  // namespace raft::runtime::neighbors::ivf_flat

cpp/src/distance/neighbors/ivf_flat_build.cu

@@ -0,0 +1,78 @@
+/*
+ * 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 <raft/neighbors/ivf_flat.cuh>
+#include <raft_runtime/neighbors/ivf_flat.hpp>
+
+namespace raft::runtime::neighbors::ivf_flat {
+
+#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                                        \
+  auto build(raft::device_resources const& handle,                                             \
+             raft::device_matrix_view<const T, uint64_t, row_major> dataset,                   \
+             const raft::neighbors::ivf_flat::index_params& params)                            \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>                                                \
+  {                                                                                            \
+    return raft::neighbors::ivf_flat::build<T, IdxT>(handle, dataset, params);                 \
+  }                                                                                            \
+  auto extend(raft::device_resources const& handle,                                            \
+              const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index,                     \
+              const T* new_vectors,                                                            \
+              const IdxT* new_indices,                                                         \
+              IdxT n_rows)                                                                     \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>                                                \
+  {                                                                                            \
+    return raft::neighbors::ivf_flat::extend<T, IdxT>(                                         \
+      handle, orig_index, new_vectors, new_indices, n_rows);                                   \
+  }                                                                                            \
+                                                                                               \
+  void build(raft::device_resources const& handle,                                             \
+             raft::device_matrix_view<const T, uint64_t, row_major> dataset,                   \
+             const raft::neighbors::ivf_flat::index_params& params,                            \
+             raft::neighbors::ivf_flat::index<T, IdxT>* idx)                                   \
+  {                                                                                            \
+    *idx = raft::neighbors::ivf_flat::build<T, IdxT>(handle, dataset, params);                 \
+  }                                                                                            \
+                                                                                               \
+  void extend(raft::device_resources const& handle,                                            \
+              raft::neighbors::ivf_flat::index<T, IdxT>* idx,                                  \
+              const T* new_vectors,                                                            \
+              const IdxT* new_indices,                                                         \
+              IdxT n_rows)                                                                     \
+  {                                                                                            \
+    raft::neighbors::ivf_flat::extend<T, IdxT>(handle, idx, new_vectors, new_indices, n_rows); \
+  }
+
+RAFT_INST_BUILD_EXTEND(float, uint64_t);
+RAFT_INST_BUILD_EXTEND(int8_t, uint64_t);
+RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t);
+
+#undef RAFT_INST_BUILD_EXTEND
+
+// void save(raft::device_resources const& handle,
+//           const std::string& filename,
+//           const raft::neighbors::ivf_flat::index<T, uint64_t>& index)
+// {
+//   raft::spatial::knn::ivf_flat::detail::save(handle, filename, index);
+// };
+
+// void load(raft::device_resources const& handle,
+//           const std::string& filename,
+//           raft::neighbors::ivf_flat::index<T, uint64_t>* index)
+// {
+//   if (!index) { RAFT_FAIL("Invalid index pointer"); }
+//   *index = raft::spatial::knn::ivf_flat::detail::load<T, uint64_t>(handle, filename);
+// };
+}  // namespace raft::runtime::neighbors::ivf_flat

cpp/src/distance/neighbors/ivf_flat_search.cu

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2022-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 <raft/neighbors/ivf_pq.cuh>
+// #include <raft/neighbors/specializations/detail/ivf_flat_search.cuh>
+#include <raft_runtime/neighbors/ivf_flat.hpp>
+
+namespace raft::runtime::neighbors::ivf_flat {
+
+#define RAFT_INST_SEARCH(T, IdxT)                                         \
+  void search(raft::device_resources const&,                              \
+              const raft::neighbors::ivf_flat::search_params&,            \
+              const raft::neighbors::ivf_flat::index<T, IdxT>&,           \
+              raft::device_matrix_view<const T, IdxT, row_major> queries, \
+              raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,  \
+              raft::device_matrix_view<T, IdxT, row_major> distances,     \
+              uint32_t k);
+
+RAFT_INST_SEARCH(float, uint64_t);
+RAFT_INST_SEARCH(int8_t, uint64_t);
+RAFT_INST_SEARCH(uint8_t, uint64_t);
+
+#undef RAFT_INST_SEARCH
+
+}  // namespace raft::runtime::neighbors::ivf_flat

python/pylibraft/pylibraft/neighbors/CMakeLists.txt

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-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
@@ -23,4 +23,5 @@ rapids_cython_create_modules(
   LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX neighbors_
 )
 
+add_subdirectory(ivf_flat)
 add_subdirectory(ivf_pq)

python/pylibraft/pylibraft/neighbors/ivf_flat/CMakeLists.txt

@@ -0,0 +1,24 @@
+# =============================================================================
+# 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.
+# =============================================================================
+
+# Set the list of Cython files to build
+set(cython_sources ivf_flat.pyx)
+set(linked_libraries raft::raft raft::distance)
+
+# Build all of the Cython targets
+rapids_cython_create_modules(
+  CXX
+  SOURCE_FILES "${cython_sources}"
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX neighbors_ivfflat_
+)

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

@@ -0,0 +1,25 @@
+# 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.
+#
+
+from .ivf_flat import Index, IndexParams, SearchParams, build, extend, search
+
+__all__ = [
+    "Index",
+    "IndexParams",
+    "SearchParams",
+    "build",
+    "extend",
+    "search",
+]

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

@@ -0,0 +1,14 @@
+# 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.
+#