diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 45eb945fbc..220ffd86ab 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -364,6 +364,7 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_"); nlohmann::json nn_descent_conf = collect_conf_with_prefix(conf, "nn_descent_"); nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_"); + nlohmann::json build_search_conf = collect_conf_with_prefix(conf, "build_search_"); // When graph_build_algo is not specified, leave graph_build_params as monostate so the // CAGRA build uses AUTO selection (NN_DESCENT or IVF_PQ based on dataset/heuristics). @@ -394,6 +395,85 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index } else if constexpr (std::is_same_v) { parse_build_param(nn_descent_conf, arg); + } else if constexpr (std::is_same_v< + U, + cuvs::neighbors::graph_build_params::iterative_search_params>) { + if (build_search_conf.contains("width")) { + arg.search_width = build_search_conf.at("width"); + } + if (build_search_conf.contains("max_iterations")) { + arg.max_iterations = build_search_conf.at("max_iterations"); + } + if (build_search_conf.contains("min_iterations")) { + arg.min_iterations = build_search_conf.at("min_iterations"); + } + if (build_search_conf.contains("itopk")) { arg.itopk_size = build_search_conf.at("itopk"); } + if (build_search_conf.contains("max_queries")) { + arg.max_queries = build_search_conf.at("max_queries"); + } + if (build_search_conf.contains("team_size")) { + arg.team_size = build_search_conf.at("team_size"); + } + if (build_search_conf.contains("thread_block_size")) { + arg.thread_block_size = build_search_conf.at("thread_block_size"); + } + if (build_search_conf.contains("hashmap_min_bitlen")) { + arg.hashmap_min_bitlen = build_search_conf.at("hashmap_min_bitlen"); + } + if (build_search_conf.contains("hashmap_max_fill_rate")) { + arg.hashmap_max_fill_rate = build_search_conf.at("hashmap_max_fill_rate"); + } + if (build_search_conf.contains("num_random_samplings")) { + arg.num_random_samplings = build_search_conf.at("num_random_samplings"); + } + if (build_search_conf.contains("persistent")) { + arg.persistent = build_search_conf.at("persistent"); + } + if (build_search_conf.contains("persistent_lifetime")) { + arg.persistent_lifetime = build_search_conf.at("persistent_lifetime"); + } + if (build_search_conf.contains("persistent_device_usage")) { + arg.persistent_device_usage = build_search_conf.at("persistent_device_usage"); + } + if (build_search_conf.contains("algo")) { + std::string algo = build_search_conf.at("algo"); + if (algo == "single_cta") { + arg.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + } else if (algo == "multi_cta") { + arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + } else if (algo == "multi_kernel") { + arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_KERNEL; + } else if (algo == "auto") { + arg.algo = cuvs::neighbors::cagra::search_algo::AUTO; + } + } + if (build_search_conf.contains("hashmap_mode")) { + std::string mode = build_search_conf.at("hashmap_mode"); + if (mode == "hash") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::HASH; + } else if (mode == "small") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::SMALL; + } else if (mode == "auto") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::AUTO; + } + } + // Precision of the codebook/query in shared memory for the VPQ search used during + // the iterative build. Accepts an integer code (0=F16, 1=E5M2) or a string. + if (build_search_conf.contains("smem_dtype")) { + const auto& sd = build_search_conf.at("smem_dtype"); + if (sd.is_number_integer()) { + arg.smem_dtype = static_cast(sd.get()); + } else { + std::string s = sd.get(); + if (s == "f16" || s == "F16" || s == "fp16" || s == "half") { + arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } else if (s == "e5m2" || s == "E5M2" || s == "fp8") { + arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::E5M2; + } else { + throw std::runtime_error("invalid value for build_search smem_dtype: " + s); + } + } + } } }, params.graph_build_params); diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 43ae7a6235..dc18a6b792 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -34,13 +34,147 @@ #include #include #include +#include #include #include +namespace CUVS_EXPORT cuvs { +namespace neighbors { +namespace cagra { + +/** + * @defgroup cagra_cpp_search_params CAGRA index search parameters + * @{ + */ + +enum class search_algo { + /** For large batch sizes. */ + SINGLE_CTA = 0, + /** For small batch sizes. */ + MULTI_CTA = 1, + MULTI_KERNEL = 2, + AUTO = 100 +}; + +enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; + +enum class internal_dtype { F16 = 0, E5M2 = 1 }; + +struct search_params : cuvs::neighbors::search_params { + /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ + size_t max_queries = 0; + + /** Number of intermediate search results retained during the search. + * + * This is the main knob to adjust trade off between accuracy and search speed. + * Higher values improve the search accuracy. + */ + size_t itopk_size = 64; + + /** Upper limit of search iterations. Auto select when 0.*/ + size_t max_iterations = 0; + + // In the following we list additional search parameters for fine tuning. + // Reasonable default values are automatically chosen. + + /** Which search implementation to use. */ + search_algo algo = search_algo::AUTO; + + /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */ + size_t team_size = 0; + + /** Number of graph nodes to select as the starting point for the search in each iteration. aka + * search width?*/ + size_t search_width = 1; + /** Lower limit of search iterations. */ + size_t min_iterations = 0; + + /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */ + size_t thread_block_size = 0; + /** Hashmap type. Auto selection when AUTO. */ + hash_mode hashmap_mode = hash_mode::AUTO; + /** Lower limit of hashmap bit length. More than 8. */ + size_t hashmap_min_bitlen = 0; + /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/ + float hashmap_max_fill_rate = 0.5; + + /** Number of iterations of initial random seed node selection. 1 or more. */ + uint32_t num_random_samplings = 1; + /** Bit mask used for initial random seed node selection. */ + uint64_t rand_xor_mask = 0x128394; + + /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ + bool persistent = false; + /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ + float persistent_lifetime = 2; + /** + * Set the fraction of maximum grid size used by persistent kernel. + * Value 1.0 means the kernel grid size is maximum possible for the selected device. + * The value must be greater than 0.0 and not greater than 1.0. + * + * One may need to run other kernels alongside this persistent kernel. This parameter can + * be used to reduce the grid size of the persistent kernel to leave a few SMs idle. + * Note: running any other work on GPU alongside with the persistent kernel makes the setup + * fragile. + * - Running another kernel in another thread usually works, but no progress guaranteed + * - Any CUDA allocations block the context (this issue may be obscured by using pools) + * - Memory copies to not-pinned host memory may block the context + * + * Even when we know there are no other kernels working at the same time, setting + * kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. + * If you suspect this is an issue, you can reduce this number to ~0.9 without a significant + * impact on the throughput. + */ + float persistent_device_usage = 1.0; + + /** + * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. + * The value must be equal to or greater than 0.0 and less than 1.0. Default value is + * negative, in which case the filtering rate is automatically calculated when possible. + * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is + * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be + * inferred from the source string. + */ + float filtering_rate = -1.0; + + /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ + * supports FP8. **/ + internal_dtype smem_dtype = internal_dtype::F16; +}; + +/** + * @} + */ + +} // namespace cagra +} // namespace neighbors +} // namespace CUVS_EXPORT cuvs + namespace CUVS_EXPORT cuvs { namespace neighbors { namespace graph_build_params { -using iterative_search_params = cuvs::neighbors::search_params; +/** + * Parameters for the iterative CAGRA graph build algorithm. + * + * Inherits from cagra::search_params so that all search tuning knobs + * (search_width, max_iterations, itopk_size, etc.) are available for + * controlling the search-and-optimize loop during graph construction. + * The defaults are tuned for the build loop (e.g. search_width=1, + * max_iterations=8) and may differ from the regular search defaults. + * + */ +struct iterative_search_params : cuvs::neighbors::cagra::search_params { + iterative_search_params() + { + this->search_width = 1; + this->max_iterations = 8; + // itopk_size controls the search during the *growing* iterations of the build loop. + // 0 (default) means auto-select per iteration (max(graph_degree + 32, 128)); a nonzero + // value overrides it for the growing iterations. The final iteration always uses a fixed + // itopk tied to the output topk, regardless of this value. + this->itopk_size = 0; + } +}; /** Specialized parameters for ACE (Augmented Core Extraction) graph build */ struct ace_params { @@ -311,110 +445,6 @@ struct index_params : cuvs::neighbors::index_params { cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded); }; -/** - * @} - */ - -/** - * @defgroup cagra_cpp_search_params CAGRA index search parameters - * @{ - */ - -enum class search_algo { - /** For large batch sizes. */ - SINGLE_CTA = 0, - /** For small batch sizes. */ - MULTI_CTA = 1, - MULTI_KERNEL = 2, - AUTO = 100 -}; - -enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; - -enum class internal_dtype { F16 = 0, E5M2 = 1 }; - -struct search_params : cuvs::neighbors::search_params { - /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ - size_t max_queries = 0; - - /** Number of intermediate search results retained during the search. - * - * This is the main knob to adjust trade off between accuracy and search speed. - * Higher values improve the search accuracy. - */ - size_t itopk_size = 64; - - /** Upper limit of search iterations. Auto select when 0.*/ - size_t max_iterations = 0; - - // In the following we list additional search parameters for fine tuning. - // Reasonable default values are automatically chosen. - - /** Which search implementation to use. */ - search_algo algo = search_algo::AUTO; - - /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */ - size_t team_size = 0; - - /** Number of graph nodes to select as the starting point for the search in each iteration. aka - * search width?*/ - size_t search_width = 1; - /** Lower limit of search iterations. */ - size_t min_iterations = 0; - - /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */ - size_t thread_block_size = 0; - /** Hashmap type. Auto selection when AUTO. */ - hash_mode hashmap_mode = hash_mode::AUTO; - /** Lower limit of hashmap bit length. More than 8. */ - size_t hashmap_min_bitlen = 0; - /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/ - float hashmap_max_fill_rate = 0.5; - - /** Number of iterations of initial random seed node selection. 1 or more. */ - uint32_t num_random_samplings = 1; - /** Bit mask used for initial random seed node selection. */ - uint64_t rand_xor_mask = 0x128394; - - /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ - bool persistent = false; - /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ - float persistent_lifetime = 2; - /** - * Set the fraction of maximum grid size used by persistent kernel. - * Value 1.0 means the kernel grid size is maximum possible for the selected device. - * The value must be greater than 0.0 and not greater than 1.0. - * - * One may need to run other kernels alongside this persistent kernel. This parameter can - * be used to reduce the grid size of the persistent kernel to leave a few SMs idle. - * Note: running any other work on GPU alongside with the persistent kernel makes the setup - * fragile. - * - Running another kernel in another thread usually works, but no progress guaranteed - * - Any CUDA allocations block the context (this issue may be obscured by using pools) - * - Memory copies to not-pinned host memory may block the context - * - * Even when we know there are no other kernels working at the same time, setting - * kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. - * If you suspect this is an issue, you can reduce this number to ~0.9 without a significant - * impact on the throughput. - */ - float persistent_device_usage = 1.0; - - /** - * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. - * The value must be equal to or greater than 0.0 and less than 1.0. Default value is - * negative, in which case the filtering rate is automatically calculated when possible. - * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is - * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be - * inferred from the source string. - */ - float filtering_rate = -1.0; - - /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ - * supports FP8. **/ - internal_dtype smem_dtype = internal_dtype::F16; -}; - /** * @} */ @@ -914,9 +944,12 @@ using vpq_f32_index = index -using cagra_index_t = index, - uint32_t, - cuvs::neighbors::dataset_view_type_t>; +using cagra_index_t = std::conditional_t< + cuvs::neighbors::is_device_vpq_f16_dataset_view_v, + index>, + index, + uint32_t, + cuvs::neighbors::dataset_view_type_t>>; /** * @} @@ -928,10 +961,11 @@ using cagra_index_t = index>` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_vpq_dataset_view const& dataset) + -> index>; + /** * @brief Build from a device padded dataset view (`float`). * @param[in] res raft resources diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 935938c9b0..eb938cfd26 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -109,6 +109,18 @@ struct vpq_params { * The max number of data points to use per VQ cluster during training. */ uint32_t max_train_points_per_vq_cluster = 1024; + + friend bool operator==(const vpq_params& a, const vpq_params& b) + { + return a.pq_bits == b.pq_bits && a.pq_dim == b.pq_dim && a.vq_n_centers == b.vq_n_centers && + a.kmeans_n_iters == b.kmeans_n_iters && + a.vq_kmeans_trainset_fraction == b.vq_kmeans_trainset_fraction && + a.pq_kmeans_trainset_fraction == b.pq_kmeans_trainset_fraction && + a.pq_kmeans_type == b.pq_kmeans_type && + a.max_train_points_per_pq_code == b.max_train_points_per_pq_code && + a.max_train_points_per_vq_cluster == b.max_train_points_per_vq_cluster; + } + friend bool operator!=(const vpq_params& a, const vpq_params& b) { return !(a == b); } }; /** @} */ // end group cagra_cpp_index_params diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 80e2f2a07e..fc1e29c033 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -296,13 +296,43 @@ template auto build(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> cuvs::neighbors::cagra::cagra_index_t { - using T = cuvs::neighbors::cagra_view_element_type_t; - using IdxT = uint32_t; + using index_type = cuvs::neighbors::cagra::cagra_index_t; + using T = typename index_type::value_type; + using IdxT = uint32_t; // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until attach_dataset(...) supplies a device-padded dataset. if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { - RAFT_FAIL("cagra::build: VPQ-compressed dataset cannot be used for dense graph construction."); + auto effective_params = params; + if (std::holds_alternative(effective_params.graph_build_params)) { + effective_params.graph_build_params = graph_build_params::iterative_search_params{}; + } + + RAFT_EXPECTS(std::holds_alternative( + effective_params.graph_build_params), + "cagra::build: a VPQ dataset requires iterative_search_params graph construction"); + RAFT_EXPECTS(effective_params.metric == cuvs::distance::DistanceType::L2Expanded, + "cagra::build: a VPQ dataset supports only L2Expanded distance"); + RAFT_EXPECTS(dataset.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); + RAFT_EXPECTS(dataset.dset().pq_bits() == 8, + "cagra::build: VPQ dataset requires pq_bits == 8, got %u", + dataset.dset().pq_bits()); + auto const pq_len = dataset.dset().pq_len(); + RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, + "cagra::build: VPQ dataset requires pq_len in {2, 4, 8}, got %u", + pq_len); + + detail::check_graph_degree(effective_params.intermediate_graph_degree, + effective_params.graph_degree, + static_cast(dataset.n_rows())); + auto cagra_graph = detail::iterative_build_graph(res, effective_params, dataset); + + index_type idx(res, effective_params.metric); + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + if (effective_params.attach_dataset_on_build) { + idx.update_device_dataset_same_layout(res, dataset); + } + return idx; } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix( res, params, dataset); diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index acaaa942c1..90d63c3ca9 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -18,6 +18,9 @@ using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view< using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; +#if @emit_vpq_build@ +using inst_device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; +#endif } // namespace namespace cuvs::neighbors::cagra { @@ -55,6 +58,13 @@ CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_padded_view_t, CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_standard_view_t, cuvs::neighbors::cagra::host_standard_index); +#if @emit_vpq_build@ +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD( + inst_device_vpq_view_t, + cuvs::neighbors::cagra:: + index>); +#endif + #undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_build_matrix.json b/cpp/src/neighbors/cagra_build_matrix.json index a7995005c4..9fae2b33f9 100644 --- a/cpp/src/neighbors/cagra_build_matrix.json +++ b/cpp/src/neighbors/cagra_build_matrix.json @@ -2,19 +2,23 @@ "_data": [ { "data_type": "float", - "data_abbrev": "f" + "data_abbrev": "f", + "emit_vpq_build": 1 }, { "data_type": "half", - "data_abbrev": "h" + "data_abbrev": "h", + "emit_vpq_build": 0 }, { "data_type": "int8_t", - "data_abbrev": "i8" + "data_abbrev": "i8", + "emit_vpq_build": 0 }, { "data_type": "uint8_t", - "data_abbrev": "u8" + "data_abbrev": "u8", + "emit_vpq_build": 0 } ], "_index": [ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index c06f9b12e3..1674ac1768 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 { @@ -1967,51 +1965,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. */ @@ -2032,7 +2157,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 @@ -2040,22 +2166,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; @@ -2070,12 +2210,6 @@ auto iterative_build_graph(raft::resources const& res, auto dev_neighbors = raft::make_device_matrix(res, max_chunk_size, topk); auto dev_distances = raft::make_device_matrix(res, max_chunk_size, topk); - std::optional> query_contiguous; - if (static_cast(logical_dim) != dev_dataset.extent(1)) { - query_contiguous.emplace( - raft::make_device_matrix(res, max_chunk_size, logical_dim)); - } - // Determine graph degree and number of search results while increasing // graph size. auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24)); @@ -2083,6 +2217,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); @@ -2103,28 +2247,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; @@ -2139,86 +2289,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_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 165e478337..34a1762e85 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -85,7 +85,6 @@ void search_main_core( search_plan_impl> plan = factory::create( res, params, dataset_desc, queries.extent(1), graph.extent(0), graph.extent(1), topk); - plan->check(topk); RAFT_LOG_DEBUG("Cagra search"); @@ -208,6 +207,7 @@ void search_main(raft::resources const& res, params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; } // Search using a plain (strided) row-major dataset + RAFT_LOG_DEBUG("Searching with strided dataset"); RAFT_EXPECTS(index.metric() != cuvs::distance::DistanceType::CosineExpanded || index.dataset_norms().has_value(), "Dataset norms must be provided for CosineExpanded metric"); @@ -238,6 +238,7 @@ void search_main(raft::resources const& res, RAFT_FAIL("FP32 VPQ dataset support is coming soon"); } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { auto const& vv = index.dataset(); + RAFT_LOG_DEBUG("Searching with VPQ dataset"); if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && raft::getComputeCapability().first < 9) { RAFT_LOG_WARN( 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);