diff --git a/cpp/src/neighbors/detail/cagra/search_plan.cuh b/cpp/src/neighbors/detail/cagra/search_plan.cuh index 8a9d79e177..33376ae43c 100644 --- a/cpp/src/neighbors/detail/cagra/search_plan.cuh +++ b/cpp/src/neighbors/detail/cagra/search_plan.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -265,12 +265,18 @@ struct search_plan_impl : public search_plan_impl_base { // table that each CTA has in the shared memory. This hash table is not // shared among CTAs. This hash table is reset and restored in each iteration. // - const uint32_t max_visited_nodes = mc_itopk_size + (graph_degree * 2); - small_hash_bitlen = 8; // 256 - while (max_visited_nodes > hashmap::get_size(small_hash_bitlen) * max_fill_rate) { + const uint32_t max_visited_nodes = mc_itopk_size + (graph_degree * 2); + constexpr size_t max_small_hash_bitlen = 14; // 16K + small_hash_bitlen = 8; // 256 + // Stop at the supported maximum rather than growing without bound: hashmap::get_size() + // is 1U << bitlen, which is undefined (and wraps to a small value) once bitlen reaches + // 32, so an unbounded loop would never satisfy its exit condition again. + while (small_hash_bitlen <= max_small_hash_bitlen && + max_visited_nodes > hashmap::get_size(small_hash_bitlen) * max_fill_rate) { small_hash_bitlen += 1; } - RAFT_EXPECTS(small_hash_bitlen <= 14, "small_hash_bitlen cannot be largen than 14 (16K)"); + RAFT_EXPECTS(small_hash_bitlen <= max_small_hash_bitlen, + "small_hash_bitlen cannot be largen than 14 (16K)"); // // [traversed_hash_table] // Whether a node has ever been used as the starting point for a traversal @@ -279,13 +285,15 @@ struct search_plan_impl : public search_plan_impl_base { // const auto max_traversed_nodes = mc_num_cta_per_query * max((size_t)mc_itopk_size, max_iterations); - unsigned min_bitlen = 11; // 2K + unsigned min_bitlen = 11; // 2K + constexpr int64_t max_hash_bitlen_mc = 25; // 32M if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; } hash_bitlen = min_bitlen; - while (max_traversed_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) { + while (hash_bitlen <= max_hash_bitlen_mc && + max_traversed_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) { hash_bitlen += 1; } - RAFT_EXPECTS(hash_bitlen <= 25, "hash_bitlen cannot be largen than 25 (32M)"); + RAFT_EXPECTS(hash_bitlen <= max_hash_bitlen_mc, "hash_bitlen cannot be largen than 25 (32M)"); } else { while (hashmap_mode == hash_mode::AUTO || hashmap_mode == hash_mode::SMALL) { // @@ -300,7 +308,8 @@ struct search_plan_impl : public search_plan_impl_base { unsigned max_bitlen = 13; // 8K if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; } hash_bitlen = min_bitlen; - while (max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) { + while (hash_bitlen <= max_bitlen && + max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) { hash_bitlen += 1; } if (hash_bitlen > max_bitlen) { @@ -337,13 +346,15 @@ struct search_plan_impl : public search_plan_impl_base { // maximum fill rate of the hash table. // uint32_t max_visited_nodes = itopk_size + (search_width * graph_degree * max_iterations); - unsigned min_bitlen = 11; // 2K + unsigned min_bitlen = 11; // 2K + constexpr int64_t max_hash_bitlen = 20; // 1M if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; } hash_bitlen = min_bitlen; - while (max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) { + while (hash_bitlen <= max_hash_bitlen && + max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) { hash_bitlen += 1; } - RAFT_EXPECTS(hash_bitlen <= 20, + RAFT_EXPECTS(hash_bitlen <= max_hash_bitlen, "hash_bitlen cannot be largen than 20 (1M). You can decrease itopk_size, " "search_width or max_iterations to reduce the required hashmap size."); } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index b8f571e4a9..07040d6c1c 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -194,6 +194,7 @@ ConfigureTest( neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu neighbors/ann_cagra/bug_iterative_cagra_build.cu neighbors/ann_cagra/bug_issue_93_reproducer.cu + neighbors/ann_cagra/bug_issue_2523_hashmap_bitlen.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/bug_issue_2523_hashmap_bitlen.cu b/cpp/tests/neighbors/ann_cagra/bug_issue_2523_hashmap_bitlen.cu new file mode 100644 index 0000000000..0d0df08f0f --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/bug_issue_2523_hashmap_bitlen.cu @@ -0,0 +1,135 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include "../cagra_padded_build_helpers.cuh" +#include + +#include +#include +#include +#include + +#include +#include + +namespace cuvs::neighbors::cagra { + +/** + * @brief Regression test for issue #2523: hash table sizing must not loop forever. + * + * The hash tables in search_plan_impl::calc_hashmap_params() are sized by loops that grow + * a bit length until the requested table fits, with the supported maximum checked only + * after the loop. hashmap::get_size() is `1U << bitlen`, which is undefined once bitlen + * reaches 32 and in practice wraps back to a small value, so a request that needs a table + * larger than 2^31 entries made the loop condition permanently true. The bit length grew + * without bound, the post-loop RAFT_EXPECTS was never reached, and the search hung on the + * host instead of failing. + * + * check_params() only caps itopk_size at 1024 for SINGLE_CTA, so MULTI_CTA and MULTI_KERNEL + * are the two algorithms that can reach the sizing loops with an oversized request. + * + * These searches must now raise the existing "hash_bitlen cannot be larger than ..." error + * rather than hanging. A regression reappears as a test timeout, not a failed assertion. + */ +class cagra_hashmap_bitlen_no_hang_test : public ::testing::Test { + public: + using data_type = float; + using index_type = uint32_t; + + protected: + // Large enough that the traversed-node hash table would need more than 2^31 entries, + // which is what previously drove the sizing loop past the wrap point. + constexpr static size_t oversized_itopk = 1'100'000'000; + constexpr static size_t valid_itopk = 64; + + constexpr static int64_t n_dataset = 1000; + constexpr static int64_t n_dim = 32; + constexpr static int64_t n_queries = 4; + constexpr static int64_t k = 10; + + void SetUp() override + { + dataset.emplace(raft::make_device_matrix(res, n_dataset, n_dim)); + queries.emplace(raft::make_device_matrix(res, n_queries, n_dim)); + neighbors.emplace(raft::make_device_matrix(res, n_queries, k)); + distances.emplace(raft::make_device_matrix(res, n_queries, k)); + + raft::random::RngState r(1234ULL); + raft::random::uniform( + res, r, dataset->data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); + raft::random::uniform( + res, r, queries->data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); + + cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + + padded_.emplace(res, raft::make_const_mdspan(dataset->view())); + index_.emplace(cagra::build(res, index_params, padded_->view)); + raft::resource::sync_stream(res); + } + + void TearDown() override + { + index_.reset(); + padded_.reset(); + dataset.reset(); + queries.reset(); + neighbors.reset(); + distances.reset(); + raft::resource::sync_stream(res); + } + + void search_with(cagra::search_algo algo, size_t itopk_size, uint32_t max_iterations) + { + cagra::search_params search_params; + search_params.algo = algo; + search_params.itopk_size = itopk_size; + search_params.search_width = 8; + search_params.max_iterations = max_iterations; + + cagra::search(res, + search_params, + *index_, + raft::make_const_mdspan(queries->view()), + neighbors->view(), + distances->view()); + raft::resource::sync_stream(res); + } + + raft::resources res; + std::optional> padded_{}; + std::optional> index_ = std::nullopt; + std::optional> dataset = std::nullopt; + std::optional> queries = std::nullopt; + std::optional> neighbors = std::nullopt; + std::optional> distances = std::nullopt; +}; + +// MULTI_CTA sizes the shared traversed-node table from itopk_size and search_width. +// This is the path confirmed to hang before the fix. +TEST_F(cagra_hashmap_bitlen_no_hang_test, MultiCtaOversizedItopkThrows) +{ + EXPECT_THROW(search_with(cagra::search_algo::MULTI_CTA, oversized_itopk, 0), raft::exception); +} + +// MULTI_KERNEL reaches the small-hash loop first and then the normal-hash loop. +// max_iterations is pinned so the sizing inputs do not depend on the auto-derived value. +TEST_F(cagra_hashmap_bitlen_no_hang_test, MultiKernelOversizedItopkThrows) +{ + EXPECT_THROW(search_with(cagra::search_algo::MULTI_KERNEL, oversized_itopk, 32), raft::exception); +} + +// Guards against the bound rejecting requests that were previously accepted. +TEST_F(cagra_hashmap_bitlen_no_hang_test, ValidItopkStillSucceeds) +{ + EXPECT_NO_THROW(search_with(cagra::search_algo::MULTI_CTA, valid_itopk, 0)); + EXPECT_NO_THROW(search_with(cagra::search_algo::MULTI_KERNEL, valid_itopk, 0)); + EXPECT_NO_THROW(search_with(cagra::search_algo::SINGLE_CTA, valid_itopk, 0)); +} + +} // namespace cuvs::neighbors::cagra