Migrate IVF-PQ from RAFT to cuVS

cpp/CMakeLists.txt

@@ -216,6 +216,7 @@ add_library(
   src/neighbors/ivf_flat/ivf_flat_serialize_int8_t_int64_t.cpp
   src/neighbors/ivf_flat/ivf_flat_serialize_uint8_t_int64_t.cpp
   src/neighbors/ivf_pq_index.cpp
+  src/neighbors/ivf_pq/ivf_pq_build_common.cu
   src/neighbors/ivf_pq/ivf_pq_build_float_int64_t.cu
   src/neighbors/ivf_pq/ivf_pq_build_int8_t_int64_t.cu
   src/neighbors/ivf_pq/ivf_pq_build_uint8_t_int64_t.cu

cpp/include/cuvs/core/nvtx.hpp

@@ -18,8 +18,8 @@
 #include <raft/core/nvtx.hpp>
 
 namespace cuvs::common::nvtx::domain {
-    /** @brief This NVTX domain is supposed to be used within cuvs.  */
-    struct cuvs {
-        static constexpr const char* name = "cuvs";
-    };
+/** @brief This NVTX domain is supposed to be used within cuvs.  */
+struct cuvs {
+  static constexpr const char* name = "cuvs";
+};
 };  // namespace cuvs::common::nvtx::domain

cpp/include/cuvs/neighbors/detail/ivf_pq_helpers.hpp

@@ -0,0 +1,268 @@
+/*
+ * raft::copyright (c) 2022-2024, 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 raft::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 <cuvs/core/nvtx.hpp>
+#include <cuvs/distance/distance_types.hpp>
+#include <cuvs/neighbors/ivf_list.hpp>
+#include <cuvs/neighbors/ivf_pq.hpp>
+
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/logger.hpp>
+#include <raft/core/operators.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/random/rng_state.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_uvector.hpp>
+#include <rmm/mr/device/managed_memory_resource.hpp>
+
+#include <memory>
+#include <variant>
+
+namespace cuvs::neighbors::ivf_pq::helpers {
+
+namespace codepacker {
+/**
+ * Unpack flat PQ codes from an existing list by the given offset.
+ *
+ * @param[out] codes flat PQ codes, one code per byte [n_rows, pq_dim]
+ * @param[in] list_data the packed ivf::list data.
+ * @param[in] offset_or_indices how many records in the list to skip or the exact indices.
+ * @param[in] pq_bits codebook size (1 << pq_bits)
+ * @param[in] stream
+ */
+void unpack_list_data(raft::resources const& res,
+                      const index<int64_t>& index,
+                      raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_codes,
+                      uint32_t label,
+                      std::variant<uint32_t, const uint32_t*> offset_or_indices);
+/**
+ * Write flat PQ codes into an existing list by the given offset.
+ *
+ * NB: no memory allocation happens here; the list must fit the data (offset + n_rows).
+ *
+ * @param[out] list_data the packed ivf::list data.
+ * @param[in] codes flat PQ codes, one code per byte [n_rows, pq_dim]
+ * @param[in] offset_or_indices how many records in the list to skip or the exact indices.
+ * @param[in] pq_bits codebook size (1 << pq_bits)
+ * @param[in] stream
+ */
+void pack_list_data(raft::resources const& res,
+                    index<int64_t>* index,
+                    raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_codes,
+                    uint32_t label,
+                    std::variant<uint32_t, const uint32_t*> offset_or_indices);
+
+};  // namespace codepacker
+
+/**
+ * @brief Fill-in a random orthogonal transformation matrix.
+ *
+ * @param handle
+ * @param force_random_rotation
+ * @param n_rows
+ * @param n_cols
+ * @param[out] rotation_matrix device pointer to a row-major matrix of size [n_rows, n_cols].
+ * @param rng random number generator state
+ */
+void make_rotation_matrix(raft::resources const& handle,
+                          bool force_random_rotation,
+                          uint32_t n_rows,
+                          uint32_t n_cols,
+                          float* rotation_matrix,
+                          raft::random::RngState rng = raft::random::RngState(7ULL));
+
+void set_centers(raft::resources const& handle,
+                 index<int64_t>* index,
+                 const float* cluster_centers);
+
+/**
+ * @brief Decode `n_take` consecutive records of a single list (cluster) in the compressed index
+ * starting at given `offset`.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   // We will reconstruct the fourth cluster
+ *   uint32_t label = 3;
+ *   // Get the list size
+ *   uint32_t list_size = 0;
+ *   raft::copy(&list_size, index.list_sizes().data_handle() + label, 1,
+ *   resource::get_cuda_stream(res)); resource::sync_stream(res);
+ *   // allocate the buffer for the output
+ *   auto decoded_vectors = raft::make_device_matrix<float>(res, list_size, index.dim());
+ *   // decode the whole list
+ *   ivf_pq::helpers::reconstruct_list_data(res, index, decoded_vectors.view(), label, 0);
+ * @endcode
+ *
+ * @param[in] res
+ * @param[in] index
+ * @param[out] out_vectors
+ *   the destination buffer [n_take, index.dim()].
+ *   The length `n_take` defines how many records to reconstruct,
+ *   it must be smaller than the list size.
+ * @param[in] label
+ *   The id of the list (cluster) to decode.
+ * @param[in] offset
+ *   How many records in the list to skip.
+ */
+void reconstruct_list_data(raft::resources const& res,
+                           const index<int64_t>& index,
+                           raft::device_matrix_view<float, uint32_t, raft::row_major> out_vectors,
+                           uint32_t label,
+                           std::variant<uint32_t, const uint32_t*> offset_or_indices);
+
+void reconstruct_list_data(raft::resources const& res,
+                           const index<int64_t>& index,
+                           raft::device_matrix_view<int8_t, uint32_t, raft::row_major> out_vectors,
+                           uint32_t label,
+                           std::variant<uint32_t, const uint32_t*> offset_or_indices);
+void reconstruct_list_data(raft::resources const& res,
+                           const index<int64_t>& index,
+                           raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_vectors,
+                           uint32_t label,
+                           std::variant<uint32_t, const uint32_t*> offset_or_indices);
+
+/**
+ * Write flat PQ codes into an existing list by the given offset. The input codes of a single vector
+ * are contiguous (not expanded to one code per byte).
+ *
+ * NB: no memory allocation happens here; the list must fit the data (offset + n_rows records).
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   raft::resources res;
+ *   auto list_data  = index.lists()[label]->data.view();
+ *   // allocate the buffer for the input codes
+ *   auto codes = raft::make_device_matrix<uint8_t>(
+ *     res, n_rows, raft::ceildiv(index.pq_dim() * index.pq_bits(), 8));
+ *   ... prepare compressed vectors to pack into the list in codes ...
+ *   // write codes into the list starting from the 42nd position. If the current size of the list
+ *   // is greater than 42, this will overwrite the codes starting at this offset.
+ *   ivf_pq::helpers::codepacker::pack_contiguous(
+ *     res, codes.data_handle(), n_rows, index.pq_dim(), index.pq_bits(), 42, list_data);
+ * @endcode
+ *
+ * @param[in] res raft resource
+ * @param[inout] index
+ * @param[in] new_codes flat PQ codes, [n_vec, ceildiv(pq_dim * pq_bits, 8)]
+ * @param[in] n_rows number of records
+ * @param[in] offset_or_indices how many records in the list to skip or the exact indices.
+ */
+void pack_contiguous_list_data(raft::resources const& res,
+                               index<int64_t>* index,
+                               const uint8_t* new_codes,
+                               uint32_t n_rows,
+                               uint32_t label,
+                               std::variant<uint32_t, const uint32_t*> offset_or_indices);
+
+/**
+ * @brief Extend one list of the index in-place, by the list label, skipping the classification and
+ * encoding steps.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   // We will extend the fourth cluster
+ *   uint32_t label = 3;
+ *   // We will fill 4 new vectors
+ *   uint32_t n_vec = 4;
+ *   // Indices of the new vectors
+ *   auto indices = raft::make_device_vector<uint32_t>(res, n_vec);
+ *   ... fill the indices ...
+ *   auto new_codes = raft::make_device_matrix<uint8_t, uint32_t, row_major> new_codes(
+ *       res, n_vec, index.pq_dim());
+ *   ... fill codes ...
+ *   // extend list with new codes
+ *   ivf_pq::helpers::extend_list_with_codes(
+ *       res, &index, codes.view(), indices.view(), label);
+ * @endcode
+ *
+ * @param[in] res
+ * @param[inout] index
+ * @param[in] new_codes flat PQ codes, one code per byte [n_rows, index.pq_dim()]
+ * @param[in] new_indices source indices [n_rows]
+ * @param[in] label the id of the target list (cluster).
+ */
+void extend_list_with_codes(
+  raft::resources const& res,
+  index<int64_t>* index,
+  raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_codes,
+  raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
+  uint32_t label);
+
+/**
+ * @brief Extend one list of the index in-place, by the list label, skipping the classification
+ * step.
+ *
+ *  Usage example:
+ * @code{.cpp}
+ *   // We will extend the fourth cluster
+ *   uint32_t label = 3;
+ *   // We will extend with 4 new vectors
+ *   uint32_t n_vec = 4;
+ *   // Indices of the new vectors
+ *   auto indices = raft::make_device_vector<uint32_t>(res, n_vec);
+ *   ... fill the indices ...
+ *   auto new_vectors = raft::make_device_matrix<float, uint32_t, row_major> new_codes(
+ *       res, n_vec, index.dim());
+ *   ... fill vectors ...
+ *   // extend list with new vectors
+ *   ivf_pq::helpers::extend_list(
+ *       res, &index, new_vectors.view(), indices.view(), label);
+ * @endcode
+ *
+ *
+ * @param[in] res
+ * @param[inout] index
+ * @param[in] new_vectors data to encode [n_rows, index.dim()]
+ * @param[in] new_indices source indices [n_rows]
+ * @param[in] label the id of the target list (cluster).
+ */
+void extend_list(raft::resources const& res,
+                 index<int64_t>* index,
+                 raft::device_matrix_view<const float, uint32_t, raft::row_major> new_vectors,
+                 raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
+                 uint32_t label);
+void extend_list(raft::resources const& res,
+                 index<int64_t>* index,
+                 raft::device_matrix_view<const int8_t, uint32_t, raft::row_major> new_vectors,
+                 raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
+                 uint32_t label);
+void extend_list(raft::resources const& res,
+                 index<int64_t>* index,
+                 raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_vectors,
+                 raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
+                 uint32_t label);
+/**
+ * @brief Remove all data from a single list (cluster) in the index.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   // We will erase the fourth cluster (label = 3)
+ *   ivf_pq::helpers::erase_list(res, &index, 3);
+ * @endcode
+ *
+ *
+ * @param[in] res
+ * @param[inout] index
+ * @param[in] label the id of the target list (cluster).
+ */
+void erase_list(raft::resources const& res, index<int64_t>* index, uint32_t label);
+
+}  // namespace cuvs::neighbors::ivf_pq::helpers