Merge pull request #1730 from rapidsai/release/26.02

Forward-merge release/26.02 into main

Author: GPUtester

c/include/cuvs/neighbors/ivf_pq.h

@@ -23,9 +23,18 @@ extern "C" {
  * @brief A type for specifying how PQ codebooks are created
  *
  */
-enum codebook_gen {  // NOLINT
-  PER_SUBSPACE = 0,  // NOLINT
-  PER_CLUSTER  = 1,  // NOLINT
+enum cuvsIvfPqCodebookGen {
+  CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE = 0,
+  CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER  = 1,
+};
+
+/**
+ * @brief A type for specifying the memory layout of IVF-PQ list data
+ *
+ */
+enum cuvsIvfPqListLayout {
+  CUVS_IVF_PQ_LIST_LAYOUT_FLAT        = 0,
+  CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED = 1,
 };
 
 /**
@@ -80,7 +89,7 @@ struct cuvsIvfPqIndexParams {
    */
   uint32_t pq_dim;
   /** How PQ codebooks are created. */
-  enum codebook_gen codebook_kind;
+  enum cuvsIvfPqCodebookGen codebook_kind;
   /**
    * Apply a random rotation matrix on the input data and queries even if `dim % pq_dim == 0`.
    *
@@ -114,6 +123,14 @@ struct cuvsIvfPqIndexParams {
    * points to train each codebook.
    */
   uint32_t max_train_points_per_pq_code;
+  /**
+   * Memory layout of the IVF-PQ list data.
+   *
+   * - CUVS_IVF_PQ_LIST_LAYOUT_FLAT: Codes are stored contiguously, one vector's codes after another.
+   * - CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED: Codes are interleaved for optimized search performance.
+   *   This is the default and recommended for search workloads.
+   */
+  enum cuvsIvfPqListLayout codes_layout;
 };
 
 typedef struct cuvsIvfPqIndexParams* cuvsIvfPqIndexParams_t;
@@ -294,8 +311,8 @@ cuvsError_t cuvsIvfPqIndexGetCentersPadded(cuvsIvfPqIndex_t index, DLManagedTens
 /**
  * @brief Get the PQ cluster centers
  *
- *   - codebook_gen::PER_SUBSPACE: [pq_dim , pq_len, pq_book_size]
- *   - codebook_gen::PER_CLUSTER:  [n_lists, pq_len, pq_book_size]
+ *   - CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE: [pq_dim , pq_len, pq_book_size]
+ *   - CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER:  [n_lists, pq_len, pq_book_size]
  *
  * @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index
  * @param[out] pq_centers Output tensor that will be populated with a non-owning view of the data
@@ -443,8 +460,8 @@ cuvsError_t cuvsIvfPqBuild(cuvsResources_t res,
  * matrices)
  * @param[in] dim dimensionality of the input data
  * @param[in] pq_centers PQ codebook on device memory with required shape:
- *   - codebook_kind PER_SUBSPACE: [pq_dim, pq_len, pq_book_size]
- *   - codebook_kind PER_CLUSTER:  [n_lists, pq_len, pq_book_size]
+ *   - codebook_kind CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE: [pq_dim, pq_len, pq_book_size]
+ *   - codebook_kind CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER:  [n_lists, pq_len, pq_book_size]
  * @param[in] centers Cluster centers in the original space [n_lists, dim_ext]
  *   where dim_ext = round_up(dim + 1, 8)
  * @param[in] centers_rot Rotated cluster centers [n_lists, rot_dim]

c/src/neighbors/cagra.cpp

@@ -367,7 +367,7 @@ static void _populate_c_ivf_pq_params(cuvsIvfPqParams* c_ivf_pq,
   c_ivf_pq->ivf_pq_build_params->kmeans_trainset_fraction = bp.kmeans_trainset_fraction;
   c_ivf_pq->ivf_pq_build_params->pq_bits = bp.pq_bits;
   c_ivf_pq->ivf_pq_build_params->pq_dim = bp.pq_dim;
-  c_ivf_pq->ivf_pq_build_params->codebook_kind = static_cast<codebook_gen>(bp.codebook_kind);
+  c_ivf_pq->ivf_pq_build_params->codebook_kind = static_cast<cuvsIvfPqCodebookGen>(bp.codebook_kind);
   c_ivf_pq->ivf_pq_build_params->force_random_rotation = bp.force_random_rotation;
   c_ivf_pq->ivf_pq_build_params->conservative_memory_allocation = bp.conservative_memory_allocation;
   c_ivf_pq->ivf_pq_build_params->max_train_points_per_pq_code = bp.max_train_points_per_pq_code;

c/src/neighbors/ivf_pq.cpp

@@ -34,6 +34,7 @@ void convert_c_index_params(cuvsIvfPqIndexParams params, cuvs::neighbors::ivf_pq
   out->force_random_rotation          = params.force_random_rotation;
   out->conservative_memory_allocation = params.conservative_memory_allocation;
   out->max_train_points_per_pq_code   = params.max_train_points_per_pq_code;
+  out->codes_layout = static_cast<cuvs::neighbors::ivf_pq::list_layout>((int)params.codes_layout);
 }
 void convert_c_search_params(cuvsIvfPqSearchParams params,
                              cuvs::neighbors::ivf_pq::search_params* out)
@@ -218,8 +219,16 @@ void _get_list_indices(cuvsIvfPqIndex index,
                        uint32_t label,
                        DLManagedTensor* out_labels)
 {
-  auto index_ptr    = reinterpret_cast<cuvs::neighbors::ivf_pq::index<IdxT>*>(index.addr);
-  cuvs::core::to_dlpack(index_ptr->lists()[label]->indices.view(), out_labels);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::ivf_pq::index<IdxT>*>(index.addr);
+  if (index_ptr->codes_layout() == cuvs::neighbors::ivf_pq::list_layout::FLAT) {
+    auto& list =
+      static_cast<cuvs::neighbors::ivf_pq::list_data_flat<IdxT>&>(*index_ptr->lists()[label]);
+    cuvs::core::to_dlpack(list.indices.view(), out_labels);
+  } else {
+    auto& list = static_cast<cuvs::neighbors::ivf_pq::list_data_interleaved<IdxT>&>(
+      *index_ptr->lists()[label]);
+    cuvs::core::to_dlpack(list.indices.view(), out_labels);
+  }
 }
 }  // namespace
 
@@ -325,10 +334,11 @@ extern "C" cuvsError_t cuvsIvfPqIndexParamsCreate(cuvsIvfPqIndexParams_t* params
                                        .kmeans_trainset_fraction       = 0.5,
                                        .pq_bits                        = 8,
                                        .pq_dim                         = 0,
-                                       .codebook_kind                  = codebook_gen::PER_SUBSPACE,
+                                       .codebook_kind                  = CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE,
                                        .force_random_rotation          = false,
                                        .conservative_memory_allocation = false,
-                                       .max_train_points_per_pq_code   = 256};
+                                       .max_train_points_per_pq_code   = 256,
+                                       .codes_layout                   = CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED};
   });
 }
 

cpp/cmake/patches/faiss-1.13-cuvs-26.02.diff

@@ -1,5 +1,5 @@
 diff --git a/faiss/gpu/impl/CuvsIVFPQ.cu b/faiss/gpu/impl/CuvsIVFPQ.cu
-index 1e2fef225..35b388147 100644
+index 1e2fef225..2ee40da46 100644
 --- a/faiss/gpu/impl/CuvsIVFPQ.cu
 +++ b/faiss/gpu/impl/CuvsIVFPQ.cu
 @@ -129,8 +129,14 @@ void CuvsIVFPQ::updateQuantizer(Index* quantizer) {
@@ -122,7 +122,35 @@ index 1e2fef225..35b388147 100644
      }
 
      setPQCentroids_();
-@@ -520,7 +583,7 @@ void CuvsIVFPQ::setPQCentroids_() {
+@@ -404,10 +467,11 @@ void CuvsIVFPQ::copyInvertedListsFrom(const InvertedLists* ivf) {
+     auto& cuvs_index_lists = cuvs_index->lists();
+
+     // conservative memory alloc for cloning cpu inverted lists
+-    cuvs::neighbors::ivf_pq::list_spec<uint32_t, idx_t> ivf_list_spec{
+-            static_cast<uint32_t>(bitsPerSubQuantizer_),
+-            static_cast<uint32_t>(numSubQuantizers_),
+-            true};
++    cuvs::neighbors::ivf_pq::list_spec_interleaved<uint32_t, idx_t>
++            ivf_list_spec{
++                    static_cast<uint32_t>(bitsPerSubQuantizer_),
++                    static_cast<uint32_t>(numSubQuantizers_),
++                    true};
+
+     for (size_t i = 0; i < nlist; ++i) {
+         size_t listSize = ivf->list_size(i);
+@@ -426,9 +490,9 @@ void CuvsIVFPQ::copyInvertedListsFrom(const InvertedLists* ivf) {
+         // This cuVS list must currently be empty
+         FAISS_ASSERT(getListLength(i) == 0);
+
+-        cuvs::neighbors::ivf::resize_list(
++        cuvs::neighbors::ivf_pq::helpers::resize_list(
+                 raft_handle,
+-                cuvs_index_lists[i],
++                cuvs_index_lists[i],
+                 ivf_list_spec,
+                 static_cast<uint32_t>(listSize),
+                 static_cast<uint32_t>(0));
+@@ -520,7 +587,7 @@ void CuvsIVFPQ::setPQCentroids_() {
      auto stream = resources_->getDefaultStreamCurrentDevice();
 
      raft::copy(

cpp/include/cuvs/neighbors/common.hpp

@@ -711,11 +711,52 @@ template <typename IdxT>
 constexpr static IdxT kInvalidRecord =
   (std::is_signed_v<IdxT> ? IdxT{0} : std::numeric_limits<IdxT>::max()) - 1;
 
+/**
+ * Abstract base class for IVF list data.
+ * This allows polymorphic access to list data regardless of the underlying layout.
+ *
+ * @tparam ValueT The data element type (e.g., uint8_t for PQ codes, float for raw vectors)
+ * @tparam IdxT The index type for source indices
+ * @tparam SizeT The size type
+ *
+ * TODO: Make this struct internal (tracking issue: https://github.com/rapidsai/cuvs/issues/1726)
+ */
+template <typename ValueT, typename IdxT, typename SizeT = uint32_t>
+struct list_base {
+  using value_type = ValueT;
+  using index_type = IdxT;
+  using size_type  = SizeT;
+
+  virtual ~list_base() = default;
+
+  /** Get the raw data pointer. */
+  virtual value_type* data_ptr() noexcept             = 0;
+  virtual const value_type* data_ptr() const noexcept = 0;
+
+  /** Get the indices pointer. */
+  virtual index_type* indices_ptr() noexcept             = 0;
+  virtual const index_type* indices_ptr() const noexcept = 0;
+
+  /** Get the current size (number of records). */
+  virtual size_type get_size() const noexcept = 0;
+
+  /** Set the current size (number of records). */
+  virtual void set_size(size_type new_size) noexcept = 0;
+
+  /** Get the total size of the data array in bytes. */
+  virtual size_t data_byte_size() const noexcept = 0;
+
+  /** Get the capacity (number of indices that can be stored). */
+  virtual size_type indices_capacity() const noexcept = 0;
+};
+
 /** The data for a single IVF list. */
 template <template <typename, typename...> typename SpecT,
           typename SizeT,
           typename... SpecExtraArgs>
-struct list {
+struct list : public list_base<typename SpecT<SizeT, SpecExtraArgs...>::value_type,
+                               typename SpecT<SizeT, SpecExtraArgs...>::index_type,
+                               SizeT> {
   using size_type    = SizeT;
   using spec_type    = SpecT<size_type, SpecExtraArgs...>;
   using value_type   = typename spec_type::value_type;
@@ -731,6 +772,18 @@ struct list {
 
   /** Allocate a new list capable of holding at least `n_rows` data records and indices. */
   list(raft::resources const& res, const spec_type& spec, size_type n_rows);
+
+  value_type* data_ptr() noexcept override { return data.data_handle(); }
+  const value_type* data_ptr() const noexcept override { return data.data_handle(); }
+
+  index_type* indices_ptr() noexcept override { return indices.data_handle(); }
+  const index_type* indices_ptr() const noexcept override { return indices.data_handle(); }
+
+  size_type get_size() const noexcept override { return size.load(); }
+  void set_size(size_type new_size) noexcept override { size.store(new_size); }
+
+  size_t data_byte_size() const noexcept override { return data.size() * sizeof(value_type); }
+  size_type indices_capacity() const noexcept override { return indices.extent(0); }
 };
 
 template <typename ListT, class T = void>
@@ -755,6 +808,10 @@ using enable_if_valid_list_t = typename enable_if_valid_list<ListT, T>::type;
 /**
  * Resize a list by the given id, so that it can contain the given number of records;
  * copy the data if necessary.
+ *
+ * @note This is an internal function that requires the concrete list type.
+ *       For IVF-PQ indexes, prefer using the helper functions in
+ *       `cuvs::neighbors::ivf_pq::helpers::resize_list` which handle type casting internally.
  */
 template <typename ListT>
 void resize_list(raft::resources const& res,
@@ -763,13 +820,23 @@ void resize_list(raft::resources const& res,
                  typename ListT::size_type new_used_size,
                  typename ListT::size_type old_used_size);
 
+/**
+ * Serialize a list to an output stream.
+ *
+ * @note This function requires the concrete list type (not the base class) because:
+ *       1. It needs access to the spec_type to determine the data layout for serialization
+ *       2. The serialized format depends on the spec's make_list_extents() method
+ *       When calling from code that only has a base class pointer, use std::static_pointer_cast
+ *       to obtain the typed pointer first.
+ */
 template <typename ListT>
 enable_if_valid_list_t<ListT> serialize_list(
   const raft::resources& handle,
   std::ostream& os,
   const ListT& ld,
   const typename ListT::spec_type& store_spec,
   std::optional<typename ListT::size_type> size_override = std::nullopt);
+
 template <typename ListT>
 enable_if_valid_list_t<ListT> serialize_list(
   const raft::resources& handle,

cpp/include/cuvs/neighbors/ivf_flat.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -2933,17 +2933,17 @@ void reset_index(const raft::resources& res, index<uint8_t, int64_t>* index);
  *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, float, int64_t>{
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
- * @param[inout] index pointer to IVF-PQ index
+ * @param[inout] index pointer to IVF-Flat index
  */
 void recompute_internal_state(const raft::resources& res, index<float, int64_t>* index);
 
@@ -2961,17 +2961,17 @@ void recompute_internal_state(const raft::resources& res, index<float, int64_t>*
  *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, half, int64_t>{
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
- * @param[inout] index pointer to IVF-PQ index
+ * @param[inout] index pointer to IVF-Flat index
  */
 void recompute_internal_state(const raft::resources& res, index<half, int64_t>* index);
 
@@ -2989,17 +2989,17 @@ void recompute_internal_state(const raft::resources& res, index<half, int64_t>*
  *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, int8_t, int64_t>{
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
- * @param[inout] index pointer to IVF-PQ index
+ * @param[inout] index pointer to IVF-Flat index
  */
 void recompute_internal_state(const raft::resources& res, index<int8_t, int64_t>* index);
 
@@ -3018,9 +3018,9 @@ void recompute_internal_state(const raft::resources& res, index<int8_t, int64_t>
  *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
  *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
- *     index->dim(), index->conservative_memory_allocation()};
+ *     index.dim(), index.conservative_memory_allocation()};
  *   uint32_t new_size = 5;
- *   ivf::resize_list(res, list, spec, new_size, 0);
+ *   ivf::resize_list(res, index.lists()[0], spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
  *   ivf_flat::helpers::recompute_internal_state(res, index);