Iterative CAGRA-Q

c/include/cuvs/core/dataset.h

@@ -20,7 +20,8 @@ extern "C" {
  */
 typedef enum {
   CUVS_DATASET_LAYOUT_STANDARD = 0,
-  CUVS_DATASET_LAYOUT_PADDED   = 1
+  CUVS_DATASET_LAYOUT_PADDED   = 1,
+  CUVS_DATASET_LAYOUT_PQ = 2
 } cuvsDatasetLayout_t;
 
 /**
@@ -48,6 +49,9 @@ typedef struct {
 } cuvsDataset;
 typedef cuvsDataset* cuvsDataset_t;
 
+struct cuvsCagraCompressionParams;
+typedef struct cuvsCagraCompressionParams* cuvsCagraCompressionParams_t;
+
 /**
  * @brief Create an empty owning dataset handle.
  *
@@ -72,6 +76,14 @@ CUVS_EXPORT cuvsError_t cuvsDatasetMakePadded(cuvsResources_t res,
                                               cuvsDatasetMemType_t target_mem_type,
                                               cuvsDataset_t* padded_dataset);
 
+/**
+ * @brief Compress a dense dataset into a device VPQ dataset.
+ */
+CUVS_EXPORT cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res,
+                                           cuvsCagraCompressionParams_t params,
+                                           cuvsDataset_t dataset,
+                                           cuvsDataset_t* vpq_dataset);
+
 /**
  * @brief Create a non-owning padded dataset view from a host- or device-resident tensor.
  *

c/include/cuvs/neighbors/cagra.h

@@ -118,8 +118,6 @@ struct cuvsCagraCompressionParams {
   double pq_kmeans_trainset_fraction;
 };
 
-typedef struct cuvsCagraCompressionParams* cuvsCagraCompressionParams_t;
-
 struct cuvsIvfPqParams {
   cuvsIvfPqIndexParams_t ivf_pq_build_params;
   cuvsIvfPqSearchParams_t ivf_pq_search_params;
@@ -685,6 +683,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res,
  * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
  * @endcode
  *
+ * A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakeVpq` builds an iterative CAGRA-Q
+ * index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the
+ * VPQ dataset must outlive the index because the index stores a non-owning view.
+ *
  * @param[in] res cuvsResources_t opaque C handle
  * @param[in] params cuvsCagraIndexParams_t used to build CAGRA index
  * @param[in] dataset cuvsDataset_t training dataset or dataset view

c/src/neighbors/cagra.cpp

@@ -30,6 +30,7 @@
 #include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/common.h>
 #include <cuvs/neighbors/cagra.hpp>
+#include <cuvs/preprocessing/quantize/pq.hpp>
 #include "../core/exceptions.hpp"
 #include "../core/interop.hpp"
 
@@ -52,7 +53,13 @@ struct cuvs_cagra_c_api_index_lifetime_holder {
 /** Owns how to delete co-located index storage; `cuvsCagraIndex::addr` points here. */
 struct sg_cagra_c_api_index_box {
   void* index_ptr;
-  enum class dataset_layout : uint8_t { device_padded, device_standard, host_padded, host_standard } layout;
+  enum class dataset_layout : uint8_t {
+    device_padded,
+    device_standard,
+    device_vpq,
+    host_padded,
+    host_standard
+  } layout;
   cuvs::neighbors::c_api::detail::owner_record owner_rec;
 };
 
@@ -63,6 +70,8 @@ constexpr auto sg_cagra_index_layout_from_view()
     return sg_cagra_c_api_index_box::dataset_layout::device_standard;
   } else if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v<DatasetViewT>) {
     return sg_cagra_c_api_index_box::dataset_layout::device_padded;
+  } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v<DatasetViewT>) {
+    return sg_cagra_c_api_index_box::dataset_layout::device_vpq;
   } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v<DatasetViewT>) {
     return sg_cagra_c_api_index_box::dataset_layout::host_standard;
   } else {
@@ -100,6 +109,13 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box,
       fn(*idx);
       break;
     }
+    case sg_cagra_c_api_index_box::dataset_layout::device_vpq: {
+      using index_t = cuvs::neighbors::cagra::
+        index<T, IdxT, cuvs::neighbors::device_vpq_dataset_view<half, int64_t>>;
+      auto* idx = reinterpret_cast<index_t*>(box->index_ptr);
+      fn(*idx);
+      break;
+    }
     case sg_cagra_c_api_index_box::dataset_layout::host_standard: {
       if constexpr (AllowHost) {
         auto* idx =
@@ -374,6 +390,77 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn)
   }
 }
 
+using device_vpq_owner_t = cuvs::neighbors::device_vpq_dataset<half, int64_t>;
+using device_vpq_view_t  = cuvs::neighbors::device_vpq_dataset_view<half, int64_t>;
+
+static void bind_vpq_owner_to_dataset(std::unique_ptr<device_vpq_owner_t> owner,
+                                      cuvsDataset_t* output)
+{
+  RAFT_EXPECTS(output != nullptr, "VPQ output dataset pointer must not be null");
+  auto* out         = new cuvsDataset{};
+  out->addr         = reinterpret_cast<uintptr_t>(owner.release());
+  out->destroy_addr = &destroy_typed_addr<device_vpq_owner_t>;
+  out->dtype.code   = kDLFloat;
+  out->dtype.bits   = 32;
+  out->dtype.lanes  = 1;
+  out->mem_type     = CUVS_DATASET_MEM_TYPE_DEVICE;
+  out->layout       = CUVS_DATASET_LAYOUT_PQ;
+  out->is_owning    = true;
+  *output           = out;
+}
+
+static auto make_cpp_vpq_params(cuvsCagraCompressionParams const& params)
+  -> cuvs::neighbors::vpq_params
+{
+  auto out                            = cuvs::neighbors::vpq_params{};
+  out.pq_bits                         = params.pq_bits;
+  out.pq_dim                          = params.pq_dim;
+  out.vq_n_centers                    = params.vq_n_centers;
+  out.kmeans_n_iters                  = params.kmeans_n_iters;
+  out.vq_kmeans_trainset_fraction     = params.vq_kmeans_trainset_fraction;
+  out.pq_kmeans_trainset_fraction     = params.pq_kmeans_trainset_fraction;
+  return out;
+}
+
+template <typename T>
+static auto make_vpq_from_dense_dataset(raft::resources* res_ptr,
+                                        cuvsCagraCompressionParams const& params,
+                                        cuvsDataset_t dataset)
+  -> std::unique_ptr<device_vpq_owner_t>
+{
+  RAFT_EXPECTS(dataset->layout == CUVS_DATASET_LAYOUT_STANDARD ||
+                 dataset->layout == CUVS_DATASET_LAYOUT_PADDED,
+               "cuvsDatasetMakeVpq: source dataset must have STANDARD or PADDED layout");
+  auto cpp_params = make_cpp_vpq_params(params);
+  std::unique_ptr<device_vpq_owner_t> owner;
+  auto make = [&](auto const& view) {
+    owner = std::make_unique<device_vpq_owner_t>(
+      cuvs::preprocessing::quantize::pq::make_vpq_dataset(*res_ptr, cpp_params, view));
+  };
+
+  const bool padded = dataset->layout == CUVS_DATASET_LAYOUT_PADDED;
+  if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE) {
+    if (padded) {
+      with_dataset_view<cuvs::neighbors::device_padded_dataset<T, int64_t>,
+                        cuvs::neighbors::device_padded_dataset_view<T, int64_t>>(dataset, make);
+    } else {
+      with_dataset_view<cuvs::neighbors::device_standard_dataset<T, int64_t>,
+                        cuvs::neighbors::device_standard_dataset_view<T, int64_t>>(dataset, make);
+    }
+  } else if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_HOST) {
+    if (padded) {
+      with_dataset_view<cuvs::neighbors::host_padded_dataset<T, int64_t>,
+                        cuvs::neighbors::host_padded_dataset_view<T, int64_t>>(dataset, make);
+    } else {
+      with_dataset_view<cuvs::neighbors::host_standard_dataset<T, int64_t>,
+                        cuvs::neighbors::host_standard_dataset_view<T, int64_t>>(dataset, make);
+    }
+  } else {
+    RAFT_FAIL("cuvsDatasetMakeVpq: invalid source dataset memory type");
+  }
+  return owner;
+}
+
 template <typename T>
 static void make_device_padded_dataset(raft::resources* res_ptr,
                                        DLManagedTensor* dataset_tensor,
@@ -552,13 +639,18 @@ static void update_dataset(raft::resources* res_ptr,
       "cuvsCagraUpdateDataset: null index handle",
       "cuvsCagraUpdateDataset: host index layout is allowed for this operation",
       [&](auto& idx) {
-        auto padded_idx =
-          cuvs::neighbors::cagra::update_dataset(*res_ptr, std::move(idx), padded_view);
-        auto* holder =
-          new cuvs_cagra_c_api_index_lifetime_holder<T, view_t>{std::move(padded_idx)};
-        destroy_sg_cagra_c_api_box(index->addr);
-        index->addr = 0;
-        bind_index_lifetime_holder_to_C_index<T, view_t>(index, index->dtype, holder);
+        using index_dataset_view_t = std::remove_cvref_t<decltype(idx.dataset())>;
+        if constexpr (cuvs::neighbors::is_vpq_dataset_view_v<index_dataset_view_t>) {
+          RAFT_FAIL("cuvsCagraUpdateDataset: replacing a VPQ dataset is not supported");
+        } else {
+          auto padded_idx =
+            cuvs::neighbors::cagra::update_dataset(*res_ptr, std::move(idx), padded_view);
+          auto* holder =
+            new cuvs_cagra_c_api_index_lifetime_holder<T, view_t>{std::move(padded_idx)};
+          destroy_sg_cagra_c_api_box(index->addr);
+          index->addr = 0;
+          bind_index_lifetime_holder_to_C_index<T, view_t>(index, index->dtype, holder);
+        }
       });
   });
 }
@@ -920,13 +1012,19 @@ void _serialize(cuvsResources_t res, const char *filename,
                       : "cuvsCagraSerializeGraph: null index handle";
   with_index_by_layout<T, uint32_t,
                        true>(box, null_handle_err, "", [&](auto &idx) {
-    if (include_dataset) {
-      RAFT_EXPECTS(
-          idx.dataset().n_rows() > 0,
-          "cuvsCagraSerializeGraphAndDataset: index has no attached dataset");
+    using index_dataset_view_t = std::remove_cvref_t<decltype(idx.dataset())>;
+    if constexpr (cuvs::neighbors::is_vpq_dataset_view_v<index_dataset_view_t>) {
+      RAFT_FAIL(
+        "CAGRA index serialization is not supported for VPQ indices");
+    } else {
+      if (include_dataset) {
+        RAFT_EXPECTS(
+            idx.dataset().n_rows() > 0,
+            "cuvsCagraSerializeGraphAndDataset: index has no attached dataset");
+      }
+      cuvs::neighbors::cagra::serialize(
+        *res_ptr, std::string(filename), idx, include_dataset);
     }
-    cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx,
-                                      include_dataset);
   });
 }
 
@@ -972,7 +1070,7 @@ static auto read_serialized_header(cuvsResources_t res, const char *filename)
       "serialization version mismatch, expected %d, got %d",
       cuvs::neighbors::cagra::cagra_serialization_version, version);
   using kind = cuvs::neighbors::cagra::serialized_dataset_kind;
-  RAFT_EXPECTS(dataset_kind_raw <= static_cast<std::uint32_t>(kind::host_standard),
+  RAFT_EXPECTS(dataset_kind_raw <= static_cast<std::uint32_t>(kind::device_vpq_f16),
                "Invalid serialized dataset kind %u in file %s",
                dataset_kind_raw, filename);
   return {output_dtype, static_cast<kind>(dataset_kind_raw)};
@@ -1019,6 +1117,12 @@ void dispatch_serialized_dataset_kind(
       fn.template operator()<
           cuvs::neighbors::device_padded_dataset_view<T, int64_t>>();
       break;
+    case serialized_kind::device_vpq_f16:
+      // A recognised file the C API has no index layout for, as opposed to an unreadable one.
+      // cuvsDatasetLayout_t covers standard and padded only, and every C entry point dispatches
+      // on that layout, so there is nothing here to hand a VPQ index to yet.
+      RAFT_FAIL("File holds a VPQ-compressed (CAGRA-Q) dataset, which the C API has no dataset "
+                "layout for; load it through the C++ API");
   }
 }
 
@@ -1083,8 +1187,13 @@ void _serialize_to_hnswlib(cuvsResources_t res, const char *filename,
       box, "cuvsCagraSerializeToHnswlib: null index handle",
       "cuvsCagraSerializeToHnswlib: host indices are allowed",
       [&](auto &idx) {
-        cuvs::neighbors::cagra::serialize_to_hnswlib(
-            *res_ptr, std::string(filename), idx);
+        using index_dataset_view_t = std::remove_cvref_t<decltype(idx.dataset())>;
+        if constexpr (cuvs::neighbors::is_vpq_dataset_view_v<index_dataset_view_t>) {
+          RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for VPQ indices");
+        } else {
+          cuvs::neighbors::cagra::serialize_to_hnswlib(
+              *res_ptr, std::string(filename), idx);
+        }
       });
 }
 template <typename T>
@@ -1159,7 +1268,14 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset)
     box,
     "cuvsCagraIndexGetDataset: null index handle",
     "cuvsCagraIndexGetDataset: host indices are allowed",
-    [&](auto& idx) { cuvs::core::to_dlpack(idx.dataset().view(), dataset); });
+    [&](auto& idx) {
+      using index_dataset_view_t = std::remove_cvref_t<decltype(idx.dataset())>;
+      if constexpr (cuvs::neighbors::is_vpq_dataset_view_v<index_dataset_view_t>) {
+        RAFT_FAIL("cuvsCagraIndexGetDataset does not expose VPQ datasets as dense DLPack tensors");
+      } else {
+        cuvs::core::to_dlpack(idx.dataset().view(), dataset);
+      }
+    });
 }
 
 template <typename T, typename IdxT>
@@ -1540,6 +1656,36 @@ extern "C" cuvsError_t cuvsDatasetMakeStandardView(cuvsResources_t res,
   });
 }
 
+extern "C" cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res,
+                                          cuvsCagraCompressionParams_t params,
+                                          cuvsDataset_t dataset,
+                                          cuvsDataset_t* vpq_dataset)
+{
+  return cuvs::core::translate_exceptions([=] {
+    RAFT_EXPECTS(params != nullptr, "cuvsDatasetMakeVpq: null compression params");
+    RAFT_EXPECTS(dataset != nullptr && dataset->addr != 0,
+                 "cuvsDatasetMakeVpq: null source dataset");
+    RAFT_EXPECTS(vpq_dataset != nullptr, "cuvsDatasetMakeVpq: null output dataset");
+    *vpq_dataset = nullptr;
+    auto* res_ptr = reinterpret_cast<raft::resources*>(res);
+    std::unique_ptr<device_vpq_owner_t> owner;
+    if (dataset->dtype.code == kDLFloat && dataset->dtype.bits == 32) {
+      owner = make_vpq_from_dense_dataset<float>(res_ptr, *params, dataset);
+    } else if (dataset->dtype.code == kDLFloat && dataset->dtype.bits == 16) {
+      owner = make_vpq_from_dense_dataset<half>(res_ptr, *params, dataset);
+    } else if (dataset->dtype.code == kDLInt && dataset->dtype.bits == 8) {
+      owner = make_vpq_from_dense_dataset<int8_t>(res_ptr, *params, dataset);
+    } else if (dataset->dtype.code == kDLUInt && dataset->dtype.bits == 8) {
+      owner = make_vpq_from_dense_dataset<uint8_t>(res_ptr, *params, dataset);
+    } else {
+      RAFT_FAIL("cuvsDatasetMakeVpq: unsupported source dtype: code=%d, bits=%d",
+                dataset->dtype.code,
+                dataset->dtype.bits);
+    }
+    bind_vpq_owner_to_dataset(std::move(owner), vpq_dataset);
+  });
+}
+
 static cuvsError_t dispatch_update_dataset(cuvsResources_t res,
                                            cuvsDataset_t device_padded_dataset,
                                            cuvsCagraIndex_t index)
@@ -1671,7 +1817,15 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
     index->addr  = 0;
     index->dtype = dtype;
 
-    if (dtype.code == kDLFloat && dtype.bits == 32) {
+    if (dataset->layout == CUVS_DATASET_LAYOUT_PQ) {
+      RAFT_EXPECTS(dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE,
+                   "cuvsCagraBuild: VPQ dataset must be device-resident");
+      RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32,
+                   "cuvsCagraBuild: VPQ dataset query dtype must be float32");
+      with_dataset_view<device_vpq_owner_t, device_vpq_view_t>(dataset, [&](auto const& view) {
+        build_index_from_dataset_view<float>(res_ptr, params, view, index);
+      });
+    } else if (dtype.code == kDLFloat && dtype.bits == 32) {
       build_dispatch_on_mem_type_and_layout<float>(res_ptr, params, dataset, index);
     } else if (dtype.code == kDLFloat && dtype.bits == 16) {
       build_dispatch_on_mem_type_and_layout<half>(res_ptr, params, dataset, index);
@@ -1769,10 +1923,12 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
     auto index = *index_c_ptr;
     auto* box  = reinterpret_cast<sg_cagra_c_api_index_box*>(index.addr);
     RAFT_EXPECTS(box != nullptr, "cuvsCagraSearch: null index handle");
-    RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded,
-                 "cuvsCagraSearch: index must be device-padded. For standard indices, call "
+    RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded ||
+                   box->layout == sg_cagra_c_api_index_box::dataset_layout::device_vpq,
+                 "cuvsCagraSearch: index must be device-padded or VPQ. For standard indices, call "
                  "cuvsCagraUpdateDataset first.");
-    RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries");
+    RAFT_EXPECTS(queries.dtype.code == index.dtype.code && queries.dtype.bits == index.dtype.bits,
+                 "type mismatch between index and queries");
 
     if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) {
       _search<float>(

c/tests/CMakeLists.txt

@@ -89,7 +89,9 @@ ConfigureTest(NAME IVF_FLAT_C_TEST PATH neighbors/run_ivf_flat_c.c neighbors/ann
 ConfigureTest(NAME IVF_PQ_C_TEST PATH neighbors/run_ivf_pq_c.c neighbors/ann_ivf_pq_c.cu)
 ConfigureTest(NAME IVF_SQ_C_TEST PATH neighbors/run_ivf_sq_c.c neighbors/ann_ivf_sq_c.cu)
 ConfigureTest(NAME CAGRA_C_TEST PATH neighbors/ann_cagra_c.cu)
-ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu)
+if(BUILD_MG_ALGOS)
+  ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu)
+endif()
 ConfigureTest(
   NAME ALL_NEIGHBORS_C_TEST PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu
 )