diff --git a/c/include/cuvs/core/dataset.h b/c/include/cuvs/core/dataset.h index 78d3547495..0d28a9977c 100644 --- a/c/include/cuvs/core/dataset.h +++ b/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_VPQ = 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. * diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 350711d069..66e222d747 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/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; @@ -655,6 +653,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); * @endcode * + * A `CUVS_DATASET_LAYOUT_VPQ` 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 diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 99e456e23c..4eb83cb7bb 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #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) { return sg_cagra_c_api_index_box::dataset_layout::device_padded; + } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::device_vpq; } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { 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>; + auto* idx = reinterpret_cast(box->index_ptr); + fn(*idx); + break; + } case sg_cagra_c_api_index_box::dataset_layout::host_standard: { if constexpr (AllowHost) { auto* idx = @@ -369,6 +385,77 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn) } } +using device_vpq_owner_t = cuvs::neighbors::device_vpq_dataset; +using device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; + +static void bind_vpq_owner_to_dataset(std::unique_ptr owner, + cuvsDataset_t* output) +{ + RAFT_EXPECTS(output != nullptr, "VPQ output dataset pointer must not be null"); + auto* out = new cuvsDataset{}; + out->addr = reinterpret_cast(owner.release()); + out->destroy_addr = &destroy_typed_addr; + 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_VPQ; + 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 +static auto make_vpq_from_dense_dataset(raft::resources* res_ptr, + cuvsCagraCompressionParams const& params, + cuvsDataset_t dataset) + -> std::unique_ptr +{ + 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 owner; + auto make = [&](auto const& view) { + owner = std::make_unique( + 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_view>(dataset, make); + } else { + with_dataset_view, + cuvs::neighbors::device_standard_dataset_view>(dataset, make); + } + } else if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_HOST) { + if (padded) { + with_dataset_view, + cuvs::neighbors::host_padded_dataset_view>(dataset, make); + } else { + with_dataset_view, + cuvs::neighbors::host_standard_dataset_view>(dataset, make); + } + } else { + RAFT_FAIL("cuvsDatasetMakeVpq: invalid source dataset memory type"); + } + return owner; +} + template static void make_device_padded_dataset(raft::resources* res_ptr, DLManagedTensor* dataset_tensor, @@ -547,12 +634,17 @@ static void attach_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::attach_dataset(*res_ptr, idx, padded_view); - auto* holder = - new cuvs_cagra_c_api_index_lifetime_holder{std::move(padded_idx)}; - destroy_sg_cagra_c_api_box(index->addr); - index->addr = 0; - bind_index_lifetime_holder_to_C_index(index, index->dtype, holder); + using index_dataset_view_t = std::remove_cvref_t; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL("cuvsCagraUpdateDataset: replacing a VPQ dataset is not supported"); + } else { + auto padded_idx = cuvs::neighbors::cagra::attach_dataset(*res_ptr, idx, padded_view); + auto* holder = + new cuvs_cagra_c_api_index_lifetime_holder{std::move(padded_idx)}; + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + bind_index_lifetime_holder_to_C_index(index, index->dtype, holder); + } }); }); } @@ -959,13 +1051,19 @@ void _serialize(cuvsResources_t res, const char *filename, : "cuvsCagraSerializeGraph: null index handle"; with_index_by_layout(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; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + 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); }); } @@ -1122,8 +1220,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; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for VPQ indices"); + } else { + cuvs::neighbors::cagra::serialize_to_hnswlib( + *res_ptr, std::string(filename), idx); + } }); } template @@ -1197,7 +1300,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; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL("cuvsCagraIndexGetDataset does not expose VPQ datasets as dense DLPack tensors"); + } else { + cuvs::core::to_dlpack(idx.dataset().view(), dataset); + } + }); } template @@ -1578,6 +1688,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(res); + std::unique_ptr owner; + if (dataset->dtype.code == kDLFloat && dataset->dtype.bits == 32) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else if (dataset->dtype.code == kDLFloat && dataset->dtype.bits == 16) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else if (dataset->dtype.code == kDLInt && dataset->dtype.bits == 8) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else if (dataset->dtype.code == kDLUInt && dataset->dtype.bits == 8) { + owner = make_vpq_from_dense_dataset(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_attach_dataset(cuvsResources_t res, cuvsDataset_t device_padded_dataset, cuvsCagraIndex_t index) @@ -1741,7 +1881,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_VPQ) { + 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(dataset, [&](auto const& view) { + build_index_from_dataset_view(res_ptr, params, view, index); + }); + } else if (dtype.code == kDLFloat && dtype.bits == 32) { build_dispatch_on_mem_type_and_layout(res_ptr, params, dataset, index); } else if (dtype.code == kDLFloat && dtype.bits == 16) { build_dispatch_on_mem_type_and_layout(res_ptr, params, dataset, index); @@ -1839,10 +1987,12 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, auto index = *index_c_ptr; auto* box = reinterpret_cast(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( diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index 7d6c588bd9..ff8f807a6a 100644 --- a/c/tests/CMakeLists.txt +++ b/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 ) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index bd669dff78..6baee5b52f 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -157,8 +157,43 @@ class algo : public algo_base { // and set_search_dataset() should save the passed-in pointer somewhere. // The client code should call set_search_dataset() before searching, // and should not release dataset before searching is finished. + // + // A compressed base set is never handed over this way, as it has no dense rows to pass, so + // needs_dataset() says nothing about one. An algorithm that cannot search the base set it was + // given with the parameters it was given has to reject them itself, from set_search_param(). virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) {}; + /* ### Base sets the benchmark cannot read ### + + Some algorithms build from a base set that has been compressed for them offline, which is + neither dense nor made of `T` values and so cannot be passed as `build`'s `const T*`. Such a + base set is handed over as a file path and the algorithm owns whatever it decodes. + + A path rather than a library type on purpose: this header is shared with the faiss, hnswlib and + diskann wrappers, and must not acquire their unrelated dependencies. + + Loading is separate from building because the benchmark times only `build_from_base_set_file`. + Deserializing a compressed base set is benchmark setup, the same as reading a dense one, and + folding it into the measured build would inflate build times by however long the file takes to + read. `set_base_set_file` is also called in search mode, before `load`, for algorithms whose + index file holds only part of the picture and needs the base set reattached. + */ + + /** + * Hand over a compressed base set as a file path. Returns the number of rows in it, which the + * benchmark has no way of reading for itself. Called outside the timed sections. + */ + virtual auto set_base_set_file(const std::string& /*file*/) -> size_t + { + throw std::runtime_error{"This algorithm cannot read a compressed base set from a file."}; + } + + /** Build the index from the base set handed over by `set_base_set_file`. */ + virtual void build_from_base_set_file() + { + throw std::runtime_error{"This algorithm cannot build from a compressed base set."}; + } + /** * Make a shallow copy of the algo wrapper that shares the resources and ensures thread-safe * access to them. */ diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index a588b1e2a6..483ed524d9 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -134,8 +134,22 @@ void bench_build(::benchmark::State& state, const auto algo_property = parse_algo_property(algo->get_preference(), index.build_param); - const T* base_set = dataset->base_set(algo_property.dataset_memory_type); - std::size_t index_size = dataset->base_set_size(); + // Loading the base set is setup, not part of the build, so a compressed one is read here rather + // than inside the timed loop below; its row count comes back from the algorithm, since the + // benchmark cannot read the file. + const bool base_compressed = dataset->base_is_compressed(); + const T* base_set = nullptr; + std::size_t index_size = 0; + try { + if (base_compressed) { + index_size = algo->set_base_set_file(dataset->base_file()); + } else { + base_set = dataset->base_set(algo_property.dataset_memory_type); + index_size = dataset->base_set_size(); + } + } catch (const std::exception& e) { + return state.SkipWithError("Failed to load the base set: " + std::string(e.what())); + } cuda_timer gpu_timer{algo}; { @@ -158,7 +172,11 @@ void bench_build(::benchmark::State& state, [[maybe_unused]] auto ntx_lap = nvtx.lap(); [[maybe_unused]] auto gpu_lap = gpu_timer.lap(!no_lap_sync); try { - algo->build(base_set, index_size); + if (base_compressed) { + algo->build_from_base_set_file(); + } else { + algo->build(base_set, index_size); + } } catch (const std::exception& e) { state.SkipWithError(std::string(e.what())); } @@ -238,6 +256,11 @@ void bench_search(::benchmark::State& state, auto ualgo = create_algo(index.algo, dataset->distance(), dataset->dim(), index.build_param); a = ualgo.get(); + // An index built from a compressed base set stores only its graph, so `load` alone would + // leave it with nothing to search over. Handing the file over first lets `load` attach the + // same rows the graph was built from and return a complete index. The row count it returns + // is of no use here; only the build reports that. + if (dataset->base_is_compressed()) { a->set_base_set_file(dataset->base_file()); } a->load(index_file); current_algo = std::move(ualgo); } @@ -250,7 +273,13 @@ void bench_search(::benchmark::State& state, current_algo_props = std::make_unique(std::move(parse_algo_property(a->get_preference(), sp_json))); - if (search_param->needs_dataset()) { + // Not a reliable signal for a compressed base set: cuvs_cagra answers true unconditionally, + // because its index file carries no dataset and the dense rows are re-attached here instead. + // There are no dense rows to attach for a compressed base, and the algorithm already has the + // file from `set_base_set_file` above. An algorithm that truly cannot search without the dense + // rows, such as CAGRA with refine_ratio > 1, has to reject that combination itself: only it + // knows which of its search parameters read them. + if (search_param->needs_dataset() && !dataset->base_is_compressed()) { try { a->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); @@ -540,6 +569,7 @@ void dispatch_benchmark(std::string cmdline, auto dataset = std::make_shared>(dataset_conf.name, base_file, + dataset_conf.base_compressed, dataset_conf.subset_first_row, dataset_conf.subset_size, query_file, @@ -552,7 +582,13 @@ void dispatch_benchmark(std::string cmdline, if (build_mode) { if (file_exists(base_file)) { log_info("Using the dataset file '%s'", base_file.c_str()); - ::benchmark::AddCustomContext("n_records", std::to_string(dataset->base_set_size())); + if (dataset_conf.base_compressed) { + // The row count sits inside the compressed file, so it is reported per benchmark as + // `index_size` once the algorithm has read it, rather than up front here. + ::benchmark::AddCustomContext("base_format", "vpq"); + } else { + ::benchmark::AddCustomContext("n_records", std::to_string(dataset->base_set_size())); + } ::benchmark::AddCustomContext("dim", std::to_string(dataset->dim())); } else { log_warn("dataset file '%s' does not exist; benchmarking index building is impossible.", @@ -718,51 +754,59 @@ inline auto run_main(int argc, char** argv) -> int log_warn("cudart library is not found, GPU-based indices won't work."); } - auto& conf = bench::configuration::initialize(conf_stream, data_prefix, index_prefix); - std::string dtype = conf.get_dataset_conf().dtype; - - if (dtype == "float") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "half") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "uint8") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "int8") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else { - log_error("datatype '%s' is not supported", dtype.c_str()); + // A rejected configuration reaches us as an exception, from the json parser or from the dataset + // itself. Reporting it here keeps that a legible error and a non-zero exit code, rather than an + // abort from an uncaught exception. + try { + auto& conf = bench::configuration::initialize(conf_stream, data_prefix, index_prefix); + std::string dtype = conf.get_dataset_conf().dtype; + + if (dtype == "float") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "half") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "uint8") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "int8") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else { + log_error("datatype '%s' is not supported", dtype.c_str()); + return -1; + } + } catch (const std::exception& e) { + log_error("%s", e.what()); return -1; } diff --git a/cpp/bench/ann/src/common/conf.hpp b/cpp/bench/ann/src/common/conf.hpp index afc7bc0a1f..0ed63d0ba6 100644 --- a/cpp/bench/ann/src/common/conf.hpp +++ b/cpp/bench/ann/src/common/conf.hpp @@ -8,12 +8,19 @@ #include #include +#include #include #include #include namespace cuvs::bench { +inline auto has_suffix(const std::string& str, const std::string& suffix) -> bool +{ + return str.size() >= suffix.size() && + str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; +} + class configuration { public: struct index { @@ -40,6 +47,12 @@ class configuration { std::string distance; std::optional groundtruth_neighbors_file{std::nullopt}; + // The base_file holds rows already compressed for the algorithm (a .vpq written by the offline + // VPQ compression tool) rather than dense vectors. The benchmark cannot read such a file: its + // rows are not `dtype` values, so they cannot travel through `algo::build`. The path is + // handed to the algorithm instead. Queries stay dense, and `dtype` keeps describing them. + bool base_compressed{false}; + // data type of input dataset, possible values ["float", "int8", "uint8"] std::string dtype; @@ -99,11 +112,30 @@ class configuration { } if (conf.contains("subset_size")) { dataset_conf_.subset_size = conf.at("subset_size"); } + // Decided separately from the dtype inference below, so that an explicit "dtype" does not stop + // us noticing that the base set is compressed. + if (conf.contains("base_format")) { + const auto base_format = conf.at("base_format").get(); + if (base_format == "vpq") { + dataset_conf_.base_compressed = true; + } else if (base_format != "dense") { + throw std::runtime_error("Unknown base_format '" + base_format + + "', expected \"vpq\" or \"dense\""); + } + } else { + dataset_conf_.base_compressed = has_suffix(dataset_conf_.base_file, ".vpq"); + } + if (conf.contains("dtype")) { dataset_conf_.dtype = conf.at("dtype"); } else { auto filename = dataset_conf_.base_file; - if (filename.size() > 6 && filename.compare(filename.size() - 6, 6, "f16bin") == 0) { + if (dataset_conf_.base_compressed) { + // A VPQ dataset stores its codebooks as half, but it is searched with float queries and + // yields a float index, so float is the type the benchmark instantiates. Keyed off the flag + // rather than the suffix, so that an explicit base_format also gets a dtype. + dataset_conf_.dtype = "float"; + } else if (filename.size() > 6 && filename.compare(filename.size() - 6, 6, "f16bin") == 0) { dataset_conf_.dtype = "half"; } else if (filename.size() > 9 && filename.compare(filename.size() - 9, 9, "fp16.fbin") == 0) { diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index 4dc43c343c..93a193cad4 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -152,6 +153,10 @@ struct dataset { private: std::string name_; std::string distance_; + std::string base_file_; + // A compressed base set is opaque to the benchmark: `base_set_` stays lazy and untouched, and the + // path is handed to the algorithm, which is the only thing able to decode it. + bool base_compressed_; blob base_set_; blob query_set_; std::optional> filter_bitset_; @@ -171,9 +176,22 @@ struct dataset { } } + // Reading a compressed base set as a .bin would not fail, it would succeed on garbage: the first + // eight bytes of a .vpq are a dtype prefix and a numpy magic, which parse as an absurd shape. + // Hence an explicit error, and callers that can proceed must ask `base_is_compressed()` first. + inline void throw_if_base_compressed(const char* what) const + { + if (base_compressed_) { + throw std::runtime_error{std::string{"dataset::"} + what + + "() is not available for the compressed base_file '" + base_file_ + + "'; only the algorithm can read that file."}; + } + } + public: dataset(std::string name, std::string base_file, + bool base_compressed, uint32_t subset_first_row, uint32_t subset_size, std::string query_file, @@ -182,9 +200,27 @@ struct dataset { std::optional filtering_rate = std::nullopt) : name_{std::move(name)}, distance_{std::move(distance)}, + base_file_{base_file}, + base_compressed_{base_compressed}, base_set_{base_file, subset_first_row, subset_size}, query_set_{query_file} { + if (base_compressed_) { + // Which rows went into the file was decided when it was written, so subsetting here is not + // merely unsupported: there is nothing left to select from. + if (subset_first_row != 0 || subset_size != 0) { + throw std::runtime_error{ + "A compressed base_file cannot be subset by the benchmark; choose the rows when " + "compressing it (the tool's --subset_first_row / --subset_size) and drop these keys."}; + } + // The bitset is sized from the base set row count, which is inside the compressed file. + if (filtering_rate.has_value()) { + throw std::runtime_error{ + "filtering_rate is not supported with a compressed base_file: generating the filter " + "bitset needs the base set size, which only the algorithm can read."}; + } + } + if (filtering_rate.has_value()) { // Generate a random bitset for filtering auto n_rows = static_cast(subset_size) + static_cast(subset_first_row); @@ -210,6 +246,8 @@ struct dataset { [[nodiscard]] auto name() const -> std::string { return name_; } [[nodiscard]] auto distance() const -> std::string { return distance_; } + [[nodiscard]] auto base_file() const -> std::string { return base_file_; } + [[nodiscard]] auto base_is_compressed() const -> bool { return base_compressed_; } [[nodiscard]] auto dim() const -> int { auto d = dim_.load(std::memory_order_relaxed); @@ -221,6 +259,14 @@ struct dataset { } catch (const std::runtime_error& e) { // Any exception raised above will re-raise next time we try to access the query set. query_set_.reset_lazy_state(); + // A compressed base set has no dense header to fall back on, and reading it as one would + // yield a nonsense dimension rather than an error. + if (base_compressed_) { + throw std::runtime_error{ + "Cannot determine the dataset dimension: the query set is not readable and the base set " + "is compressed. " + + std::string{e.what()}}; + } // If the query set is not accessible, use the base set. // Don't catch the exception here, because we have nothing else to do anyway. d = static_cast(base_set_.n_cols()); @@ -235,6 +281,7 @@ struct dataset { } [[nodiscard]] auto base_set_size() const -> size_t { + throw_if_base_compressed("base_set_size"); std::lock_guard lock(mutex_); auto r = base_set_.n_rows(); cache_dim(base_set_); @@ -272,6 +319,7 @@ struct dataset { [[nodiscard]] auto base_set() const -> const DataT* { + throw_if_base_compressed("base_set"); std::lock_guard lock(mutex_); auto* r = base_set_.data(); cache_dim(base_set_); @@ -281,6 +329,7 @@ struct dataset { HugePages request_hugepages_2mb = HugePages::kDisable) const -> const DataT* { + throw_if_base_compressed("base_set"); std::lock_guard lock(mutex_); auto* r = base_set_.data(memory_type, request_hugepages_2mb); cache_dim(base_set_); diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 45eb945fbc..220ffd86ab 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -364,6 +364,7 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_"); nlohmann::json nn_descent_conf = collect_conf_with_prefix(conf, "nn_descent_"); nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_"); + nlohmann::json build_search_conf = collect_conf_with_prefix(conf, "build_search_"); // When graph_build_algo is not specified, leave graph_build_params as monostate so the // CAGRA build uses AUTO selection (NN_DESCENT or IVF_PQ based on dataset/heuristics). @@ -394,6 +395,85 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index } else if constexpr (std::is_same_v) { parse_build_param(nn_descent_conf, arg); + } else if constexpr (std::is_same_v< + U, + cuvs::neighbors::graph_build_params::iterative_search_params>) { + if (build_search_conf.contains("width")) { + arg.search_width = build_search_conf.at("width"); + } + if (build_search_conf.contains("max_iterations")) { + arg.max_iterations = build_search_conf.at("max_iterations"); + } + if (build_search_conf.contains("min_iterations")) { + arg.min_iterations = build_search_conf.at("min_iterations"); + } + if (build_search_conf.contains("itopk")) { arg.itopk_size = build_search_conf.at("itopk"); } + if (build_search_conf.contains("max_queries")) { + arg.max_queries = build_search_conf.at("max_queries"); + } + if (build_search_conf.contains("team_size")) { + arg.team_size = build_search_conf.at("team_size"); + } + if (build_search_conf.contains("thread_block_size")) { + arg.thread_block_size = build_search_conf.at("thread_block_size"); + } + if (build_search_conf.contains("hashmap_min_bitlen")) { + arg.hashmap_min_bitlen = build_search_conf.at("hashmap_min_bitlen"); + } + if (build_search_conf.contains("hashmap_max_fill_rate")) { + arg.hashmap_max_fill_rate = build_search_conf.at("hashmap_max_fill_rate"); + } + if (build_search_conf.contains("num_random_samplings")) { + arg.num_random_samplings = build_search_conf.at("num_random_samplings"); + } + if (build_search_conf.contains("persistent")) { + arg.persistent = build_search_conf.at("persistent"); + } + if (build_search_conf.contains("persistent_lifetime")) { + arg.persistent_lifetime = build_search_conf.at("persistent_lifetime"); + } + if (build_search_conf.contains("persistent_device_usage")) { + arg.persistent_device_usage = build_search_conf.at("persistent_device_usage"); + } + if (build_search_conf.contains("algo")) { + std::string algo = build_search_conf.at("algo"); + if (algo == "single_cta") { + arg.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + } else if (algo == "multi_cta") { + arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + } else if (algo == "multi_kernel") { + arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_KERNEL; + } else if (algo == "auto") { + arg.algo = cuvs::neighbors::cagra::search_algo::AUTO; + } + } + if (build_search_conf.contains("hashmap_mode")) { + std::string mode = build_search_conf.at("hashmap_mode"); + if (mode == "hash") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::HASH; + } else if (mode == "small") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::SMALL; + } else if (mode == "auto") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::AUTO; + } + } + // Precision of the codebook/query in shared memory for the VPQ search used during + // the iterative build. Accepts an integer code (0=F16, 1=E5M2) or a string. + if (build_search_conf.contains("smem_dtype")) { + const auto& sd = build_search_conf.at("smem_dtype"); + if (sd.is_number_integer()) { + arg.smem_dtype = static_cast(sd.get()); + } else { + std::string s = sd.get(); + if (s == "f16" || s == "F16" || s == "fp16" || s == "half") { + arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } else if (s == "e5m2" || s == "E5M2" || s == "fp8") { + arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::E5M2; + } else { + throw std::runtime_error("invalid value for build_search smem_dtype: " + s); + } + } + } } }, params.graph_build_params); diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 43ae7a6235..dc18a6b792 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -34,13 +34,147 @@ #include #include #include +#include #include #include +namespace CUVS_EXPORT cuvs { +namespace neighbors { +namespace cagra { + +/** + * @defgroup cagra_cpp_search_params CAGRA index search parameters + * @{ + */ + +enum class search_algo { + /** For large batch sizes. */ + SINGLE_CTA = 0, + /** For small batch sizes. */ + MULTI_CTA = 1, + MULTI_KERNEL = 2, + AUTO = 100 +}; + +enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; + +enum class internal_dtype { F16 = 0, E5M2 = 1 }; + +struct search_params : cuvs::neighbors::search_params { + /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ + size_t max_queries = 0; + + /** Number of intermediate search results retained during the search. + * + * This is the main knob to adjust trade off between accuracy and search speed. + * Higher values improve the search accuracy. + */ + size_t itopk_size = 64; + + /** Upper limit of search iterations. Auto select when 0.*/ + size_t max_iterations = 0; + + // In the following we list additional search parameters for fine tuning. + // Reasonable default values are automatically chosen. + + /** Which search implementation to use. */ + search_algo algo = search_algo::AUTO; + + /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */ + size_t team_size = 0; + + /** Number of graph nodes to select as the starting point for the search in each iteration. aka + * search width?*/ + size_t search_width = 1; + /** Lower limit of search iterations. */ + size_t min_iterations = 0; + + /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */ + size_t thread_block_size = 0; + /** Hashmap type. Auto selection when AUTO. */ + hash_mode hashmap_mode = hash_mode::AUTO; + /** Lower limit of hashmap bit length. More than 8. */ + size_t hashmap_min_bitlen = 0; + /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/ + float hashmap_max_fill_rate = 0.5; + + /** Number of iterations of initial random seed node selection. 1 or more. */ + uint32_t num_random_samplings = 1; + /** Bit mask used for initial random seed node selection. */ + uint64_t rand_xor_mask = 0x128394; + + /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ + bool persistent = false; + /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ + float persistent_lifetime = 2; + /** + * Set the fraction of maximum grid size used by persistent kernel. + * Value 1.0 means the kernel grid size is maximum possible for the selected device. + * The value must be greater than 0.0 and not greater than 1.0. + * + * One may need to run other kernels alongside this persistent kernel. This parameter can + * be used to reduce the grid size of the persistent kernel to leave a few SMs idle. + * Note: running any other work on GPU alongside with the persistent kernel makes the setup + * fragile. + * - Running another kernel in another thread usually works, but no progress guaranteed + * - Any CUDA allocations block the context (this issue may be obscured by using pools) + * - Memory copies to not-pinned host memory may block the context + * + * Even when we know there are no other kernels working at the same time, setting + * kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. + * If you suspect this is an issue, you can reduce this number to ~0.9 without a significant + * impact on the throughput. + */ + float persistent_device_usage = 1.0; + + /** + * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. + * The value must be equal to or greater than 0.0 and less than 1.0. Default value is + * negative, in which case the filtering rate is automatically calculated when possible. + * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is + * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be + * inferred from the source string. + */ + float filtering_rate = -1.0; + + /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ + * supports FP8. **/ + internal_dtype smem_dtype = internal_dtype::F16; +}; + +/** + * @} + */ + +} // namespace cagra +} // namespace neighbors +} // namespace CUVS_EXPORT cuvs + namespace CUVS_EXPORT cuvs { namespace neighbors { namespace graph_build_params { -using iterative_search_params = cuvs::neighbors::search_params; +/** + * Parameters for the iterative CAGRA graph build algorithm. + * + * Inherits from cagra::search_params so that all search tuning knobs + * (search_width, max_iterations, itopk_size, etc.) are available for + * controlling the search-and-optimize loop during graph construction. + * The defaults are tuned for the build loop (e.g. search_width=1, + * max_iterations=8) and may differ from the regular search defaults. + * + */ +struct iterative_search_params : cuvs::neighbors::cagra::search_params { + iterative_search_params() + { + this->search_width = 1; + this->max_iterations = 8; + // itopk_size controls the search during the *growing* iterations of the build loop. + // 0 (default) means auto-select per iteration (max(graph_degree + 32, 128)); a nonzero + // value overrides it for the growing iterations. The final iteration always uses a fixed + // itopk tied to the output topk, regardless of this value. + this->itopk_size = 0; + } +}; /** Specialized parameters for ACE (Augmented Core Extraction) graph build */ struct ace_params { @@ -311,110 +445,6 @@ struct index_params : cuvs::neighbors::index_params { cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded); }; -/** - * @} - */ - -/** - * @defgroup cagra_cpp_search_params CAGRA index search parameters - * @{ - */ - -enum class search_algo { - /** For large batch sizes. */ - SINGLE_CTA = 0, - /** For small batch sizes. */ - MULTI_CTA = 1, - MULTI_KERNEL = 2, - AUTO = 100 -}; - -enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; - -enum class internal_dtype { F16 = 0, E5M2 = 1 }; - -struct search_params : cuvs::neighbors::search_params { - /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ - size_t max_queries = 0; - - /** Number of intermediate search results retained during the search. - * - * This is the main knob to adjust trade off between accuracy and search speed. - * Higher values improve the search accuracy. - */ - size_t itopk_size = 64; - - /** Upper limit of search iterations. Auto select when 0.*/ - size_t max_iterations = 0; - - // In the following we list additional search parameters for fine tuning. - // Reasonable default values are automatically chosen. - - /** Which search implementation to use. */ - search_algo algo = search_algo::AUTO; - - /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */ - size_t team_size = 0; - - /** Number of graph nodes to select as the starting point for the search in each iteration. aka - * search width?*/ - size_t search_width = 1; - /** Lower limit of search iterations. */ - size_t min_iterations = 0; - - /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */ - size_t thread_block_size = 0; - /** Hashmap type. Auto selection when AUTO. */ - hash_mode hashmap_mode = hash_mode::AUTO; - /** Lower limit of hashmap bit length. More than 8. */ - size_t hashmap_min_bitlen = 0; - /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/ - float hashmap_max_fill_rate = 0.5; - - /** Number of iterations of initial random seed node selection. 1 or more. */ - uint32_t num_random_samplings = 1; - /** Bit mask used for initial random seed node selection. */ - uint64_t rand_xor_mask = 0x128394; - - /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ - bool persistent = false; - /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ - float persistent_lifetime = 2; - /** - * Set the fraction of maximum grid size used by persistent kernel. - * Value 1.0 means the kernel grid size is maximum possible for the selected device. - * The value must be greater than 0.0 and not greater than 1.0. - * - * One may need to run other kernels alongside this persistent kernel. This parameter can - * be used to reduce the grid size of the persistent kernel to leave a few SMs idle. - * Note: running any other work on GPU alongside with the persistent kernel makes the setup - * fragile. - * - Running another kernel in another thread usually works, but no progress guaranteed - * - Any CUDA allocations block the context (this issue may be obscured by using pools) - * - Memory copies to not-pinned host memory may block the context - * - * Even when we know there are no other kernels working at the same time, setting - * kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. - * If you suspect this is an issue, you can reduce this number to ~0.9 without a significant - * impact on the throughput. - */ - float persistent_device_usage = 1.0; - - /** - * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. - * The value must be equal to or greater than 0.0 and less than 1.0. Default value is - * negative, in which case the filtering rate is automatically calculated when possible. - * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is - * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be - * inferred from the source string. - */ - float filtering_rate = -1.0; - - /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ - * supports FP8. **/ - internal_dtype smem_dtype = internal_dtype::F16; -}; - /** * @} */ @@ -914,9 +944,12 @@ using vpq_f32_index = index -using cagra_index_t = index, - uint32_t, - cuvs::neighbors::dataset_view_type_t>; +using cagra_index_t = std::conditional_t< + cuvs::neighbors::is_device_vpq_f16_dataset_view_v, + index>, + index, + uint32_t, + cuvs::neighbors::dataset_view_type_t>>; /** * @} @@ -928,10 +961,11 @@ using cagra_index_t = index>` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_vpq_dataset_view const& dataset) + -> index>; + /** * @brief Build from a device padded dataset view (`float`). * @param[in] res raft resources diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 112341f2ad..f6456624d6 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -14,6 +14,9 @@ #include #include +#include +#include +#include #include #include @@ -331,6 +334,82 @@ template } } +/** Current VPQ dataset serialization format version. */ +inline constexpr int vpq_serialization_version = 1; + +/** + * @brief Write a VPQ dataset (both codebooks plus the encoded rows) to a stream. + * + * Lets compression be done once, offline, and reused: the encoded rows are what CAGRA-Q builds and + * searches over, so a stored VPQ dataset removes the need to keep the dense vectors around or + * re-quantize them on every run. + * + * The file opens with the same preamble as `cagra::serialize` — a 4-byte NumPy dtype prefix then + * `vpq_serialization_version` — followed by a dataset kind tag and the codebook element type. A file + * of the wrong kind, or one written by an older format, is rejected rather than misread. Bump the + * version whenever the encoded row layout changes, since that layout is a library convention and is + * not otherwise described by the file. + * + * @code{.cpp} + * #include + * #include + * + * // Offline, once. + * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, rows); + * cuvs::preprocessing::quantize::pq::serialize(res, "base.vpq", vpq); + * + * // Later, per run: load the compressed rows and build a CAGRA-Q graph over them. + * std::unique_ptr> loaded; + * cuvs::preprocessing::quantize::pq::deserialize(res, "base.vpq", &loaded); + * auto index = cuvs::neighbors::cagra::build(res, index_params, loaded->as_dataset_view()); + * // `loaded` must outlive `index`, which only holds a view of it. + * @endcode + * + * @param[in] res raft resource + * @param[in] os output stream, opened in binary mode + * @param[in] dataset the VPQ dataset to write + */ +void serialize(raft::resources const& res, + std::ostream& os, + const cuvs::neighbors::device_vpq_dataset& dataset); + +/** + * @copydoc serialize + * + * @param[in] res raft resource + * @param[in] filename path to write, truncated if it exists + * @param[in] dataset the VPQ dataset to write + */ +void serialize(raft::resources const& res, + const std::string& filename, + const cuvs::neighbors::device_vpq_dataset& dataset); + +/** + * @brief Read a VPQ dataset written by `serialize`. + * + * Returned through an out-parameter because the dataset owns device allocations and has no default + * constructor, matching how `cagra::deserialize` hands back its dataset. Throws if the blob was not + * written by `serialize` or holds codebooks of a different element type. + * + * @param[in] res raft resource + * @param[in] is input stream, opened in binary mode + * @param[out] out_dataset receives the loaded dataset; must not be null + */ +void deserialize(raft::resources const& res, + std::istream& is, + std::unique_ptr>* out_dataset); + +/** + * @copydoc deserialize + * + * @param[in] res raft resource + * @param[in] filename path to read + * @param[out] out_dataset receives the loaded dataset; must not be null + */ +void deserialize(raft::resources const& res, + const std::string& filename, + std::unique_ptr>* out_dataset); + /** @} */ // end of group product } // namespace pq diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 36b4078a06..ea5620fa23 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -296,13 +296,43 @@ template auto build(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> cuvs::neighbors::cagra::cagra_index_t { - using T = cuvs::neighbors::cagra_view_element_type_t; - using IdxT = uint32_t; + using index_type = cuvs::neighbors::cagra::cagra_index_t; + using T = typename index_type::value_type; + using IdxT = uint32_t; // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until attach_dataset(...) supplies a device-padded dataset. if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { - RAFT_FAIL("cagra::build: VPQ-compressed dataset cannot be used for dense graph construction."); + auto effective_params = params; + if (std::holds_alternative(effective_params.graph_build_params)) { + effective_params.graph_build_params = graph_build_params::iterative_search_params{}; + } + + RAFT_EXPECTS(std::holds_alternative( + effective_params.graph_build_params), + "cagra::build: a VPQ dataset requires iterative_search_params graph construction"); + RAFT_EXPECTS(effective_params.metric == cuvs::distance::DistanceType::L2Expanded, + "cagra::build: a VPQ dataset supports only L2Expanded distance"); + RAFT_EXPECTS(dataset.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); + RAFT_EXPECTS(dataset.dset().pq_bits() == 8, + "cagra::build: VPQ dataset requires pq_bits == 8, got %u", + dataset.dset().pq_bits()); + auto const pq_len = dataset.dset().pq_len(); + RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, + "cagra::build: VPQ dataset requires pq_len in {2, 4, 8}, got %u", + pq_len); + + detail::check_graph_degree(effective_params.intermediate_graph_degree, + effective_params.graph_degree, + static_cast(dataset.n_rows())); + auto cagra_graph = detail::iterative_build_graph(res, effective_params, dataset); + + index_type idx(res, effective_params.metric); + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + if (effective_params.attach_dataset_on_build) { + idx.update_device_dataset_same_layout(res, dataset); + } + return idx; } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix( res, params, dataset); diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index acaaa942c1..90d63c3ca9 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -18,6 +18,9 @@ using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view< using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; +#if @emit_vpq_build@ +using inst_device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; +#endif } // namespace namespace cuvs::neighbors::cagra { @@ -55,6 +58,13 @@ CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_padded_view_t, CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_standard_view_t, cuvs::neighbors::cagra::host_standard_index); +#if @emit_vpq_build@ +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD( + inst_device_vpq_view_t, + cuvs::neighbors::cagra:: + index>); +#endif + #undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_build_matrix.json b/cpp/src/neighbors/cagra_build_matrix.json index a7995005c4..9fae2b33f9 100644 --- a/cpp/src/neighbors/cagra_build_matrix.json +++ b/cpp/src/neighbors/cagra_build_matrix.json @@ -2,19 +2,23 @@ "_data": [ { "data_type": "float", - "data_abbrev": "f" + "data_abbrev": "f", + "emit_vpq_build": 1 }, { "data_type": "half", - "data_abbrev": "h" + "data_abbrev": "h", + "emit_vpq_build": 0 }, { "data_type": "int8_t", - "data_abbrev": "i8" + "data_abbrev": "i8", + "emit_vpq_build": 0 }, { "data_type": "uint8_t", - "data_abbrev": "u8" + "data_abbrev": "u8", + "emit_vpq_build": 0 } ], "_index": [ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 8705926a41..32786b4e1b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -7,7 +7,6 @@ #include "../../../core/nvtx.hpp" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" #include "graph_core.cuh" -#include #include #include @@ -46,7 +45,6 @@ #include #include -#include #include namespace cuvs::neighbors::cagra::detail { @@ -1994,51 +1992,178 @@ void optimize( res, knn_graph_internal, new_graph_internal, guarantee_connectivity); } -// RAII wrapper for allocating memory with Transparent HugePage -struct mmap_owner { - // Allocate a new memory (not backed by a file) - mmap_owner(size_t size) : size_{size} - { - int flags = MAP_ANONYMOUS | MAP_PRIVATE; - ptr_ = mmap(nullptr, size, PROT_READ | PROT_WRITE, flags, -1, 0); - if (ptr_ == MAP_FAILED) { - ptr_ = nullptr; - throw std::runtime_error("cuvs::mmap_owner error"); - } - if (madvise(ptr_, size, MADV_HUGEPAGE) != 0) { - munmap(ptr_, size); - ptr_ = nullptr; - throw std::runtime_error("cuvs::mmap_owner error"); - } +template +__global__ void kern_reconstruct_vpq_queries(const uint8_t* encoded_data, + uint32_t encoded_row_len, + const MathT* vq_codebook, + const MathT* pq_codebook, + uint32_t dim, + uint32_t pq_len, + uint64_t offset, + uint32_t batch_size, + T* output) +{ + const uint64_t batch_idx = blockIdx.x; + if (batch_idx >= batch_size) return; + const uint64_t vec_idx = offset + batch_idx; + const uint8_t* vec_data = encoded_data + vec_idx * encoded_row_len; + const uint32_t vq_code = *reinterpret_cast(vec_data); + const uint8_t* pq_codes = vec_data + sizeof(uint32_t); + const MathT* vq_centroid_ptr = vq_codebook + static_cast(vq_code) * dim; + + for (uint32_t d = threadIdx.x; d < dim; d += blockDim.x) { + uint32_t j = d / pq_len; + uint32_t k = d % pq_len; + float val = static_cast(vq_centroid_ptr[d]) + + static_cast(pq_codebook[static_cast(pq_codes[j]) * pq_len + k]); + output[batch_idx * dim + d] = static_cast(val); } +} - ~mmap_owner() noexcept - { - if (ptr_ != nullptr) { munmap(ptr_, size_); } - } +template +void reconstruct_vpq_queries(raft::resources const& res, + const cuvs::neighbors::device_vpq_dataset& vpq_dset, + uint64_t offset, + uint32_t batch_size, + raft::device_matrix_view output) +{ + const uint32_t dim = vpq_dset.dim(); + const uint32_t pq_len = vpq_dset.pq_len(); + const uint32_t threads = std::min(dim, 256u); + + kern_reconstruct_vpq_queries + <<>>( + vpq_dset.data.data_handle(), + vpq_dset.encoded_row_length(), + vpq_dset.vq_code_book.data_handle(), + vpq_dset.pq_code_book.data_handle(), + dim, + pq_len, + offset, + batch_size, + output.data_handle()); +} - // No copies for owning struct - mmap_owner(const mmap_owner& res) = delete; - auto operator=(const mmap_owner& other) -> mmap_owner& = delete; - // Moving is fine - mmap_owner(mmap_owner&& other) - : ptr_{std::exchange(other.ptr_, nullptr)}, size_{std::exchange(other.size_, 0)} - { - } - auto operator=(mmap_owner&& other) -> mmap_owner& - { - std::swap(this->ptr_, other.ptr_); - std::swap(this->size_, other.size_); - return *this; +// Runs CAGRA search for `curr_query_size` queries against `idx` in chunks of `max_chunk_size`, +// stacks the results into a kNN graph, and optimizes it into the next graph (returned). +// +// Query source: +// - `vpq_queries == nullptr`: queries are read directly from `dev_query_view` (uncompressed +// build; the view is a slice of the resident device dataset). +// - `vpq_queries != nullptr`: `dev_query_view` is ignored and each chunk of queries is +// reconstructed on the fly from the VPQ codes into a small reusable scratch buffer, so we +// never materialize the whole (up to N x dim) reconstructed dataset. +template +raft::device_matrix search_and_optimize( + raft::resources const& res, + const cuvs::neighbors::cagra::search_params& search_params, + const cuvs::neighbors::cagra::index& idx, + raft::device_matrix_view dev_query_view, + raft::device_matrix_view dev_neighbors, + raft::device_matrix_view dev_distances, + raft::device_matrix prev_graph, + const cuvs::neighbors::device_vpq_dataset* vpq_queries, + size_t curr_query_size, + size_t next_graph_degree, + size_t curr_topk, + uint64_t max_chunk_size, + int64_t query_dim, + bool guarantee_connectivity) +{ + auto stream = raft::resource::get_cuda_stream(res); + + // These buffers scale with N (e.g. N * (intermediate_degree+1) for the kNN graph). Allocate them + // from the default device resource (a pool over device memory): allocating from the large + // workspace resource here would use an unpooled managed_memory_resource, paying a synchronous + // cudaMallocManaged/cudaFree every iteration for multi-GB buffers. + auto dev_knn_graph = raft::make_device_matrix(res, curr_query_size, curr_topk); + + // Scratch for one reconstructed or depadded chunk. Reused across chunks; safe because all + // reconstruct/search/copy work is serialized on `stream`. + auto batch_queries = + vpq_queries != nullptr || dev_query_view.extent(1) != query_dim + ? raft::make_device_matrix(res, static_cast(max_chunk_size), query_dim) + : raft::make_device_matrix(res, 0, 0); + + auto run_batch = [&](int64_t offset, + int64_t batch_size, + raft::device_matrix_view batch_query_view) { + auto batch_dev_neighbors_view = raft::make_device_matrix_view( + dev_neighbors.data_handle(), batch_size, curr_topk); + auto batch_dev_distances_view = raft::make_device_matrix_view( + dev_distances.data_handle(), batch_size, curr_topk); + + cuvs::neighbors::cagra::search(res, + search_params, + idx, + batch_query_view, + batch_dev_neighbors_view, + batch_dev_distances_view); + + raft::copy(dev_knn_graph.data_handle() + offset * curr_topk, + batch_dev_neighbors_view.data_handle(), + batch_size * curr_topk, + stream); + }; + + if (vpq_queries != nullptr) { + // Reconstruct-and-search one chunk at a time: reconstruct source rows [offset, offset+bs) into + // the scratch, then search that chunk. + for (int64_t offset = 0; offset < static_cast(curr_query_size); + offset += static_cast(max_chunk_size)) { + const int64_t batch_size = std::min(static_cast(max_chunk_size), + static_cast(curr_query_size) - offset); + reconstruct_vpq_queries(res, + *vpq_queries, + static_cast(offset), + static_cast(batch_size), + batch_queries.view()); + auto batch_query_view = raft::make_device_matrix_view( + batch_queries.data_handle(), batch_size, query_dim); + run_batch(offset, batch_size, batch_query_view); + } + } else { + const int64_t source_row_width = dev_query_view.extent(1); + auto query_batch = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( + res, + dev_query_view.data_handle(), + static_cast(curr_query_size), + source_row_width, + max_chunk_size, + stream, + raft::resource::get_workspace_resource_ref(res)); + for (const auto& batch : query_batch) { + raft::device_matrix_view batch_query_view; + if (source_row_width != query_dim) { + raft::copy_matrix(batch_queries.data_handle(), + query_dim, + batch.data(), + source_row_width, + query_dim, + batch.size(), + stream); + batch_query_view = raft::make_device_matrix_view( + batch_queries.data_handle(), static_cast(batch.size()), query_dim); + } else { + batch_query_view = raft::make_device_matrix_view( + batch.data(), static_cast(batch.size()), query_dim); + } + run_batch( + static_cast(batch.offset()), static_cast(batch.size()), batch_query_view); + } } - [[nodiscard]] auto data() const -> void* { return ptr_; } - [[nodiscard]] auto size() const -> size_t { return size_; } + // The previous-iteration graph (which `idx` was built on) is no longer needed now that the + // search has produced `dev_knn_graph`. Release it before allocating the full-size output graph + // so we never hold two large graph buffers at once. + prev_graph = raft::make_device_matrix(res, 0, 0); - private: - void* ptr_; - size_t size_; -}; + auto dev_output_graph = + raft::make_device_matrix(res, curr_query_size, next_graph_degree); + + graph::optimize(res, dev_knn_graph.view(), dev_output_graph.view(), guarantee_connectivity); + return dev_output_graph; +} /** Upload and/or pad `dataset` to a device-resident CAGRA-aligned view for iterative internal * search. */ @@ -2059,7 +2184,8 @@ auto ensure_device_padded_for_iterative_search( } template - requires cuvs::neighbors::is_dense_row_major_dataset_view_v + requires(cuvs::neighbors::is_dense_row_major_dataset_view_v || + cuvs::neighbors::is_device_vpq_f16_dataset_view_v) auto iterative_build_graph(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> raft::host_matrix @@ -2067,22 +2193,36 @@ auto iterative_build_graph(raft::resources const& res, size_t intermediate_degree = params.intermediate_graph_degree; size_t graph_degree = params.graph_degree; + const auto& iter_params = + std::get(params.graph_build_params); + RAFT_LOG_INFO("Build search params: search_width=%zu, max_iterations=%zu", + iter_params.search_width, + iter_params.max_iterations); + auto cagra_graph = raft::make_host_matrix(0, 0); - // Iteratively improve the accuracy of the graph by repeatedly running - // CAGRA's search() and optimize(). Host or non-CAGRA-aligned device inputs are uploaded - // and padded here only for the internal search loop — same role as main's - // make_aligned_dataset() inside iterative_build_graph. IVF-PQ / NN-descent never take this path. + // Iteratively improve the graph by repeatedly running CAGRA search and optimize. Dense inputs are + // padded on device; VPQ inputs are searched directly and reconstructed per query batch. RAFT_LOG_INFO("Iteratively creating/improving graph index using CAGRA's search() and optimize()"); std::unique_ptr> padded_own; - auto search_dataset = ensure_device_padded_for_iterative_search(res, dataset, padded_own); - - auto dev_dataset = search_dataset.view(); - uint32_t logical_dim = search_dataset.dim(); + auto dev_dataset = + raft::make_device_matrix_view(static_cast(nullptr), 0, 0); + uint32_t logical_dim = dataset.dim(); + uint64_t final_graph_size; + const cuvs::neighbors::device_vpq_dataset* vpq_dataset = nullptr; + + if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + final_graph_size = static_cast(dataset.n_rows()); + vpq_dataset = &dataset.dset(); + } else { + auto search_dataset = ensure_device_padded_for_iterative_search(res, dataset, padded_own); + dev_dataset = search_dataset.view(); + logical_dim = search_dataset.dim(); + final_graph_size = static_cast(search_dataset.n_rows()); + } // Determine initial graph size. - uint64_t final_graph_size = (uint64_t)search_dataset.n_rows(); uint64_t initial_graph_size = (final_graph_size + 1) / 2; while (initial_graph_size > graph_degree * 64) { initial_graph_size = (initial_graph_size + 1) / 2; @@ -2097,12 +2237,6 @@ auto iterative_build_graph(raft::resources const& res, auto dev_neighbors = raft::make_device_matrix(res, max_chunk_size, topk); auto dev_distances = raft::make_device_matrix(res, max_chunk_size, topk); - std::optional> query_contiguous; - if (static_cast(logical_dim) != dev_dataset.extent(1)) { - query_contiguous.emplace( - raft::make_device_matrix(res, max_chunk_size, logical_dim)); - } - // Determine graph degree and number of search results while increasing // graph size. auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24)); @@ -2110,6 +2244,16 @@ auto iterative_build_graph(raft::resources const& res, RAFT_LOG_DEBUG("# graph_degree = %lu", (uint64_t)graph_degree); RAFT_LOG_DEBUG("# topk = %lu", (uint64_t)topk); + // A fixed itopk_size (0 = auto) governs the growing iterations, which build graphs of degree + // ~graph_degree/2 and thus request topk ~= graph_degree/2 + 1; the search planner requires + // topk <= itopk_size. (The full-size iterations override itopk internally, so they are not + // constrained by this value.) + RAFT_EXPECTS(iter_params.itopk_size == 0 || iter_params.itopk_size >= graph_degree / 2 + 1, + "iterative build search itopk_size (%zu) must be 0 (auto) or >= " + "graph_degree / 2 + 1 (%zu)", + (size_t)iter_params.itopk_size, + (size_t)(graph_degree / 2 + 1)); + // Create an initial graph. The initial graph created here is not suitable for // searching, but connectivity is guaranteed. auto offset = raft::make_host_vector(small_graph_degree); @@ -2130,28 +2274,34 @@ auto iterative_build_graph(raft::resources const& res, } } - // Allocate memory for neighbors list using Transparent HugePage - constexpr size_t thp_size = 2 * 1024 * 1024; - size_t byte_size = sizeof(IdxT) * final_graph_size * topk; - if (byte_size % thp_size) { byte_size += thp_size - (byte_size % thp_size); } - mmap_owner neighbors_list(byte_size); - IdxT* neighbors_ptr = (IdxT*)neighbors_list.data(); - memset(neighbors_ptr, 0, byte_size); - bool flag_last = false; auto curr_graph_size = initial_graph_size; + + auto dev_graph = raft::make_device_matrix(res, 0, 0); + bool use_device_graph = false; + while (true) { auto start = std::chrono::high_resolution_clock::now(); auto curr_query_size = std::min(2 * curr_graph_size, final_graph_size); auto next_graph_degree = small_graph_degree; if (curr_graph_size == final_graph_size) { next_graph_degree = graph_degree; } + RAFT_LOG_INFO("Current graph size %lu: # current graph degree = %lu", + (uint64_t)curr_graph_size, + (uint64_t)next_graph_degree); // The search count (topk) is set to the next graph degree + 1, because // pruning is not used except in the last iteration. // (*) The appropriate setting for itopk_size requires careful consideration. - auto curr_topk = next_graph_degree + 1; - auto curr_itopk_size = next_graph_degree + 32; + auto curr_topk = next_graph_degree + 1; + // The configurable itopk (iter_params.itopk_size, 0 = auto) applies only to the true growing + // iterations, where the degree being built is small_graph_degree. When the graph reaches its + // full size the search builds a graph_degree-degree graph (topk = graph_degree + 1); that + // iteration needs a larger itopk, so it overrides the configured value with the auto formula. + // The final iteration (flag_last) uses a fixed itopk tied to the output topk. + auto curr_itopk_size = (iter_params.itopk_size > 0 && next_graph_degree == small_graph_degree) + ? (uint64_t)iter_params.itopk_size + : std::max(next_graph_degree + 32, (uint64_t)128); if (flag_last) { curr_topk = topk; curr_itopk_size = curr_topk + 32; @@ -2166,86 +2316,90 @@ auto iterative_build_graph(raft::resources const& res, (uint64_t)curr_itopk_size, (uint64_t)curr_topk); - cuvs::neighbors::cagra::search_params search_params; - search_params.algo = cuvs::neighbors::cagra::search_algo::AUTO; - search_params.max_queries = max_chunk_size; - search_params.itopk_size = curr_itopk_size; - - // Create an index (idx), a query view (dev_query_view), and a mdarray for - // search results (neighbors). - auto dev_dataset_view = raft::make_device_matrix_view( - dev_dataset.data_handle(), (int64_t)curr_graph_size, dev_dataset.extent(1)); - cuvs::neighbors::device_padded_dataset_view sub_padded(dev_dataset_view, - logical_dim); - - auto idx = cuvs::neighbors::cagra::device_padded_index( - res, params.metric, sub_padded, raft::make_const_mdspan(cagra_graph.view())); - - auto dev_query_view = raft::make_device_matrix_view( - dev_dataset.data_handle(), (int64_t)curr_query_size, dev_dataset.extent(1)); - - auto neighbors_view = - raft::make_host_matrix_view(neighbors_ptr, curr_query_size, curr_topk); + cuvs::neighbors::cagra::search_params search_params = iter_params; + search_params.max_queries = max_chunk_size; + search_params.itopk_size = curr_itopk_size; + + // Each index holds non-owning dataset and graph views. The local dataset owner and the graph + // passed to search_and_optimize keep those views alive for the duration of the search. + if (vpq_dataset != nullptr) { + auto idx = cuvs::neighbors::cagra::vpq_f16_index(res, params.metric); + idx.update_device_dataset_same_layout(res, vpq_dataset->as_dataset_view()); + if (use_device_graph) { + idx.update_graph(res, raft::make_const_mdspan(dev_graph.view())); + } else { + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + } - // Search. - // Since there are many queries, divide them into batches and search them. - auto query_batch = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - res, - dev_query_view.data_handle(), - static_cast(curr_query_size), - static_cast(dev_query_view.extent(1)), - max_chunk_size, - raft::resource::get_cuda_stream(res), - raft::resource::get_workspace_resource_ref(res)); - for (const auto& batch : query_batch) { - raft::device_matrix_view batch_dev_query_view; - if (query_contiguous) { - raft::copy_matrix(query_contiguous->data_handle(), - static_cast(logical_dim), - batch.data(), - dev_query_view.extent(1), - static_cast(logical_dim), - batch.size(), - raft::resource::get_cuda_stream(res)); - batch_dev_query_view = raft::make_device_matrix_view( - query_contiguous->data_handle(), batch.size(), static_cast(logical_dim)); + auto empty_query_view = + raft::make_device_matrix_view(static_cast(nullptr), 0, 0); + dev_graph = search_and_optimize(res, + search_params, + idx, + empty_query_view, + dev_neighbors.view(), + dev_distances.view(), + std::move(dev_graph), + vpq_dataset, + curr_query_size, + next_graph_degree, + curr_topk, + max_chunk_size, + static_cast(logical_dim), + flag_last && params.guarantee_connectivity); + } else { + auto dev_dataset_view = raft::make_device_matrix_view( + dev_dataset.data_handle(), static_cast(curr_graph_size), dev_dataset.extent(1)); + cuvs::neighbors::device_padded_dataset_view sub_padded(dev_dataset_view, + logical_dim); + auto idx = cuvs::neighbors::cagra::device_padded_index(res, params.metric); + idx.update_device_dataset_same_layout(res, sub_padded); + if (use_device_graph) { + idx.update_graph(res, raft::make_const_mdspan(dev_graph.view())); } else { - batch_dev_query_view = raft::make_device_matrix_view( - batch.data(), batch.size(), dev_query_view.extent(1)); + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); } - auto batch_dev_neighbors_view = raft::make_device_matrix_view( - dev_neighbors.data_handle(), batch.size(), curr_topk); - auto batch_dev_distances_view = raft::make_device_matrix_view( - dev_distances.data_handle(), batch.size(), curr_topk); - - cuvs::neighbors::cagra::search(res, - search_params, - idx, - batch_dev_query_view, - batch_dev_neighbors_view, - batch_dev_distances_view); - - auto batch_neighbors_view = raft::make_host_matrix_view( - neighbors_view.data_handle() + batch.offset() * curr_topk, batch.size(), curr_topk); - raft::copy(res, batch_neighbors_view, batch_dev_neighbors_view); - } - // Optimize graph - auto next_graph_size = curr_query_size; - cagra_graph = raft::make_host_matrix(0, 0); // delete existing grahp - cagra_graph = raft::make_host_matrix(next_graph_size, next_graph_degree); - optimize( - res, neighbors_view, cagra_graph.view(), flag_last ? params.guarantee_connectivity : 0); + auto dev_query_view = raft::make_device_matrix_view( + dev_dataset.data_handle(), static_cast(curr_query_size), dev_dataset.extent(1)); + dev_graph = search_and_optimize( + res, + search_params, + idx, + dev_query_view, + dev_neighbors.view(), + dev_distances.view(), + std::move(dev_graph), + static_cast*>(nullptr), + curr_query_size, + next_graph_degree, + curr_topk, + max_chunk_size, + static_cast(logical_dim), + flag_last && params.guarantee_connectivity); + } + use_device_graph = true; - auto end = std::chrono::high_resolution_clock::now(); - auto elapsed_ms = std::chrono::duration_cast(end - start).count(); + auto end = std::chrono::high_resolution_clock::now(); + [[maybe_unused]] auto elapsed_ms = + std::chrono::duration_cast(end - start).count(); RAFT_LOG_DEBUG("# elapsed time: %.3lf sec", (double)elapsed_ms / 1000); if (flag_last) { break; } - flag_last = (curr_graph_size == final_graph_size); - curr_graph_size = next_graph_size; + flag_last = (curr_graph_size == final_graph_size); + auto next_graph_size = curr_query_size; + curr_graph_size = next_graph_size; } + auto stream = raft::resource::get_cuda_stream(res); + + cagra_graph = raft::make_host_matrix(dev_graph.extent(0), dev_graph.extent(1)); + raft::copy(cagra_graph.data_handle(), + dev_graph.data_handle(), + dev_graph.extent(0) * dev_graph.extent(1), + stream); + raft::resource::sync_stream(res); + return cagra_graph; } diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp index 72f118e5c3..161ad34321 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp @@ -101,6 +101,7 @@ using search_single_cta_p_kernel_func_t = const std::uint32_t, const std::uint32_t, const dataset_descriptor_base_t*, + const IndexT, cagra_sample_filter); } // namespace single_cta_search diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh index 302bf4e9d8..47f02fb6ef 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh @@ -555,6 +555,7 @@ __device__ void search_single_cta_p_impl( const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id_offset, // Offset to add to query_id when calling filter const dataset_descriptor_base_t* dataset_desc, + const IndexT graph_size, cagra_sample_filter filter_payload) { using job_desc_type = job_desc_t>; @@ -629,7 +630,8 @@ __device__ void search_single_cta_p_impl( query_id, query_id_offset, dataset_desc, - filter_payload); + filter_payload, + graph_size); // make sure all writes are visible even for the host // (e.g. when result buffers are in pinned memory) diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in index 9986f7abc1..b003220497 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in @@ -51,6 +51,7 @@ extern "C" __global__ __launch_bounds__(1024, 1) void search_single_cta_p( const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id_offset, const dataset_desc_base* dataset_desc, + const index_t graph_size, cagra_sample_filter_t filter_payload) { search_single_cta_p_impl(this->dataset_size)); + // Bound random seed selection to the graph size, not the dataset size. + // During iterative / CAGRA-Q build the graph is smaller than the dataset, + // so using dataset_size here selects seeds that index past the graph end + // (out-of-bounds access). See https://github.com/rapidsai/cuvs/pull/1780. + static_cast(graph.extent(0))); std::shared_ptr compute_distance_to_child_nodes_launcher = make_cagra_multi_kernel_jit_launcher(graph.extent(0)), filter_payload); last_touch.store(std::chrono::system_clock::now(), std::memory_order_relaxed); diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 05c71e0213..6e73f36d10 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -279,6 +279,40 @@ auto deserialize_host_dense(raft::resources const& res, std::istream& is) return std::make_unique(std::move(storage), metadata.dim); } +/** VPQ codebooks are floating point; the encoded rows are always uint8 and carry no dtype. */ +template +constexpr auto vpq_wire_dtype() -> cudaDataType_t +{ + static_assert(std::is_same_v || std::is_same_v, + "serialize_vpq: codebook element type must be float or half"); + return std::is_same_v ? CUDA_R_16F : CUDA_R_32F; +} + +/** + * Write the payload of a VPQ dataset: six scalars followed by the two codebooks and the encoded + * rows. + * + * Stays on `raft::serialize_mdspan` rather than the `write_dense_bytes` scheme used by the dense + * path above, because `deserialize_vpq` reads with `raft::deserialize_mdspan`, which expects the + * NumPy header that helper embeds per matrix. The scalar types must also match the reader exactly: + * `n_rows` is `IdxT` and the remaining five are `uint32_t`. + */ +template +void serialize_vpq(raft::resources const& res, + std::ostream& os, + device_vpq_dataset const& dataset) +{ + raft::serialize_scalar(res, os, dataset.n_rows()); + raft::serialize_scalar(res, os, dataset.dim()); + raft::serialize_scalar(res, os, dataset.vq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_len()); + raft::serialize_scalar(res, os, dataset.encoded_row_length()); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.vq_code_book.view())); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.pq_code_book.view())); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.data.view())); +} + template auto deserialize_vpq(raft::resources const& res, std::istream& is) -> std::unique_ptr> @@ -305,6 +339,41 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) std::move(vq_code_book), std::move(pq_code_book), std::move(data)); } +/** + * Write a self-describing VPQ dataset blob: tag + codebook dtype + payload. + * + * The tag and dtype are deliberately written here rather than inside `serialize_vpq`, mirroring how + * `serialize_cagra_dense_dataset` wraps the dense payload, so that a reader can identify the blob + * before committing to a `DataT`. + */ +template +void serialize_vpq_dataset(raft::resources const& res, + std::ostream& os, + device_vpq_dataset const& dataset) +{ + raft::serialize_scalar(res, os, kSerializeVPQDataset); + raft::serialize_scalar(res, os, vpq_wire_dtype()); + serialize_vpq(res, os, dataset); +} + +/** Read a blob written by `serialize_vpq_dataset`, validating the tag and codebook dtype. */ +template +auto deserialize_vpq_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + const auto tag = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(tag == kSerializeVPQDataset, + "deserialize_vpq_dataset: expected VPQ tag (%u), got %u", + static_cast(kSerializeVPQDataset), + static_cast(tag)); + const auto dtype = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dtype == vpq_wire_dtype(), + "deserialize_vpq_dataset: codebook dtype (%d) does not match expected (%d)", + static_cast(dtype), + static_cast(vpq_wire_dtype())); + return deserialize_vpq(res, is); +} + template auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) -> std::unique_ptr diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..673e49759b 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -3,13 +3,20 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../neighbors/detail/dataset_serialize.hpp" +#include "../../util/serialize_validation.hpp" #include "./detail/pq.cuh" #include +#include #include #include +#include +#include +#include + namespace cuvs::preprocessing::quantize::pq { #define CUVS_INST_QUANTIZATION(T, QuantI) \ @@ -76,6 +83,55 @@ CUVS_INST_VPQ_BUILD(uint8_t); #undef CUVS_INST_VPQ_BUILD +void serialize(raft::resources const& res, + std::ostream& os, + const cuvs::neighbors::device_vpq_dataset& dataset) +{ + // Same file preamble as cagra::serialize. The nested blob carries only a kind tag and dtype, + // matching serialize_cagra_dense_dataset, because a nested blob relies on its enclosing file for + // the version; a standalone .vpq has no enclosing file, so the version is written here. + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + os << dtype_string; + raft::serialize_scalar(res, os, vpq_serialization_version); + ::cuvs::neighbors::detail::serialize_vpq_dataset(res, os, dataset); +} + +void serialize(raft::resources const& res, + const std::string& filename, + const cuvs::neighbors::device_vpq_dataset& dataset) +{ + std::ofstream os(filename, std::ios::out | std::ios::binary | std::ios::trunc); + RAFT_EXPECTS(os.good(), "pq::serialize: cannot open %s for writing", filename.c_str()); + serialize(res, os, dataset); +} + +void deserialize(raft::resources const& res, + std::istream& is, + std::unique_ptr>* out_dataset) +{ + RAFT_EXPECTS(out_dataset != nullptr, "pq::deserialize: out_dataset must not be null"); + char dtype_string[4]; + RAFT_EXPECTS(is.read(dtype_string, 4), "pq::deserialize: failed to read the dtype prefix"); + RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), + "pq::deserialize: dtype prefix does not match a VPQ dataset with half codebooks"); + auto const version = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(version == vpq_serialization_version, + "pq::deserialize: serialization version mismatch, expected %d, got %d", + vpq_serialization_version, + version); + *out_dataset = ::cuvs::neighbors::detail::deserialize_vpq_dataset(res, is); +} + +void deserialize(raft::resources const& res, + const std::string& filename, + std::unique_ptr>* out_dataset) +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.good(), "pq::deserialize: cannot open %s for reading", filename.c_str()); + deserialize(res, is, out_dataset); +} + namespace detail { template diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index b4a657c90c..c8e64a7ab6 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -199,7 +199,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST - PATH neighbors/ann_cagra/test_float_uint32_t.cu + PATH neighbors/ann_cagra/test_float_uint32_t.cu neighbors/ann_cagra/test_iterative_cagra_q.cu GPUS 1 PERCENT 100 ) @@ -415,6 +415,7 @@ ConfigureTest( preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu preprocessing/product_quantization.cu + preprocessing/vpq_serialization.cu preprocessing/pca.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu index 07835c50e8..6d10797ef3 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu @@ -7,6 +7,7 @@ #include "../cagra_padded_build_helpers.cuh" #include +#include #include #include @@ -23,9 +24,10 @@ class CagraIterativeBuildBugTest : public ::testing::Test { using data_type = DataT; protected: - void run() + // The bug manifests when graph_degree is equal to intermediate_graph_degree + // see issue https://github.com/rapidsai/cuvs/issues/1818 + static auto bug_index_params() -> cagra::index_params { - // Set up iterative CAGRA graph building cagra::index_params index_params; // The bug manifests when graph_degree is equal to intermediate_graph_degree // see issue https://github.com/nvidia/cuvs/issues/1818 @@ -34,6 +36,12 @@ class CagraIterativeBuildBugTest : public ::testing::Test { // Use iterative CAGRA search for graph building index_params.graph_build_params = graph_build_params::iterative_search_params(); + return index_params; + } + + void run() + { + auto index_params = bug_index_params(); cuvs::neighbors::test::padded_device_matrix_for_cagra padded( res, raft::make_const_mdspan(dataset->view())); @@ -46,6 +54,30 @@ class CagraIterativeBuildBugTest : public ::testing::Test { ASSERT_EQ(cagra_index.dim(), n_dim); } + // Same bug, reached through iterative CAGRA-Q: the graph is built from a PQ-compressed dataset, + // so the searches driving the build run on compressed rows instead of dense ones. + void run_compressed() + { + cuvs::neighbors::vpq_params vpq_params; + // pq_len = n_dim / pq_dim must be 2, 4 or 8 for CAGRA-Q. Codebook quality is irrelevant here, + // since only graph construction is under test, so training stays short. + vpq_params.pq_dim = static_cast(n_dim / 4); + vpq_params.vq_n_centers = 64; + vpq_params.kmeans_n_iters = 5; + + auto compressed = cuvs::preprocessing::quantize::pq::make_vpq_dataset( + res, vpq_params, raft::make_const_mdspan(dataset->view())); + + // No padding and no attach step: a compressed dataset is read through its own view, which the + // index holds on to, so `compressed` has to outlive the index. + auto cagra_index = cagra::build(res, bug_index_params(), compressed.as_dataset_view()); + raft::resource::sync_stream(res); + + ASSERT_GT(cagra_index.size(), 0); + ASSERT_EQ(cagra_index.dim(), n_dim); + ASSERT_EQ(cagra_index.graph_degree(), 16u); + } + void SetUp() override { dataset.emplace(raft::make_device_matrix(res, n_samples, n_dim)); @@ -85,4 +117,9 @@ TYPED_TEST_SUITE(CagraIterativeBuildBugTest, TestTypes); TYPED_TEST(CagraIterativeBuildBugTest, IterativeBuildTest) { this->run(); } +TYPED_TEST(CagraIterativeBuildBugTest, IterativeBuildFromCompressedDatasetTest) +{ + this->run_compressed(); +} + } // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu new file mode 100644 index 0000000000..bad504b589 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -0,0 +1,280 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * Iterative CAGRA-Q: building a CAGRA graph directly from a PQ-compressed dataset. + * + * This path takes a `device_vpq_dataset_view` instead of dense rows, so the caller owns + * compression and the inner searches of the iterative build run against the compressed data. + * It only accepts `L2Expanded`, `pq_bits == 8` and `pq_len` in {2, 4, 8}, so it does not fit + * the dtype-templated suites in ann_cagra.cuh and lives in its own file. + * + * What is checked here is that a compressed dataset, freshly compressed or loaded from disk, + * builds a usable graph, and that the constraints above are rejected rather than accepted and + * quietly ignored. Serialization fidelity itself is covered by preprocessing/vpq_serialization.cu. + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +namespace { + +auto compress(const raft::resources& res, + raft::device_matrix_view dataset, + uint32_t pq_dim, + uint32_t pq_bits = 8) -> vpq_dataset_t +{ + cuvs::neighbors::vpq_params params; + params.pq_dim = pq_dim; + params.pq_bits = pq_bits; + params.vq_n_centers = 32; + params.kmeans_n_iters = 5; // Codebooks need to be well defined here, not optimal. + return cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, params, dataset); +} + +auto iterative_params(uint32_t graph_degree = 32) -> index_params +{ + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = graph_degree; + params.intermediate_graph_degree = graph_degree * 2; + params.graph_build_params = graph_build_params::iterative_search_params(); + return params; +} + +/** + * Fraction of queries that retrieve their own row, where the queries are dataset rows. + * + * A sanity signal rather than a quality metric: search quality is a benchmark's job, and an + * iterative build varies by a few points run to run over identical input, so callers assert a + * loose floor that only a broken dataset would miss. + */ +template +auto self_recall_at_1(const raft::resources& res, + const IndexT& idx, + raft::device_matrix_view queries) -> double +{ + constexpr int64_t k = 10; + const auto n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + + search_params params; + params.itopk_size = 64; + search(res, params, idx, queries, neighbors.view(), distances.view()); + + std::vector ids(static_cast(n_queries * k)); + raft::copy(ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + int64_t hits = 0; + for (int64_t q = 0; q < n_queries; q++) { + if (ids[q * k] == static_cast(q)) { hits++; } + } + return static_cast(hits) / static_cast(n_queries); +} + +} // namespace + +struct CagraQInputs { + int64_t n_rows; + int64_t dim; + uint32_t pq_dim; // pq_len = dim / pq_dim, which must land in {2, 4, 8} +}; + +std::ostream& operator<<(std::ostream& os, const CagraQInputs& in) +{ + return os << "n_rows:" << in.n_rows << " dim:" << in.dim << " pq_dim:" << in.pq_dim + << " pq_len:" << (in.dim / in.pq_dim); +} + +/** Shared clustered dataset; each test compresses it itself so nothing leaks between cases. */ +class CagraQCompressedTestBase : public ::testing::Test { + protected: + void make_dataset(int64_t n_rows, int64_t dim) + { + dataset_.emplace(raft::make_device_matrix(res_, n_rows, dim)); + auto labels = raft::make_device_vector(res_, n_rows); + raft::random::make_blobs(res_, + dataset_->view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + 1234ULL); + raft::resource::sync_stream(res_); + } + + auto dataset() -> raft::device_matrix_view + { + return raft::make_const_mdspan(dataset_->view()); + } + + /** The first rows of the dataset, reused as queries. */ + auto queries(int64_t n_queries) -> raft::device_matrix_view + { + return raft::make_device_matrix_view( + dataset_->data_handle(), std::min(n_queries, dataset_->extent(0)), dataset_->extent(1)); + } + + void TearDown() override + { + dataset_.reset(); + raft::resource::sync_stream(res_); + } + + raft::resources res_; + std::optional> dataset_ = std::nullopt; +}; + +class CagraQBuildTest : public CagraQCompressedTestBase, + public ::testing::WithParamInterface { + protected: + void SetUp() override + { + params_ = GetParam(); + make_dataset(params_.n_rows, params_.dim); + } + + CagraQInputs params_{}; +}; + +TEST_P(CagraQBuildTest, BuildsAndSearchesAFreshlyCompressedDataset) +{ + auto compressed = compress(res_, dataset(), params_.pq_dim); + ASSERT_EQ(compressed.pq_len(), static_cast(params_.dim / params_.pq_dim)); + + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + ASSERT_EQ(idx.size(), params_.n_rows); + ASSERT_EQ(idx.dim(), params_.dim); + ASSERT_EQ(idx.graph_degree(), 32u); + + // Searchable straight after build: the index keeps the compressed view it was built from, so + // unlike a dense standard-layout dataset there is nothing to attach first. + EXPECT_GT(self_recall_at_1(res_, idx, queries(1000)), 0.5); +} + +TEST_P(CagraQBuildTest, BuildsFromADeserializedDataset) +{ + auto compressed = compress(res_, dataset(), params_.pq_dim); + + // The path the benchmarks take: compress offline, store, then build from the file. + std::stringstream stored; + cuvs::preprocessing::quantize::pq::serialize(res_, stored, compressed); + std::unique_ptr loaded; + cuvs::preprocessing::quantize::pq::deserialize(res_, stored, &loaded); + ASSERT_NE(loaded, nullptr); + + auto idx = cagra::build(res_, iterative_params(), loaded->as_dataset_view()); + ASSERT_EQ(idx.size(), params_.n_rows); + ASSERT_EQ(idx.dim(), params_.dim); + EXPECT_GT(self_recall_at_1(res_, idx, queries(1000)), 0.5); +} + +TEST_P(CagraQBuildTest, PromotesUnsetGraphBuildParamsToIterative) +{ + auto compressed = compress(res_, dataset(), params_.pq_dim); + + // Iterative search is the only construction a compressed dataset supports, so leaving + // graph_build_params at its default must select it, not fall back to IVF-PQ or NN-descent. + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = 32; + ASSERT_TRUE(std::holds_alternative(params.graph_build_params)); + + auto idx = cagra::build(res_, params, compressed.as_dataset_view()); + EXPECT_EQ(idx.size(), params_.n_rows); +} + +INSTANTIATE_TEST_CASE_P(CagraQBuildTests, + CagraQBuildTest, + ::testing::ValuesIn(std::vector{ + {2000, 64, 32}, // pq_len 2 + {2000, 128, 32}, // pq_len 4 + {2000, 256, 32}, // pq_len 8 + })); + +/** The constraints the VPQ build overload documents, each of which must be rejected loudly. */ +class CagraQContractTest : public CagraQCompressedTestBase { + protected: + void SetUp() override { make_dataset(n_rows, dim); } + + static constexpr int64_t n_rows = 1000; + static constexpr int64_t dim = 64; +}; + +TEST_F(CagraQContractTest, RejectsNonIterativeGraphBuilder) +{ + auto compressed = compress(res_, dataset(), 32); + auto params = iterative_params(); + params.graph_build_params = + graph_build_params::nn_descent_params(params.intermediate_graph_degree); + EXPECT_THROW(cagra::build(res_, params, compressed.as_dataset_view()), raft::exception); +} + +TEST_F(CagraQContractTest, RejectsMetricOtherThanL2Expanded) +{ + auto compressed = compress(res_, dataset(), 32); + auto params = iterative_params(); + params.metric = cuvs::distance::DistanceType::InnerProduct; + EXPECT_THROW(cagra::build(res_, params, compressed.as_dataset_view()), raft::exception); +} + +TEST_F(CagraQContractTest, RejectsPqBitsOtherThan8) +{ + auto compressed = compress(res_, dataset(), 32, /* pq_bits */ 6); + ASSERT_EQ(compressed.pq_bits(), 6u); + EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), + raft::exception); +} + +TEST_F(CagraQContractTest, RejectsPqLenOutsideSupportedSet) +{ + auto compressed = compress(res_, dataset(), /* pq_dim */ 4); // pq_len = 64 / 4 = 16 + ASSERT_EQ(compressed.pq_len(), 16u); + EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), + raft::exception); +} + +TEST_F(CagraQContractTest, RejectsEmptyDataset) +{ + // Hand-built rather than compressed, since make_vpq_dataset rejects an empty input of its own + // accord. Every other constraint is satisfied so that only the emptiness can trip. + const auto width = static_cast(dim); + auto vq_code_book = raft::make_device_matrix(res_, 1, width); + auto pq_code_book = raft::make_device_matrix(res_, 256, 2); + auto codes = raft::make_device_matrix(res_, 0, 4 + dim / 2); + vpq_dataset_t empty{std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; + ASSERT_EQ(empty.n_rows(), 0); + + EXPECT_THROW(cagra::build(res_, iterative_params(), empty.as_dataset_view()), raft::exception); +} + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/preprocessing/vpq_serialization.cu b/cpp/tests/preprocessing/vpq_serialization.cu new file mode 100644 index 0000000000..6e41cb48c7 --- /dev/null +++ b/cpp/tests/preprocessing/vpq_serialization.cu @@ -0,0 +1,260 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../neighbors/vpq_utils.cuh" +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::pq { + +using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +struct VpqSerializationInputs { + int64_t n_rows; + int64_t dim; + uint32_t pq_bits; + uint32_t pq_dim; + uint32_t vq_n_centers; // 0 lets the heuristic choose + uint64_t seed; +}; + +std::ostream& operator<<(std::ostream& os, const VpqSerializationInputs& in) +{ + return os << "n_rows:" << in.n_rows << " dim:" << in.dim << " pq_bits:" << in.pq_bits + << " pq_dim:" << in.pq_dim << " vq_n_centers:" << in.vq_n_centers + << " seed:" << in.seed; +} + +template +auto to_host(const raft::resources& res, raft::device_matrix_view m) -> std::vector +{ + std::vector host(static_cast(m.extent(0)) * static_cast(m.extent(1))); + raft::copy(host.data(), m.data_handle(), host.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return host; +} + +/** Bitwise, not approximate: serialization is expected not to perturb a single bit. */ +template +void expect_same_bits(const raft::resources& res, + raft::device_matrix_view expected, + raft::device_matrix_view actual, + const char* what) +{ + ASSERT_EQ(expected.extent(0), actual.extent(0)) << what; + ASSERT_EQ(expected.extent(1), actual.extent(1)) << what; + const auto lhs = to_host(res, expected); + const auto rhs = to_host(res, actual); + EXPECT_EQ(0, std::memcmp(lhs.data(), rhs.data(), lhs.size() * sizeof(T))) << what; +} + +class VpqSerializationTest : public ::testing::TestWithParam { + public: + VpqSerializationTest() + : params_(::testing::TestWithParam::GetParam()), + dataset_(raft::make_device_matrix(res_, params_.n_rows, params_.dim)) + { + } + + protected: + void SetUp() override + { + auto labels = raft::make_device_vector(res_, params_.n_rows); + raft::random::make_blobs(res_, + dataset_.view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + params_.seed); + raft::resource::sync_stream(res_); + } + + auto compress() -> vpq_dataset_t + { + cuvs::neighbors::vpq_params vpq; + vpq.pq_bits = params_.pq_bits; + vpq.pq_dim = params_.pq_dim; + vpq.vq_n_centers = params_.vq_n_centers; + // The codebooks only have to be well defined here, not good, so keep training short. + vpq.kmeans_n_iters = 5; + return make_vpq_dataset(res_, vpq, raft::make_const_mdspan(dataset_.view())); + } + + void expect_equivalent(const vpq_dataset_t& expected, const vpq_dataset_t& actual) + { + ASSERT_EQ(expected.n_rows(), actual.n_rows()); + ASSERT_EQ(expected.dim(), actual.dim()); + ASSERT_EQ(expected.vq_n_centers(), actual.vq_n_centers()); + ASSERT_EQ(expected.pq_n_centers(), actual.pq_n_centers()); + ASSERT_EQ(expected.pq_len(), actual.pq_len()); + ASSERT_EQ(expected.encoded_row_length(), actual.encoded_row_length()); + ASSERT_EQ(expected.pq_bits(), actual.pq_bits()); + ASSERT_EQ(expected.pq_dim(), actual.pq_dim()); + + expect_same_bits(res_, + raft::make_const_mdspan(expected.vq_code_book.view()), + raft::make_const_mdspan(actual.vq_code_book.view()), + "vq_code_book"); + expect_same_bits(res_, + raft::make_const_mdspan(expected.pq_code_book.view()), + raft::make_const_mdspan(actual.pq_code_book.view()), + "pq_code_book"); + expect_same_bits(res_, + raft::make_const_mdspan(expected.data.view()), + raft::make_const_mdspan(actual.data.view()), + "encoded rows"); + } + + /** + * Decodes both datasets and compares the reconstructions, which checks that a kernel can consume + * the deserialized extents and strides rather than only that the numbers match. + */ + void expect_same_decoded(const vpq_dataset_t& expected, const vpq_dataset_t& actual) + { + if (expected.pq_bits() != 8) { return; } // decode_vpq_dataset implements pq_bits == 8 only + auto stream = raft::resource::get_cuda_stream(res_); + auto lhs = raft::make_device_matrix(res_, expected.n_rows(), expected.dim()); + auto rhs = raft::make_device_matrix(res_, actual.n_rows(), actual.dim()); + cuvs::neighbors::decode_vpq_dataset(lhs.view(), expected, stream); + cuvs::neighbors::decode_vpq_dataset(rhs.view(), actual, stream); + raft::resource::sync_stream(res_); + expect_same_bits(res_, + raft::make_const_mdspan(lhs.view()), + raft::make_const_mdspan(rhs.view()), + "decoded rows"); + } + + raft::resources res_; + VpqSerializationInputs params_; + raft::device_matrix dataset_; +}; + +TEST_P(VpqSerializationTest, RoundTrip) +{ + auto original = compress(); + + { + SCOPED_TRACE("through a stream"); + std::stringstream stream; + serialize(res_, stream, original); + std::unique_ptr restored; + deserialize(res_, stream, &restored); + ASSERT_NE(restored, nullptr); + expect_equivalent(original, *restored); + expect_same_decoded(original, *restored); + } + + { + SCOPED_TRACE("through a file"); + const std::string path = "cuvs_vpq_serialization_test.bin"; + serialize(res_, path, original); + std::unique_ptr restored; + deserialize(res_, path, &restored); + std::remove(path.c_str()); + ASSERT_NE(restored, nullptr); + expect_equivalent(original, *restored); + } +} + +// Named for this suite rather than `inputs`: product_quantization.cu declares a variable template of +// that name in this same namespace, which would collide under a unity build. +const std::vector vpq_serialization_inputs = { + // pq_len = dim / pq_dim of 2, 4 and 8: the three values CAGRA-Q accepts. + {1000, 64, 8, 32, 0, 42ULL}, + {1000, 128, 8, 32, 0, 42ULL}, + {1000, 256, 8, 32, 0, 42ULL}, + // An explicit VQ codebook size rather than the heuristic. + {2000, 128, 8, 64, 64, 42ULL}, + // pq_bits below 8 packs several codes per byte, so encoded_row_length stops being pq_dim. + {500, 96, 6, 24, 0, 42ULL}, + {500, 32, 4, 16, 0, 42ULL}, +}; + +INSTANTIATE_TEST_CASE_P(VpqSerializationTests, + VpqSerializationTest, + ::testing::ValuesIn(vpq_serialization_inputs)); + +/** Writes the preamble that `serialize` emits, so only the field under test differs. */ +static void write_preamble(const raft::resources& res, std::ostream& os, int version) +{ + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + os << dtype_string; + raft::serialize_scalar(res, os, version); +} + +TEST(VpqSerialization, RejectsEmptyStream) +{ + raft::resources res; + std::stringstream stream; + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsForeignDtypePrefix) +{ + raft::resources res; + std::stringstream stream; + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + stream << dtype_string; + raft::serialize_scalar(res, stream, vpq_serialization_version); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsFutureVersion) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version + 1); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsTruncatedPayload) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version); + // A correct preamble followed by nothing: the payload reader must fail rather than return a + // dataset built from whatever the scalars happened to deserialize to. + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsNullOutParameter) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version); + EXPECT_THROW(deserialize(res, stream, nullptr), raft::exception); +} + +} // namespace cuvs::preprocessing::quantize::pq diff --git a/python/cuvs/cuvs/common/dataset.pxd b/python/cuvs/cuvs/common/dataset.pxd index ac2d76ec18..742f783766 100644 --- a/python/cuvs/cuvs/common/dataset.pxd +++ b/python/cuvs/cuvs/common/dataset.pxd @@ -14,6 +14,7 @@ cdef extern from "cuvs/core/dataset.h" nogil: ctypedef enum cuvsDatasetLayout_t: CUVS_DATASET_LAYOUT_STANDARD CUVS_DATASET_LAYOUT_PADDED + CUVS_DATASET_LAYOUT_VPQ ctypedef enum cuvsDatasetMemType_t: CUVS_DATASET_MEM_TYPE_HOST @@ -23,6 +24,10 @@ cdef extern from "cuvs/core/dataset.h" nogil: pass ctypedef cuvsDataset* cuvsDataset_t + cdef struct cuvsCagraCompressionParams: + pass + ctypedef cuvsCagraCompressionParams* cuvsCagraCompressionParams_t + cuvsError_t cuvsDatasetCreate(cuvsDataset_t* dataset) cuvsError_t cuvsDatasetMakePadded(cuvsResources_t res, @@ -38,6 +43,12 @@ cdef extern from "cuvs/core/dataset.h" nogil: DLManagedTensor* dataset, cuvsDataset_t* standard_dataset) + cuvsError_t cuvsDatasetMakeVpq( + cuvsResources_t res, + cuvsCagraCompressionParams_t params, + cuvsDataset_t dataset, + cuvsDataset_t* vpq_dataset) + cuvsError_t cuvsDatasetDestroy(cuvsDataset_t dataset) cuvsError_t cuvsDatasetGetMemType(cuvsDataset_t dataset, diff --git a/python/cuvs/cuvs/common/dataset.pyx b/python/cuvs/cuvs/common/dataset.pyx index 0c83633d13..b27f061b56 100644 --- a/python/cuvs/cuvs/common/dataset.pyx +++ b/python/cuvs/cuvs/common/dataset.pyx @@ -46,6 +46,8 @@ cdef class Dataset: if self.dataset == NULL: return None check_cuvs(cuvsDatasetGetLayout(self.dataset, &layout)) + if layout == CUVS_DATASET_LAYOUT_VPQ: + return "vpq" if layout == CUVS_DATASET_LAYOUT_PADDED: return "padded" return "standard" diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 9e4dbdb6f3..bd838c8754 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -43,6 +43,20 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: ITERATIVE_CAGRA_SEARCH ACE + ctypedef struct cuvsCagraCompressionParams: + uint32_t pq_bits + uint32_t pq_dim + uint32_t vq_n_centers + uint32_t kmeans_n_iters + double vq_kmeans_trainset_fraction + double pq_kmeans_trainset_fraction + + ctypedef cuvsCagraCompressionParams* cuvsCagraCompressionParams_t + + cuvsError_t cuvsCagraCompressionParamsCreate( + cuvsCagraCompressionParams_t* params) + cuvsError_t cuvsCagraCompressionParamsDestroy( + cuvsCagraCompressionParams_t params) ctypedef struct cuvsIvfPqParams: cuvsIvfPqIndexParams_t ivf_pq_build_params cuvsIvfPqSearchParams_t ivf_pq_search_params diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index dd481df259..beeff0b93a 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -474,6 +474,9 @@ def build(IndexParams index_params, dataset, resources=None): Supported dtype [float, half, int8, uint8] **Note:** For ACE build algorithm, the dataset MUST be in host memory. Use NumPy arrays or call .get() on CuPy arrays before passing. + A ``Dataset`` with ``layout == "vpq"`` builds an iterative CAGRA-Q + index and requires ``metric="sqeuclidean"`` plus + ``build_algo="iterative_cagra_search"``. {resources_docstring} Returns @@ -527,7 +530,10 @@ def build(IndexParams index_params, dataset, resources=None): dl_data_type_to_numpy(idx.index.dtype)).name idx._dataset_source = dataset_obj - if not is_ace_build: + if dataset_obj.layout == "vpq": + _keep_dataset_alive(idx, dataset_obj) + idx._dataset_source = None + elif not is_ace_build: if (dataset_obj.layout == "padded" and dataset_obj.memory_type == "device" and dataset_obj.is_owning): diff --git a/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py b/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py index 7db0c383fd..833caedead 100644 --- a/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py +++ b/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py @@ -1,12 +1,17 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from .pq import Quantizer, QuantizerParams, build, transform, inverse_transform +from .pq import ( + Quantizer, QuantizerParams, VpqParams, build, inverse_transform, + make_vpq_dataset, transform, +) __all__ = [ "Quantizer", "QuantizerParams", + "VpqParams", "build", "transform", - "inverse_transform" + "inverse_transform", + "make_vpq_dataset", ] diff --git a/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx b/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx index fd1a5326e0..67749fae41 100644 --- a/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx +++ b/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -7,6 +7,16 @@ import numpy as np from cuvs.common cimport cydlpack +from cuvs.common.dataset cimport ( + Dataset, + cuvsDatasetMakeStandardView, + cuvsDatasetMakeVpq, +) +from cuvs.neighbors.cagra.cagra cimport ( + cuvsCagraCompressionParams, + cuvsCagraCompressionParamsCreate, + cuvsCagraCompressionParamsDestroy, +) from pylibraft.common import auto_convert_output, device_ndarray from pylibraft.common.cai_wrapper import wrap_array @@ -22,6 +32,32 @@ PQ_KMEANS_TYPES = { PQ_KMEANS_NAMES = {v: k for k, v in PQ_KMEANS_TYPES.items()} + +cdef class VpqParams: + """Parameters for creating a CAGRA-Q VPQ dataset.""" + + cdef cuvsCagraCompressionParams* params + + def __cinit__(self): + self.params = NULL + check_cuvs(cuvsCagraCompressionParamsCreate(&self.params)) + + def __dealloc__(self): + if self.params != NULL: + cuvsCagraCompressionParamsDestroy(self.params) + + def __init__(self, *, pq_bits=8, pq_dim=0, vq_n_centers=0, + kmeans_n_iters=25, vq_kmeans_trainset_fraction=0.0, + pq_kmeans_trainset_fraction=0.0): + self.params.pq_bits = pq_bits + self.params.pq_dim = pq_dim + self.params.vq_n_centers = vq_n_centers + self.params.kmeans_n_iters = kmeans_n_iters + self.params.vq_kmeans_trainset_fraction = \ + vq_kmeans_trainset_fraction + self.params.pq_kmeans_trainset_fraction = \ + pq_kmeans_trainset_fraction + cdef class QuantizerParams: """ Parameters for product quantization @@ -377,3 +413,29 @@ def inverse_transform(Quantizer quantizer, codes, output=None, vq_labels=None, r vq_labels_dlpack)) return output + + +@auto_sync_resources +def make_vpq_dataset(VpqParams params, dataset, resources=None): + """Create an owning device VPQ dataset for iterative CAGRA-Q.""" + cdef Dataset dense + cdef Dataset vpq = Dataset() + cdef cuvsResources_t res = resources.get_c_obj() + cdef cydlpack.DLManagedTensor* dataset_dlpack = NULL + + if isinstance(dataset, Dataset): + dense = dataset + else: + dataset_ai = wrap_array(dataset) + _check_input_array( + dataset_ai, + [np.dtype("float32"), np.dtype("float16"), + np.dtype("int8"), np.dtype("uint8")]) + dataset_dlpack = cydlpack.dlpack_c(dataset_ai) + dense = Dataset() + check_cuvs(cuvsDatasetMakeStandardView( + res, dataset_dlpack, &dense.dataset)) + + check_cuvs(cuvsDatasetMakeVpq( + res, params.params, dense.dataset, &vpq.dataset)) + return vpq