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..68275ad7d0 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); }); } @@ -1011,7 +1109,7 @@ static auto read_serialized_header(cuvsResources_t res, const char *filename) "serialization version mismatch, expected %d, got %d", cuvs::neighbors::cagra::cagra_serialization_version, version); using kind = cuvs::neighbors::cagra::serialized_dataset_kind; - RAFT_EXPECTS(dataset_kind_raw <= static_cast(kind::host_standard), + RAFT_EXPECTS(dataset_kind_raw <= static_cast(kind::device_vpq_f16), "Invalid serialized dataset kind %u in file %s", dataset_kind_raw, filename); return {output_dtype, static_cast(dataset_kind_raw)}; @@ -1058,6 +1156,12 @@ void dispatch_serialized_dataset_kind( fn.template operator()< cuvs::neighbors::device_padded_dataset_view>(); break; + case serialized_kind::device_vpq_f16: + // A recognised file the C API has no index layout for, as opposed to an unreadable one. + // cuvsDatasetLayout_t covers standard and padded only, and every C entry point dispatches + // on that layout, so there is nothing here to hand a VPQ index to yet. + RAFT_FAIL("File holds a VPQ-compressed (CAGRA-Q) dataset, which the C API has no dataset " + "layout for; load it through the C++ API"); } } @@ -1122,8 +1226,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 +1306,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 +1694,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 +1887,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 +1993,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/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..06de5a6a83 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 @@ -2252,7 +2307,7 @@ void search( * @{ */ -/** Dense dataset storage kind recorded in a serialized CAGRA index. */ +/** Dataset storage kind recorded in a serialized CAGRA index. */ enum class serialized_dataset_kind : std::uint32_t { /** The serialized index does not contain a dataset payload. */ none = 0, @@ -2264,16 +2319,19 @@ enum class serialized_dataset_kind : std::uint32_t { host_padded = 3, /** Host-resident dataset using its standard row layout. */ host_standard = 4, + /** Device-resident VPQ-compressed dataset with f16 codebooks (CAGRA-Q). */ + device_vpq_f16 = 5, }; /** Current experimental CAGRA serialization format version. */ inline constexpr int cagra_serialization_version = 6; -// Serialize and deserialize are overloaded for device/host and padded/standard dense indexes. -// They use the same strided dataset payload; the serialized dataset kind selects the matching -// owning dataset type during deserialization. To support a new dataset kind (e.g. vpq_f16_index), -// add matching overloads here and a corresponding deserialize_ in -// detail/dataset_serialize.hpp (dense views use serialize_cagra_dense_dataset). +// Serialize and deserialize are overloaded for device/host and padded/standard dense indexes, +// which share the same strided dataset payload, and for vpq_f16_index, which writes a VPQ payload +// instead. The serialized dataset kind selects the matching owning dataset type during +// deserialization. To support a further kind, add matching overloads here and a corresponding +// serialize_/deserialize_ in detail/dataset_serialize.hpp (dense views use +// serialize_cagra_dense_dataset, VPQ ones serialize_vpq_dataset). /** * Save the index to file. @@ -2824,6 +2882,106 @@ void deserialize(raft::resources const& handle, std::unique_ptr>* out_dataset = nullptr); +/* vpq_f16_index overloads (CAGRA-Q). + * + * The compressed rows travel with the index, so that a deserialized index can be searched without + * the dense dataset it was compressed from and without retraining the codebooks. As everywhere + * else, the index holds a view: `deserialize` returns the owning dataset through `out_dataset`, + * which the caller has to keep alive for as long as the index is used. + * + * Unlike the dense overloads, `out_dataset` is required. Nothing can be searched in a VPQ index + * whose rows were dropped, so there is no use for a graph-only load, and asking for one is an + * error rather than a silently unusable index. For the same reason `include_dataset = false` + * produces an index that only `update_dataset` can make searchable again. + */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + /** @copydoc serialize */ void serialize(raft::resources const& handle, const std::string& filename, @@ -4613,14 +4771,37 @@ namespace helpers { * @param[in] intermediate_graph_degree degree of the input graph for the optimization process * @param[in] index_size * @param[in] mst_optimize whether to use MST optimization + * @param[in] device_resident_graphs whether the input and output graphs are already device + * resident. In that case `optimize` reads and writes them in place instead of staging + * them through device buffers, so those staging allocations are left out of the + * estimate. * @return tuple of [host_size, device_size, host_fixed_size, device_fixed_size] memory sizes in * bytes */ -std::tuple optimize_workspace_size(size_t n_rows, - size_t graph_degree, - size_t intermediate_degree, - size_t index_size, - bool mst_optimize = false); +std::tuple optimize_workspace_size( + size_t n_rows, + size_t graph_degree, + size_t intermediate_degree, + size_t index_size, + bool mst_optimize = false, + bool device_resident_graphs = false); + +/** + * Calculate the device memory footprint of a VPQ-compressed (CAGRA-Q) dataset. + * + * The footprint is the sum of the VQ codebook, the PQ codebook and the encoded rows. Parameters + * left at 0 are resolved with the same heuristics that `vpq_build` applies. + * + * @param[in] dataset shape of the uncompressed dataset + * @param[in] params VPQ compression parameters + * @param[in] codebook_element_size size in bytes of a codebook element (2 for the f16 codebooks + * used by CAGRA-Q) + * + * @return compressed dataset size in bytes + */ +size_t vpq_dataset_size(raft::matrix_extent dataset, + cuvs::neighbors::vpq_params params, + size_t codebook_element_size = 2); /** * Calculate memory usage of CAGRA build. @@ -4630,13 +4811,17 @@ std::tuple optimize_workspace_size(size_t n_rows * @param[in] dtype element type of the dataset * (e.g. `CUDA_R_32F`, `CUDA_R_16F`, `CUDA_R_8I`, `CUDA_R_8U`) * @param[in] cparams CAGRA index building parameters + * @param[in] compression when set, the build consumes a VPQ-compressed (CAGRA-Q) dataset with + * these parameters rather than the dense dataset described by `dataset` and `dtype` * * @return pair of [host_size, device_size] memory sizes in bytes */ -std::pair cagra_build_mem_usage(raft::resources const& res, - raft::matrix_extent dataset, - cudaDataType_t dtype, - cuvs::neighbors::cagra::index_params cparams); +std::pair cagra_build_mem_usage( + raft::resources const& res, + raft::matrix_extent dataset, + cudaDataType_t dtype, + cuvs::neighbors::cagra::index_params cparams, + std::optional compression = std::nullopt); /** * @brief Optimize a KNN graph into a CAGRA graph. 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/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index 83d047560b..9d7614e498 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -155,6 +155,43 @@ namespace cuvs::neighbors::cagra { cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::vpq_f16_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, filename, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::vpq_f16_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, out_dataset); \ + } \ + \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::vpq_f16_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, os, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::vpq_f16_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ + } \ + \ void serialize_to_hnswlib( \ raft::resources const& handle, \ std::ostream& os, \ diff --git a/cpp/src/neighbors/cagra_serialize_inst.cu.in b/cpp/src/neighbors/cagra_serialize_inst.cu.in index 3d34adb36f..58e555d17e 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/cpp/src/neighbors/cagra_serialize_inst.cu.in @@ -12,6 +12,7 @@ namespace { using data_t = @data_type@; 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_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; } // namespace @@ -21,6 +22,8 @@ extern template void index::compute raft::resources const&); extern template void index::compute_dataset_norms_( raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); CUVS_INST_CAGRA_SERIALIZE(data_t); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 8705926a41..95fca2ccc0 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; @@ -2090,19 +2230,13 @@ auto iterative_build_graph(raft::resources const& res, RAFT_LOG_DEBUG("# initial graph size = %lu", (uint64_t)initial_graph_size); // Allocate memory for search results. - constexpr uint64_t max_chunk_size = 8192; + constexpr uint64_t max_chunk_size = helpers::kIterativeBuildChunkSize; // +1 because the search may return the query node itself as a neighbor; // this is consistent with the per-iteration curr_topk = next_graph_degree + 1 auto topk = intermediate_degree + 1; 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/cagra_helpers.cpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp index fa31bccbc3..6c391e5a0b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp @@ -6,10 +6,12 @@ #include "cagra_helpers.hpp" #include +#include #include #include #include #include +#include #include namespace cuvs::neighbors::cagra::helpers { @@ -35,7 +37,8 @@ std::tuple optimize_workspace_size(size_t n_rows size_t graph_degree, size_t intermediate_degree, size_t index_size, - bool mst_optimize) + bool mst_optimize, + bool device_resident_graphs) { RAFT_EXPECTS(graph_degree > 0, "graph_degree must be greater than 0"); RAFT_EXPECTS(intermediate_degree >= graph_degree, @@ -59,17 +62,25 @@ std::tuple optimize_workspace_size(size_t n_rows // Prune stage memory // We neglect 8 bytes (both on host and device) for stats - size_t prune_dev_fixed = batch_size * intermediate_degree; // detour count (uint8_t) - prune_dev_fixed += batch_size * sizeof(uint32_t); // d_num_detour_edges - prune_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) - - size_t prune_dev = n_rows * intermediate_degree * index_size; // d_input_graph + size_t prune_dev_fixed = batch_size * intermediate_degree; // detour count (uint8_t) + prune_dev_fixed += batch_size * sizeof(uint32_t); // d_num_detour_edges + + // Buffers that only exist to stage host-resident graphs to the device. When the caller already + // owns device graphs, batch_load_iterator passes them through and the kernels read and write + // them in place. + size_t prune_dev = 0; + if (!device_resident_graphs) { + prune_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) + prune_dev += n_rows * intermediate_degree * index_size; // d_input_graph + } prune_dev += prune_dev_fixed; // Reverse graph stage memory size_t rev_dev = n_rows * graph_degree * index_size; // d_rev_graph rev_dev += n_rows * sizeof(uint32_t); // d_rev_graph_count - rev_dev += n_rows * index_size; // d_dest_nodes + if (!device_resident_graphs) { + rev_dev += n_rows * index_size; // d_dest_nodes + } // Memory for merging graphs (host only optional) size_t combine_host_fixed = graph_degree * sizeof(uint32_t); // histogram @@ -77,7 +88,10 @@ std::tuple optimize_workspace_size(size_t n_rows combine_host += combine_host_fixed; // additional memory for combine stage on device (3 batches) - size_t combine_dev_fixed = 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) + size_t combine_dev_fixed = 0; + if (!device_resident_graphs) { + combine_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) + } if (mst_optimize) { combine_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_mst_graph(2*batch) combine_dev_fixed += 2 * batch_size * sizeof(uint32_t); // d_mst_graph_num_edges(2*batch) @@ -101,6 +115,128 @@ std::tuple optimize_workspace_size(size_t n_rows return std::make_tuple(total_host, total_dev, total_host_fixed, total_dev_fixed); } +size_t vpq_dataset_size(raft::matrix_extent dataset, + cuvs::neighbors::vpq_params params, + size_t codebook_element_size) +{ + const size_t n_rows = dataset.extent(0); + const size_t dim = dataset.extent(1); + + // Mirror detail::fill_missing_params_heuristics for the fields that affect the footprint. + const size_t pq_bits = params.pq_bits == 0 ? 8 : params.pq_bits; + const size_t pq_dim = + params.pq_dim == 0 ? raft::div_rounding_up_safe(dim, size_t{4}) : params.pq_dim; + const size_t vq_n_centers = + params.vq_n_centers == 0 + ? raft::round_up_safe(static_cast(std::sqrt(static_cast(n_rows))), 8) + : params.vq_n_centers; + + const size_t pq_len = raft::div_rounding_up_safe(dim, pq_dim); + const size_t pq_n_centers = size_t{1} << pq_bits; + + // Every encoded row starts with its inlined VQ label (vpq_build always requests inline labels) + // and is followed by the bit-packed PQ codes. + using label_type = uint32_t; + constexpr size_t kLabelBits = 8 * sizeof(label_type); + const size_t encoded_row_length = + sizeof(label_type) * (1 + raft::div_rounding_up_safe(pq_dim * pq_bits, kLabelBits)); + + return vq_n_centers * dim * codebook_element_size // vq_code_book + + pq_n_centers * pq_len * codebook_element_size // pq_code_book + + n_rows * encoded_row_length; // encoded rows +} + +namespace { + +/** + * Device memory held by a single CAGRA search plan, in bytes. + * + * Mirrors `search_plan_impl_base` (algorithm selection), `search_plan_impl::adjust_search_params` + * and `search_plan_impl::calc_hashmap_params` together with the per-algorithm buffers of + * detail/cagra/search_single_cta.cuh and detail/cagra/search_multi_cta.cuh. + */ +size_t search_plan_mem_usage(cuvs::neighbors::cagra::search_params params, + size_t max_queries, + size_t itopk_size, + size_t graph_degree, + size_t dataset_size, + size_t index_size) +{ + // The iterative build always searches with max_queries = kIterativeBuildChunkSize, which is far + // above the occupancy threshold of the AUTO heuristic, so the algorithm is decided by itopk + // alone. + const bool multi_cta = + params.algo == cuvs::neighbors::cagra::search_algo::MULTI_CTA || + (params.algo == cuvs::neighbors::cagra::search_algo::AUTO && itopk_size > 512); + + constexpr size_t kMultiCtaItopkSize = 32; + + const size_t search_width = std::max(1, params.search_width); + + size_t max_iterations = params.max_iterations; + if (params.max_iterations == 0) { + max_iterations = multi_cta ? kMultiCtaItopkSize : itopk_size / search_width; + size_t reachable_nodes = 1; + while (reachable_nodes < dataset_size) { + reachable_nodes *= std::max(2, graph_degree / 2); + max_iterations += 1; + } + } + if (params.max_iterations < params.min_iterations) { max_iterations = params.min_iterations; } + max_iterations = std::max(max_iterations, params.max_iterations); + + // The internal topk is rounded up to a multiple of 32. + if (itopk_size % 32 != 0) { itopk_size += 32 - (itopk_size % 32); } + + // Smallest hash table that keeps the expected number of entries under the maximum fill rate. + auto hash_bitlen = [&](size_t min_bitlen, size_t expected_nodes) { + size_t bitlen = std::max(min_bitlen, params.hashmap_min_bitlen); + while (static_cast(expected_nodes) > + static_cast(size_t{1} << bitlen) * params.hashmap_max_fill_rate) { + bitlen += 1; + } + return bitlen; + }; + + // num_executed_iterations, allocated for every non-persistent plan. + size_t dev = max_queries * sizeof(uint32_t); + + if (!multi_cta) { + // AUTO and SMALL hash modes keep the visited-node table in shared memory, so nothing is + // allocated globally. AUTO falls back to a global table once the small one would exceed 8K + // entries. + if (params.hashmap_mode != cuvs::neighbors::cagra::hash_mode::HASH) { + const size_t small_bitlen = hash_bitlen(8, itopk_size + search_width * graph_degree); + if (small_bitlen <= 13) { return dev; } + } + const size_t bitlen = + hash_bitlen(11, itopk_size + search_width * graph_degree * max_iterations); + return dev + index_size * max_queries * (size_t{1} << bitlen); // hashmap + } + + // Multi-CTA keeps the per-CTA visited table in shared memory but shares the traversed-node + // table across the CTAs of a query, so that one is always global. + const size_t num_cta_per_query = + std::max(search_width, raft::div_rounding_up_safe(itopk_size, kMultiCtaItopkSize)); + const size_t num_intermediate = num_cta_per_query * kMultiCtaItopkSize; + const size_t bitlen = + hash_bitlen(11, num_cta_per_query * std::max(kMultiCtaItopkSize, max_iterations)); + + dev += index_size * max_queries * (size_t{1} << bitlen); // hashmap + // intermediate_indices / intermediate_distances + dev += num_intermediate * max_queries * (index_size + sizeof(float)); + // topk_workspace: one state byte per 8 candidates per thread of a 1024-thread block + // (_cuann_find_topk_bufferSize). + constexpr size_t kTopkThreads = 1024; + constexpr size_t kTopkStateBits = 8; + dev += raft::div_rounding_up_safe( + raft::div_rounding_up_safe(num_intermediate, kTopkThreads), kTopkStateBits) * + kTopkThreads * max_queries; + return dev; +} + +} // namespace + // All sizes are in bytes inline std::pair ivf_pq_build_mem_usage( raft::resources const& res, @@ -208,10 +344,97 @@ inline std::pair nn_descent_build_mem_usage(raft::resources cons return std::make_pair(total_host, total_dev); } -std::pair cagra_build_mem_usage(raft::resources const& res, - raft::matrix_extent dataset, - cudaDataType_t dtype, - cuvs::neighbors::cagra::index_params cparams) +// All sizes are in bytes +inline std::pair iterative_build_mem_usage( + raft::matrix_extent dataset, + cudaDataType_t dtype, + cuvs::neighbors::graph_build_params::iterative_search_params params, + size_t graph_degree, + size_t intermediate_graph_degree, + bool guarantee_connectivity, + std::optional compression) +{ + // Mirrors detail::iterative_build_graph and detail::search_and_optimize. The build grows the + // graph by repeatedly searching the graph it has so far and optimizing the result; every + // allocation peaks on the final iteration, where the query set, the kNN graph and the output + // graph all span the whole dataset. + constexpr size_t kIndexSize = sizeof(uint32_t); // IdxT + + const size_t n_rows = dataset.extent(0); + const size_t dim = dataset.extent(1); + const size_t chunk = kIterativeBuildChunkSize; + const size_t dtype_size = cuda_data_type_size(dtype); + // The search may return the query node itself, hence the extra column. + const size_t topk = intermediate_graph_degree + 1; + + // The dataset stays resident on the device for the whole build: either VPQ-compressed, or padded + // to CAGRA's row alignment. + size_t dataset_dev; + size_t query_scratch; + if (compression.has_value()) { + dataset_dev = vpq_dataset_size(dataset, compression.value()); + // Queries are reconstructed from the codes one chunk at a time rather than materialized for + // the whole dataset. + query_scratch = chunk * dim * dtype_size; + } else { + const size_t stride = + cuvs::neighbors::cagra_required_row_width(static_cast(dim), dtype_size); + dataset_dev = n_rows * stride * dtype_size; + // Padded rows are depadded into a per-chunk scratch buffer before being used as queries. + query_scratch = stride == dim ? 0 : chunk * dim * dtype_size; + } + + // Search results for one chunk, live for the whole loop. + size_t results_dev = chunk * topk * kIndexSize; // dev_neighbors + results_dev += chunk * topk * sizeof(float); // dev_distances + + // The graph produced by the previous iteration is still alive while the current search fills the + // kNN graph, and the kNN graph is still alive while optimize writes the output graph. So one + // full-size graph and one full-size kNN graph always coexist. search_and_optimize releases the + // previous graph before allocating the output graph, so a third full-size buffer never appears. + const size_t graph_dev = n_rows * graph_degree * kIndexSize; // dev_graph / dev_output_graph + const size_t knn_dev = n_rows * topk * kIndexSize; // dev_knn_graph + + // The final iteration searches a graph_degree graph, requests topk neighbors and derives its + // internal topk from that. + cuvs::neighbors::cagra::search_params search_params = params; + search_params.max_queries = chunk; + const size_t search_dev = + search_plan_mem_usage(search_params, chunk, topk + 32, graph_degree, n_rows, kIndexSize); + + auto [host_workspace_size, gpu_workspace_size, host_ws_fixed, gpu_ws_fixed] = + optimize_workspace_size(n_rows, + graph_degree, + std::max(topk, graph_degree), + kIndexSize, + guarantee_connectivity, + /* device_resident_graphs = */ true); + + // Searching and optimizing run sequentially within an iteration, so the search plan and the + // optimize workspace never coexist and are combined with max() rather than summed. The query + // scratch is not part of that: search_and_optimize holds it at function scope, so it is still + // alive while optimize runs. + // + // Two transients are left out. The dataset copy made by make_device_padded_dataset briefly + // coexists with its source, and cagra::search re-pads a query chunk when its rows are not + // CAGRA-aligned; both are caller-owned or chunk-sized, and counting them would inflate the + // estimate enough to push callers to an out-of-core build unnecessarily. + size_t total_dev = dataset_dev + results_dev + graph_dev + knn_dev + query_scratch + + std::max(search_dev, gpu_workspace_size); + + // On the host the optimize workspace of the last iteration and the returned graph are also + // sequential: the workspace is released before the device graph is copied back. + size_t total_host = std::max(host_workspace_size, n_rows * graph_degree * kIndexSize); + + return std::make_pair(total_host, total_dev); +} + +std::pair cagra_build_mem_usage( + raft::resources const& res, + raft::matrix_extent dataset, + cudaDataType_t dtype, + cuvs::neighbors::cagra::index_params cparams, + std::optional compression) { using namespace cuvs::neighbors; @@ -237,9 +460,21 @@ std::pair cagra_build_mem_usage(raft::resources const& res, cparams.graph_degree, cparams.intermediate_graph_degree, cparams.guarantee_connectivity); + } else if (std::holds_alternative( + cparams.graph_build_params)) { + RAFT_LOG_INFO("Considering CAGRA in memory build with iterative CAGRA search"); + std::tie(total_host, total_dev) = iterative_build_mem_usage( + dataset, + dtype, + std::get(cparams.graph_build_params), + cparams.graph_degree, + cparams.intermediate_graph_degree, + cparams.guarantee_connectivity, + compression); } else { - // iterative build - // TODO(tfeher): proper estimate + // No graph build algorithm selected yet (std::monostate) or an out-of-core (ACE) build, whose + // requirements are modelled by check_ace_memory_requirements instead. Fall back to the size of + // the dataset plus the graphs. total_host = dataset.extent(0) * dataset.extent(1) * cuda_data_type_size(dtype) + dataset.extent(0) * (cparams.graph_degree + cparams.intermediate_graph_degree) * sizeof(uint32_t); diff --git a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp index ee78930970..1e12fec983 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp @@ -11,4 +11,7 @@ namespace cuvs::neighbors::cagra::helpers { /** The batch size for the CAGRA optimize stage. */ constexpr static size_t kOptimizeBatchSize = 256 * 1024; +/** The query chunk size (search batch size) of the iterative CAGRA graph build. */ +constexpr static size_t kIterativeBuildChunkSize = 8192; + } // namespace cuvs::neighbors::cagra::helpers diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f2e0c4f07b..add4bb2532 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -60,6 +60,8 @@ constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::ser return kind::host_padded; } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { return kind::host_standard; + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + return kind::device_vpq_f16; } else { static_assert(sizeof(DatasetViewT) == 0, "serialized_dataset_kind_for_view: unsupported dataset view type"); @@ -69,7 +71,7 @@ constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::ser constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw) { using kind = cuvs::neighbors::cagra::serialized_dataset_kind; - return raw <= static_cast(kind::host_standard); + return raw <= static_cast(kind::device_vpq_f16); } /** @@ -123,9 +125,14 @@ void serialize(raft::resources const& res, RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v) { neighbors::detail::serialize_cagra_dense_dataset(res, os, index_.dataset()); + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + // The payload describes its own codebook type, which is `half` here regardless of T: the + // dtype prefix written above is the type of the queries this index answers, not of its rows. + // `dset()` is safe to call because a view over no rows left include_dataset false above. + neighbors::detail::serialize_vpq_dataset(res, os, index_.dataset().dset()); } else { - // Future dataset types (e.g. VPQ) require a new branch here and a corresponding - // deserialize overload. Use static_assert to catch unsupported types at compile time. + // A further dataset type requires a new branch here and a corresponding deserialize branch. + // Use static_assert to catch unsupported types at compile time. static_assert( sizeof(DatasetViewT) == 0, "serialize: dataset serialization is not yet implemented for this DatasetViewT"); @@ -401,7 +408,16 @@ void deserialize( std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - cuvs::neighbors::detail::skip_dense_dataset(res, is); + // Dropping the rows leaves a searchable index for a dense view, whose dataset can be + // reattached from the caller's own copy, but not for a VPQ one: the compressed rows exist + // nowhere else. Refuse rather than hand back an index that cannot answer a query, and skip + // the dense payload only when it is in fact dense. + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL( + "cagra::deserialize: a VPQ index cannot be loaded without its dataset; pass out_dataset"); + } else { + cuvs::neighbors::detail::skip_dense_dataset(res, is); + } } else { auto const expected_kind = serialized_dataset_kind_for_view(); RAFT_EXPECTS( @@ -419,6 +435,8 @@ void deserialize( } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_host_standard_dataset(res, is); + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + dataset_owner = cuvs::neighbors::detail::deserialize_vpq_dataset(res, is); } else { static_assert(sizeof(DatasetViewT) == 0, "deserialize: dataset deserialization is not implemented for this view"); 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..377494946f 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 ) @@ -214,6 +214,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST PATH neighbors/ann_cagra/test_optimize_uint32_t.cu neighbors/ann_cagra/test_batch_load_iterator.cu + neighbors/ann_cagra/test_mem_usage.cu GPUS 1 PERCENT 100 ) @@ -415,6 +416,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..5ac32b72b4 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -0,0 +1,389 @@ +/* + * 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, that such an index survives a trip through a file with its rows, and that + * the constraints above are rejected rather than accepted and quietly ignored. Fidelity of the + * dataset payload 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; +} + +constexpr int64_t kSearchK = 10; + +/** Neighbour ids for `queries`, row-major [n_queries, kSearchK]. */ +template +auto neighbor_ids(const raft::resources& res, + const IndexT& idx, + raft::device_matrix_view queries) -> std::vector +{ + const auto n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(res, n_queries, kSearchK); + auto distances = raft::make_device_matrix(res, n_queries, kSearchK); + + search_params params; + params.itopk_size = 64; + search(res, params, idx, queries, neighbors.view(), distances.view()); + + std::vector ids(static_cast(n_queries * kSearchK)); + raft::copy(ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return ids; +} + +/** + * 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 +{ + auto ids = neighbor_ids(res, idx, queries); + const auto n_queries = queries.extent(0); + int64_t hits = 0; + for (int64_t q = 0; q < n_queries; q++) { + if (ids[q * kSearchK] == 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 + })); + +/** + * An index over compressed rows is serialized with those rows, so that a loaded index can be + * searched without the dense dataset it came from and without retraining the codebooks. The + * ownership split is the usual one: the file yields an owning dataset, the index only views it. + */ +class CagraQSerializeTest : public CagraQCompressedTestBase { + protected: + void SetUp() override { make_dataset(n_rows, dim); } + + static constexpr int64_t n_rows = 2000; + static constexpr int64_t dim = 128; + static constexpr uint32_t pq_dim = 32; // pq_len 4 +}; + +TEST_F(CagraQSerializeTest, RoundTripsThroughAFileWithItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + auto before = neighbor_ids(res_, idx, queries(500)); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + vpq_f16_index restored{res_}; + std::unique_ptr owner; + cagra::deserialize(res_, stored, &restored, &owner); + + ASSERT_NE(owner, nullptr); + EXPECT_EQ(owner->n_rows(), compressed.n_rows()); + EXPECT_EQ(owner->dim(), compressed.dim()); + EXPECT_EQ(owner->pq_len(), compressed.pq_len()); + EXPECT_EQ(owner->pq_bits(), compressed.pq_bits()); + EXPECT_EQ(owner->vq_n_centers(), compressed.vq_n_centers()); + EXPECT_EQ(owner->encoded_row_length(), compressed.encoded_row_length()); + + ASSERT_EQ(restored.size(), idx.size()); + ASSERT_EQ(restored.dim(), idx.dim()); + ASSERT_EQ(restored.graph_degree(), idx.graph_degree()); + EXPECT_EQ(restored.metric(), idx.metric()); + + // Same graph over the same rows, so the results are identical rather than merely comparable. + auto after = neighbor_ids(res_, restored, queries(500)); + ASSERT_EQ(after.size(), before.size()); + size_t mismatches = 0; + for (size_t i = 0; i < before.size(); i++) { + mismatches += static_cast(after[i] != before[i]); + } + EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; +} + +TEST_F(CagraQSerializeTest, RefusesToLoadWithoutItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // Dropping the rows on load is fine for a dense index, whose caller can attach its own copy, but + // it would leave a VPQ index unsearchable with no way back: the rows exist nowhere else. + vpq_f16_index restored{res_}; + EXPECT_THROW(cagra::deserialize(res_, stored, &restored, nullptr), raft::exception); +} + +TEST_F(CagraQSerializeTest, SerializesTheGraphAloneWhenAsked) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(res_, stored, idx, /* include_dataset */ false); + + vpq_f16_index restored{res_}; + std::unique_ptr owner; + cagra::deserialize(res_, stored, &restored, &owner); + + // Nothing to own, and a graph that only update_dataset() can make searchable again. + EXPECT_EQ(owner, nullptr); + EXPECT_EQ(restored.size(), idx.size()); + EXPECT_EQ(restored.graph_degree(), idx.graph_degree()); +} + +TEST_F(CagraQSerializeTest, RejectsLoadingACompressedIndexAsDense) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // The dtype prefix says float either way, so it is the recorded dataset kind that has to stop + // the dense reader from interpreting VPQ codes as rows of floats. + device_padded_index dense{res_}; + std::unique_ptr> dense_owner; + EXPECT_THROW(cagra::deserialize(res_, stored, &dense, &dense_owner), raft::exception); +} + +/** 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/neighbors/ann_cagra/test_mem_usage.cu b/cpp/tests/neighbors/ann_cagra/test_mem_usage.cu new file mode 100644 index 0000000000..0d7223d553 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_mem_usage.cu @@ -0,0 +1,272 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * Memory-usage estimators for the CAGRA build. + * + * These predict how much host and device memory a build needs without running it, so that callers + * such as hnsw::build and the ACE partition heuristic can choose between an in-memory and an + * out-of-core build. An estimate cannot be checked against itself, so each test here compares one + * against something derived independently: the VPQ footprint against a dataset that was actually + * compressed, and the iterative estimate against the buffers that provably coexist during the build + * and against the peak that raft::memory_stats_resources measures while a real build runs. + */ + +#include + +#include "../cagra_padded_build_helpers.cuh" + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +namespace { + +using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +/** Size of the graph index type the CAGRA build uses internally. */ +constexpr size_t kIndexSize = sizeof(uint32_t); + +auto extents_of(int64_t n_rows, int64_t dim) { return raft::make_extents(n_rows, dim); } + +/** Device bytes actually owned by a compressed dataset: both codebooks plus the encoded rows. */ +auto compressed_bytes(const vpq_dataset_t& v) -> size_t +{ + return static_cast(v.vq_code_book.extent(0)) * v.vq_code_book.extent(1) * sizeof(half) + + static_cast(v.pq_code_book.extent(0)) * v.pq_code_book.extent(1) * sizeof(half) + + static_cast(v.data.extent(0)) * v.data.extent(1); +} + +auto make_clustered(const raft::resources& res, int64_t n_rows, int64_t dim) + -> raft::device_matrix +{ + auto dataset = 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); + return dataset; +} + +auto iterative_params(size_t graph_degree, size_t intermediate_graph_degree) -> index_params +{ + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = graph_degree; + params.intermediate_graph_degree = intermediate_graph_degree; + params.graph_build_params = graph_build_params::iterative_search_params(); + return params; +} + +auto device_estimate(const raft::resources& res, + int64_t n_rows, + int64_t dim, + const index_params& params, + std::optional compression = std::nullopt) -> size_t +{ + return helpers::cagra_build_mem_usage( + res, extents_of(n_rows, dim), CUDA_R_32F, params, compression) + .second; +} + +} // namespace + +// --------------------------------------------------------------------------- +// vpq_dataset_size: predicted footprint vs a dataset that was really compressed +// --------------------------------------------------------------------------- + +TEST(CagraMemUsage, VpqDatasetSizeMatchesCompressedDataset) +{ + raft::resources res; + constexpr int64_t n_rows = 512; + constexpr int64_t dim = 32; + + cuvs::neighbors::vpq_params params; + params.pq_bits = 8; + params.pq_dim = 8; // dim must be divisible by pq_dim + params.vq_n_centers = 32; + params.kmeans_n_iters = 2; + + auto dataset = make_clustered(res, n_rows, dim); + auto compressed = cuvs::preprocessing::quantize::pq::make_vpq_dataset( + res, params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(res); + + EXPECT_EQ(helpers::vpq_dataset_size(extents_of(n_rows, dim), params), + compressed_bytes(compressed)); +} + +TEST(CagraMemUsage, VpqDatasetSizeResolvesUnsetParams) +{ + raft::resources res; + constexpr int64_t n_rows = 1024; + constexpr int64_t dim = 32; + + // pq_dim and vq_n_centers left at 0, so both the estimator and the build have to derive them. + cuvs::neighbors::vpq_params params; + params.pq_bits = 8; + params.kmeans_n_iters = 2; + + auto dataset = make_clustered(res, n_rows, dim); + auto compressed = cuvs::preprocessing::quantize::pq::make_vpq_dataset( + res, params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(res); + + EXPECT_EQ(helpers::vpq_dataset_size(extents_of(n_rows, dim), params), + compressed_bytes(compressed)); +} + +// --------------------------------------------------------------------------- +// optimize_workspace_size: device-resident graphs need no staging buffers +// --------------------------------------------------------------------------- + +TEST(CagraMemUsage, OptimizeWorkspaceSkipsStagingForDeviceResidentGraphs) +{ + constexpr size_t n_rows = 1000000; + constexpr size_t graph_degree = 32; + constexpr size_t intermediate_degree = 256; + + auto [host_staged, dev_staged, host_fixed_staged, dev_fixed_staged] = + helpers::optimize_workspace_size( + n_rows, graph_degree, intermediate_degree, kIndexSize, false, false); + auto [host_resident, dev_resident, host_fixed_resident, dev_fixed_resident] = + helpers::optimize_workspace_size( + n_rows, graph_degree, intermediate_degree, kIndexSize, false, true); + + // The flag only removes device-side staging buffers, so the host estimate must not move. + EXPECT_EQ(host_resident, host_staged); + EXPECT_EQ(host_fixed_resident, host_fixed_staged); + + // At this shape the prune stage dominates the device total, and the d_input_graph staging copy + // it no longer needs is the largest single term the flag drops. + EXPECT_LT(dev_resident, dev_staged); + EXPECT_LE(dev_fixed_resident, dev_fixed_staged); +} + +TEST(CagraMemUsage, OptimizeWorkspaceDefaultsToStagedGraphs) +{ + constexpr size_t n_rows = 4096; + EXPECT_EQ(helpers::optimize_workspace_size(n_rows, 32, 64, kIndexSize), + helpers::optimize_workspace_size(n_rows, 32, 64, kIndexSize, false, false)); +} + +// --------------------------------------------------------------------------- +// cagra_build_mem_usage: the iterative graph build +// --------------------------------------------------------------------------- + +TEST(CagraMemUsage, IterativeEstimateCoversCoexistingBuffers) +{ + raft::resources res; + constexpr int64_t n_rows = 1000000; + constexpr int64_t dim = 128; // 16-byte aligned for float, so never padded + constexpr size_t graph_degree = 64; + constexpr size_t intermediate_degree = 128; + + const size_t dataset_bytes = static_cast(n_rows) * dim * sizeof(float); + const size_t graph_bytes = static_cast(n_rows) * graph_degree * kIndexSize; + const size_t knn_bytes = static_cast(n_rows) * (intermediate_degree + 1) * kIndexSize; + + const auto estimated = + device_estimate(res, n_rows, dim, iterative_params(graph_degree, intermediate_degree)); + + // The resident dataset, the graph carried over from the previous iteration and the kNN graph the + // final search fills are all live at the same moment, so the estimate cannot be below their sum. + EXPECT_GE(estimated, dataset_bytes + graph_bytes + knn_bytes); +} + +TEST(CagraMemUsage, IterativeEstimateGrowsWithDatasetAndDegree) +{ + raft::resources res; + constexpr int64_t dim = 128; + const auto params = iterative_params(64, 128); + + EXPECT_GT(device_estimate(res, 2000000, dim, params), + device_estimate(res, 1000000, dim, params)); + EXPECT_GT(device_estimate(res, 1000000, 256, params), + device_estimate(res, 1000000, dim, params)); + EXPECT_GT(device_estimate(res, 1000000, dim, iterative_params(128, 256)), + device_estimate(res, 1000000, dim, params)); +} + +TEST(CagraMemUsage, IterativeEstimateIsSmallerForCompressedDataset) +{ + raft::resources res; + constexpr int64_t n_rows = 1000000; + constexpr int64_t dim = 128; + const auto params = iterative_params(64, 128); + + cuvs::neighbors::vpq_params compression; + compression.pq_bits = 8; + compression.pq_dim = 32; + compression.vq_n_centers = 1024; + + // CAGRA-Q keeps codes instead of dense rows, so the resident dataset term shrinks by an order of + // magnitude even though the per-chunk reconstruction scratch is new. + EXPECT_LT(device_estimate(res, n_rows, dim, params, compression), + device_estimate(res, n_rows, dim, params)); +} + +TEST(CagraMemUsage, IterativeEstimateAcceptsEqualDegrees) +{ + raft::resources res; + // graph_degree == intermediate_graph_degree once broke the iterative build (issue #1818). The + // estimator has to accept it rather than trip the degree check inside optimize_workspace_size. + EXPECT_GT(device_estimate(res, 100000, 64, iterative_params(16, 16)), 0U); +} + +TEST(CagraMemUsage, IterativeEstimateBoundsMeasuredPeak) +{ + raft::resources res; + constexpr int64_t n_rows = 20000; + constexpr int64_t dim = 64; + const auto params = iterative_params(32, 64); + + const auto estimated = device_estimate(res, n_rows, dim, params); + + size_t measured = 0; + { + // The tracking handle replaces the global device resource, so it has to outlive everything + // allocated below; declaration order here gives it exactly that. + raft::memory_stats_resources tracked{res}; + auto dataset = make_clustered(tracked, n_rows, dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra padded( + tracked, raft::make_const_mdspan(dataset.view())); + auto built = cagra::build(tracked, params, padded.view); + raft::resource::sync_stream(tracked); + ASSERT_GT(built.size(), 0); + + const auto peak = tracked.get_bytes_peak(); + measured = peak.device_global + peak.device_workspace + peak.device_large_workspace + + peak.device_managed; + } + + // One-sided on purpose: hnsw::build uses this estimate to decide whether an in-memory build + // fits, so under-predicting is the failure that matters. No upper bound is asserted because the + // estimate also carries the whole workspace allowance of the resources handle. + EXPECT_GE(estimated, measured); +} + +} // 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