IVF-Flat Python wrappers

.gitignore

@@ -24,6 +24,7 @@ log
 .DS_Store
 dask-worker-space/
 *.egg-info/
+*.bin
 
 ## scikit-build
 _skbuild

cpp/CMakeLists.txt

@@ -370,6 +370,17 @@ 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/ivfflat_build_float_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu
+    src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.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/neighbors/ivf_flat.cuh

@@ -102,8 +102,8 @@ auto build(raft::device_resources const& handle,
  * @tparam matrix_idx_t matrix indexing type
  *
  * @param[in] handle
- * @param[in] params configure the index building
  * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim]
+ * @param[in] params configure the index building
  *
  * @return the constructed ivf-flat index
  */
@@ -119,6 +119,52 @@ auto build(raft::device_resources const& handle,
                                                      static_cast<idx_t>(dataset.extent(1)));
 }
 
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * NB: Currently, the following distance metrics are supported:
+ * - L2Expanded
+ * - L2Unexpanded
+ * - InnerProduct
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace raft::neighbors;
+ *   // use default index parameters
+ *   ivf_flat::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   ivf_flat::index<decltype(dataset::value_type), decltype(dataset::index_type)> *index_ptr;
+ *   ivf_flat::build(handle, dataset, index_params, index_ptr);
+ *   // use default search parameters
+ *   ivf_flat::search_params search_params;
+ *   // search K nearest neighbours for each of the N queries
+ *   ivf_flat::search(handle, index, queries, out_inds, out_dists, search_params, k);
+ * @endcode
+ *
+ * @tparam value_t data element type
+ * @tparam idx_t type of the indices in the source dataset
+ * @tparam int_t precision / type of integral arguments
+ * @tparam matrix_idx_t matrix indexing type
+ *
+ * @param[in] handle
+ * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim]
+ * @param[in] params configure the index building
+ * @param[out] idx pointer to ivf_flat::index
+ *
+ */
+template <typename value_t, typename idx_t>
+void build(raft::device_resources const& handle,
+           raft::device_matrix_view<const value_t, idx_t, row_major> dataset,
+           const index_params& params,
+           raft::neighbors::ivf_flat::index<value_t, idx_t>* idx)
+{
+  *idx = raft::spatial::knn::ivf_flat::detail::build(handle,
+                                                     params,
+                                                     dataset.data_handle(),
+                                                     static_cast<idx_t>(dataset.extent(0)),
+                                                     static_cast<idx_t>(dataset.extent(1)));
+}
+
 /** @} */
 
 /**
@@ -397,24 +443,21 @@ void search(raft::device_resources const& handle,
  * [n_queries, k]
  * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k]
  * @param[in] params configure the search
- * @param[in] k the number of neighbors to find for each query.
  */
-template <typename value_t, typename idx_t, typename int_t>
+template <typename value_t, typename idx_t>
 void search(raft::device_resources const& handle,
             const index<value_t, idx_t>& index,
             raft::device_matrix_view<const value_t, idx_t, row_major> queries,
             raft::device_matrix_view<idx_t, idx_t, row_major> neighbors,
             raft::device_matrix_view<float, idx_t, row_major> distances,
-            const search_params& params,
-            int_t k)
+            const search_params& params)
 {
   RAFT_EXPECTS(
     queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0),
     "Number of rows in output neighbors and distances matrices must equal the number of queries.");
 
-  RAFT_EXPECTS(
-    neighbors.extent(1) == distances.extent(1) && neighbors.extent(1) == static_cast<idx_t>(k),
-    "Number of columns in output neighbors and distances matrices must equal k");
+  RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1),
+               "Number of columns in output neighbors and distances matrices must be equal");
 
   RAFT_EXPECTS(queries.extent(1) == index.dim(),
                "Number of query dimensions should equal number of dimensions in the index.");
@@ -424,7 +467,7 @@ void search(raft::device_resources const& handle,
                 index,
                 queries.data_handle(),
                 static_cast<std::uint32_t>(queries.extent(0)),
-                static_cast<std::uint32_t>(k),
+                static_cast<std::uint32_t>(neighbors.extent(1)),
                 neighbors.data_handle(),
                 distances.data_handle(),
                 nullptr);

cpp/include/raft/neighbors/specializations.cuh

@@ -21,6 +21,7 @@
 
 #include <raft/neighbors/specializations/ball_cover.cuh>
 #include <raft/neighbors/specializations/fused_l2_knn.cuh>
+#include <raft/neighbors/specializations/ivf_flat.cuh>
 #include <raft/neighbors/specializations/ivf_pq.cuh>
 #include <raft/neighbors/specializations/knn.cuh>
 

cpp/include/raft/neighbors/specializations/ivf_flat.cuh

@@ -0,0 +1,59 @@
+/*
+ * 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.cuh>
+
+namespace raft::neighbors::ivf_flat {
+
+#define RAFT_INST(T, IdxT)                                                                   \
+  extern template auto build(raft::device_resources const& handle,                           \
+                             raft::device_matrix_view<const T, uint64_t, row_major> dataset, \
+                             const index_params& params)                                     \
+    ->index<T, IdxT>;                                                                        \
+                                                                                             \
+  extern template auto extend(                                                               \
+    raft::device_resources const& handle,                                                    \
+    const index<T, IdxT>& orig_index,                                                        \
+    raft::device_matrix_view<const T, IdxT, row_major> new_vectors,                          \
+    std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices)                   \
+    ->index<T, IdxT>;                                                                        \
+                                                                                             \
+  extern template void build(raft::device_resources const& handle,                           \
+                             raft::device_matrix_view<const T, uint64_t, row_major> dataset, \
+                             const index_params& params,                                     \
+                             index<T, IdxT>* idx);                                           \
+                                                                                             \
+  extern template void extend(                                                               \
+    raft::device_resources const& handle,                                                    \
+    index<T, IdxT>* idx,                                                                     \
+    raft::device_matrix_view<const T, IdxT, row_major> new_vectors,                          \
+    std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices);                  \
+                                                                                             \
+  extern template void search(raft::device_resources const&,                                 \
+                              const index<T, IdxT>&,                                         \
+                              raft::device_matrix_view<const T, IdxT, row_major>,            \
+                              raft::device_matrix_view<IdxT, IdxT, row_major>,               \
+                              raft::device_matrix_view<float, IdxT, row_major>,              \
+                              search_params const&);
+
+RAFT_INST(float, uint64_t);
+RAFT_INST(int8_t, uint64_t);
+RAFT_INST(uint8_t, uint64_t);
+
+#undef RAFT_INST
+}  // namespace raft::neighbors::ivf_flat

cpp/include/raft_runtime/neighbors/ivf_flat.hpp

@@ -0,0 +1,68 @@
+/*
+ * 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,           \
+              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>;                                     \
+                                                                                     \
+  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,                        \
+              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,        \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices);
+
+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
+
+#define RAFT_INST_SEARCH(T, IdxT)                                 \
+  void search(raft::device_resources const&,                      \
+              const raft::neighbors::ivf_flat::index<T, IdxT>&,   \
+              raft::device_matrix_view<const T, IdxT, row_major>, \
+              raft::device_matrix_view<IdxT, IdxT, row_major>,    \
+              raft::device_matrix_view<float, IdxT, row_major>,   \
+              raft::neighbors::ivf_flat::search_params const&);
+
+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

cpp/src/distance/neighbors/ivf_flat_build.cu

@@ -0,0 +1,62 @@
+/*
+ * 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/specializations/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,             \
+              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>                                        \
+  {                                                                                    \
+    return raft::neighbors::ivf_flat::extend<T, IdxT>(                                 \
+      handle, orig_index, new_vectors, new_indices);                                   \
+  }                                                                                    \
+                                                                                       \
+  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,                          \
+              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::extend<T, IdxT>(handle, idx, new_vectors, new_indices); \
+  }
+
+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
+
+}  // namespace raft::runtime::neighbors::ivf_flat

cpp/src/distance/neighbors/ivf_flat_search.cu

@@ -0,0 +1,40 @@
+/*
+ * 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/specializations/ivf_flat.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& handle,                       \
+              const raft::neighbors::ivf_flat::index<T, IdxT>& index,     \
+              raft::device_matrix_view<const T, IdxT, row_major> queries, \
+              raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,  \
+              raft::device_matrix_view<float, IdxT, row_major> distances, \
+              raft::neighbors::ivf_flat::search_params const& params)     \
+  {                                                                       \
+    raft::neighbors::ivf_flat::search<T, IdxT>(                           \
+      handle, index, queries, neighbors, distances, params);              \
+  }
+
+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