From 7d5f8270be5d04ac873d409483fd248e270f10cf Mon Sep 17 00:00:00 2001 From: Dmitriy Vesnin Date: Fri, 3 Jul 2026 22:59:03 +0300 Subject: [PATCH 1/2] Fix IVF extend list resize --- cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh | 3 +- cpp/src/neighbors/ivf_list.cuh | 19 ++++- cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh | 3 +- .../ann_ivf_flat/test_float_int64_t.cu | 74 ++++++++++++++++++- .../ann_ivf_sq/test_float_int64_t.cu | 62 +++++++++++++++- 5 files changed, 154 insertions(+), 7 deletions(-) diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh index fffe5134ae..fa7c0c36a2 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -286,6 +286,7 @@ void extend(raft::resources const& handle, lists[label], list_device_spec, new_list_sizes[label], + old_list_sizes[label], raft::Pow2::roundUp(old_list_sizes[label])); } } diff --git a/cpp/src/neighbors/ivf_list.cuh b/cpp/src/neighbors/ivf_list.cuh index 24691463ff..7c6f5d7e66 100644 --- a/cpp/src/neighbors/ivf_list.cuh +++ b/cpp/src/neighbors/ivf_list.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -66,13 +66,16 @@ CUVS_EXPORT void resize_list(raft::resources const& res, std::shared_ptr& orig_list, // NOLINT const typename ListT::spec_type& spec, typename ListT::size_type new_used_size, + typename ListT::size_type old_logical_size, typename ListT::size_type old_used_size) { + // old_logical_size is the value stored in list::size. old_used_size is the copy extent from + // the old allocation and may include padding slots for interleaved list layouts. bool skip_resize = false; if (orig_list) { if (new_used_size <= orig_list->indices.extent(0)) { - auto shared_list_size = old_used_size; - if (new_used_size <= old_used_size || + auto shared_list_size = old_logical_size; + if (new_used_size <= old_logical_size || orig_list->size.compare_exchange_strong(shared_list_size, new_used_size)) { // We don't need to resize the list if: // 1. The list exists @@ -104,6 +107,16 @@ CUVS_EXPORT void resize_list(raft::resources const& res, new_list.swap(orig_list); } +template +CUVS_EXPORT void resize_list(raft::resources const& res, + std::shared_ptr& orig_list, // NOLINT + const typename ListT::spec_type& spec, + typename ListT::size_type new_used_size, + typename ListT::size_type old_used_size) +{ + resize_list(res, orig_list, spec, new_used_size, old_used_size, old_used_size); +} + template enable_if_valid_list_t serialize_list(const raft::resources& handle, std::ostream& os, diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh index 649a4f0720..6f9268997c 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -358,6 +358,7 @@ void extend_inplace(raft::resources const& handle, lists[label], list_device_spec, new_list_sizes[label], + old_list_sizes[label], raft::Pow2::roundUp(old_list_sizes[label])); } } diff --git a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu index f425de000c..9e26942bd1 100644 --- a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,6 +7,9 @@ #include "../ann_ivf_flat.cuh" +#include +#include + namespace cuvs::neighbors::ivf_flat { typedef AnnIVFFlatTest AnnIVFFlatTestF_float; @@ -19,4 +22,73 @@ TEST_P(AnnIVFFlatTestF_float, AnnIVFFlat) INSTANTIATE_TEST_CASE_P(AnnIVFFlatTest, AnnIVFFlatTestF_float, ::testing::ValuesIn(inputs)); +TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t base_rows = 100; + constexpr int64_t grow_rows = 20; + constexpr int64_t rows = base_rows + grow_rows; + constexpr int64_t dim = 4; + + std::vector host_data(rows * dim); + for (int64_t row = 0; row < rows; row++) { + for (int64_t col = 0; col < dim; col++) { + host_data[row * dim + col] = static_cast(row + col); + } + } + + auto data = raft::make_device_matrix(handle, rows, dim); + raft::copy(data.data_handle(), host_data.data(), host_data.size(), stream); + + index_params params; + params.n_lists = 1; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.add_data_on_build = false; + params.kmeans_trainset_fraction = 1.0; + params.adaptive_centers = false; + params.conservative_memory_allocation = false; + + auto all_data_view = + raft::make_device_matrix_view(data.data_handle(), rows, dim); + auto empty_index = build(handle, params, all_data_view); + + auto base_data_view = + raft::make_device_matrix_view(data.data_handle(), base_rows, dim); + auto base_index = extend(handle, base_data_view, std::nullopt, empty_index); + raft::resource::sync_stream(handle); + + ASSERT_EQ(base_index.lists()[0]->get_size(), base_rows); + ASSERT_GE(base_index.lists()[0]->indices_capacity(), rows); + + std::vector host_indices(grow_rows); + std::iota(host_indices.begin(), host_indices.end(), base_rows); + auto indices = raft::make_device_vector(handle, grow_rows); + raft::copy(indices.data_handle(), host_indices.data(), host_indices.size(), stream); + + auto grow_data_view = raft::make_device_matrix_view( + data.data_handle() + base_rows * dim, grow_rows, dim); + auto grow_indices_view = + raft::make_device_vector_view(indices.data_handle(), grow_rows); + auto first_grown_index = + extend(handle, + grow_data_view, + std::make_optional>(grow_indices_view), + base_index); + raft::resource::sync_stream(handle); + + ASSERT_EQ(first_grown_index.lists()[0]->get_size(), rows); + + auto second_grown_index = + extend(handle, + grow_data_view, + std::make_optional>(grow_indices_view), + base_index); + raft::resource::sync_stream(handle); + + EXPECT_NE(first_grown_index.lists()[0].get(), second_grown_index.lists()[0].get()); + EXPECT_EQ(second_grown_index.lists()[0]->get_size(), rows); +} + } // namespace cuvs::neighbors::ivf_flat diff --git a/cpp/tests/neighbors/ann_ivf_sq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_sq/test_float_int64_t.cu index 8831ae720a..565b336c03 100644 --- a/cpp/tests/neighbors/ann_ivf_sq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_sq/test_float_int64_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,6 +7,9 @@ #include "../ann_ivf_sq.cuh" +#include +#include + namespace cuvs::neighbors::ivf_sq { typedef AnnIVFSQTest AnnIVFSQTestF_float; @@ -14,4 +17,61 @@ TEST_P(AnnIVFSQTestF_float, AnnIVFSQ) { this->testAll(); } INSTANTIATE_TEST_CASE_P(AnnIVFSQTest, AnnIVFSQTestF_float, ::testing::ValuesIn(inputs)); +TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t base_rows = 100; + constexpr int64_t grow_rows = 20; + constexpr int64_t rows = base_rows + grow_rows; + constexpr int64_t dim = 4; + + std::vector host_data(rows * dim); + for (int64_t row = 0; row < rows; row++) { + for (int64_t col = 0; col < dim; col++) { + host_data[row * dim + col] = static_cast(row + col); + } + } + + auto data = raft::make_device_matrix(handle, rows, dim); + raft::copy(data.data_handle(), host_data.data(), host_data.size(), stream); + + index_params params; + params.n_lists = 1; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.add_data_on_build = false; + params.max_train_points_per_cluster = 256; + params.conservative_memory_allocation = false; + + auto all_data_view = + raft::make_device_matrix_view(data.data_handle(), rows, dim); + auto index = build(handle, params, all_data_view); + + auto base_data_view = + raft::make_device_matrix_view(data.data_handle(), base_rows, dim); + extend(handle, base_data_view, std::nullopt, &index); + raft::resource::sync_stream(handle); + + ASSERT_EQ(index.lists()[0]->get_size(), base_rows); + ASSERT_GE(index.lists()[0]->indices_capacity(), rows); + + std::vector host_indices(grow_rows); + std::iota(host_indices.begin(), host_indices.end(), base_rows); + auto indices = raft::make_device_vector(handle, grow_rows); + raft::copy(indices.data_handle(), host_indices.data(), host_indices.size(), stream); + + auto grow_data_view = raft::make_device_matrix_view( + data.data_handle() + base_rows * dim, grow_rows, dim); + auto grow_indices_view = + raft::make_device_vector_view(indices.data_handle(), grow_rows); + extend(handle, + grow_data_view, + std::make_optional>(grow_indices_view), + &index); + raft::resource::sync_stream(handle); + + EXPECT_EQ(index.lists()[0]->get_size(), rows); +} + } // namespace cuvs::neighbors::ivf_sq From f4cee4f9b7d92e1a2699aa3f199e87fff9261d77 Mon Sep 17 00:00:00 2001 From: Dmitriy Vesnin Date: Sat, 4 Jul 2026 16:29:23 +0300 Subject: [PATCH 2/2] better comments + hpp --- cpp/include/cuvs/neighbors/common.hpp | 8 ++++++++ cpp/src/neighbors/ivf_list.cuh | 5 +++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2fd804f115..2b85af78ee 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -853,6 +853,14 @@ using enable_if_valid_list_t = typename enable_if_valid_list::type; * `cuvs::neighbors::ivf_pq::helpers::resize_list` which handle type casting internally. */ template +CUVS_EXPORT void resize_list(raft::resources const& res, + std::shared_ptr& orig_list, // NOLINT + const typename ListT::spec_type& spec, + typename ListT::size_type new_used_size, + typename ListT::size_type old_logical_size, + typename ListT::size_type old_used_size); + +template CUVS_EXPORT void resize_list(raft::resources const& res, std::shared_ptr& orig_list, // NOLINT const typename ListT::spec_type& spec, diff --git a/cpp/src/neighbors/ivf_list.cuh b/cpp/src/neighbors/ivf_list.cuh index 7c6f5d7e66..a251b7befc 100644 --- a/cpp/src/neighbors/ivf_list.cuh +++ b/cpp/src/neighbors/ivf_list.cuh @@ -69,8 +69,9 @@ CUVS_EXPORT void resize_list(raft::resources const& res, typename ListT::size_type old_logical_size, typename ListT::size_type old_used_size) { - // old_logical_size is the value stored in list::size. old_used_size is the copy extent from - // the old allocation and may include padding slots for interleaved list layouts. + // old_logical_size is the previous visible size from this index's list_sizes(). + // old_used_size is the old allocation copy extent and may include padded slots + // required by interleaved list layouts. bool skip_resize = false; if (orig_list) { if (new_used_size <= orig_list->indices.extent(0)) {