From 7d5f8270be5d04ac873d409483fd248e270f10cf Mon Sep 17 00:00:00 2001 From: Dmitriy Vesnin Date: Fri, 3 Jul 2026 22:59:03 +0300 Subject: [PATCH 1/5] 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/5] 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)) { From 7cd89ac74881bc3c06e97ea55a7f7ba8a8e906ba Mon Sep 17 00:00:00 2001 From: Dmitriy Vesnin Date: Tue, 7 Jul 2026 18:42:52 +0300 Subject: [PATCH 3/5] improve tests --- .../ann_ivf_flat/test_float_int64_t.cu | 23 +++++++----------- .../ann_ivf_sq/test_float_int64_t.cu | 24 ++++++++----------- 2 files changed, 19 insertions(+), 28 deletions(-) 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 9e26942bd1..d4499ed6d9 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 @@ -7,8 +7,11 @@ #include "../ann_ivf_flat.cuh" -#include -#include +#include +#include +#include + +#include namespace cuvs::neighbors::ivf_flat { @@ -32,15 +35,8 @@ TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) 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); + raft::matrix::fill(handle, data.view(), 0.0f); index_params params; params.n_lists = 1; @@ -60,12 +56,10 @@ TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) raft::resource::sync_stream(handle); ASSERT_EQ(base_index.lists()[0]->get_size(), base_rows); - ASSERT_GE(base_index.lists()[0]->indices_capacity(), rows); + ASSERT_EQ(base_index.lists()[0]->indices_capacity(), raft::Pow2::roundUp(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); + raft::linalg::range(indices.data_handle(), base_rows, rows, stream); auto grow_data_view = raft::make_device_matrix_view( data.data_handle() + base_rows * dim, grow_rows, dim); @@ -78,6 +72,7 @@ TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) base_index); raft::resource::sync_stream(handle); + ASSERT_EQ(first_grown_index.lists()[0].get(), base_index.lists()[0].get()); ASSERT_EQ(first_grown_index.lists()[0]->get_size(), rows); auto second_grown_index = 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 565b336c03..171e953e01 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 @@ -7,8 +7,11 @@ #include "../ann_ivf_sq.cuh" -#include -#include +#include +#include +#include + +#include namespace cuvs::neighbors::ivf_sq { @@ -27,15 +30,8 @@ TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity) 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); + raft::matrix::fill(handle, data.view(), 0.0f); index_params params; params.n_lists = 1; @@ -54,12 +50,11 @@ TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity) raft::resource::sync_stream(handle); ASSERT_EQ(index.lists()[0]->get_size(), base_rows); - ASSERT_GE(index.lists()[0]->indices_capacity(), rows); + ASSERT_EQ(index.lists()[0]->indices_capacity(), raft::Pow2::roundUp(rows)); + auto* base_list = index.lists()[0].get(); - 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); + raft::linalg::range(indices.data_handle(), base_rows, rows, stream); auto grow_data_view = raft::make_device_matrix_view( data.data_handle() + base_rows * dim, grow_rows, dim); @@ -71,6 +66,7 @@ TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity) &index); raft::resource::sync_stream(handle); + EXPECT_EQ(index.lists()[0].get(), base_list); EXPECT_EQ(index.lists()[0]->get_size(), rows); } From 4fad9342b6986dca975578c9be6b3ffeaa980af2 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 8 Jul 2026 15:58:21 +0200 Subject: [PATCH 4/5] Pad the list extents to account for interleaved storage --- cpp/include/cuvs/neighbors/common.hpp | 8 -------- cpp/include/cuvs/neighbors/ivf_flat.hpp | 6 ++++-- cpp/include/cuvs/neighbors/ivf_sq.hpp | 5 +++-- cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh | 8 ++------ cpp/src/neighbors/ivf_list.cuh | 18 ++---------------- cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh | 8 ++------ 6 files changed, 13 insertions(+), 40 deletions(-) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2b85af78ee..2fd804f115 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -853,14 +853,6 @@ 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/include/cuvs/neighbors/ivf_flat.hpp b/cpp/include/cuvs/neighbors/ivf_flat.hpp index d2e0015498..33a268c094 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.hpp +++ b/cpp/include/cuvs/neighbors/ivf_flat.hpp @@ -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 */ @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -110,7 +111,8 @@ struct list_spec { /** Determine the extents of an array enough to hold a given amount of data. */ constexpr auto make_list_extents(SizeT n_rows) const -> list_extents { - return raft::make_extents(n_rows, dim); + return raft::make_extents(raft::round_up_safe(n_rows, SizeT{kIndexGroupSize}), + dim); } }; diff --git a/cpp/include/cuvs/neighbors/ivf_sq.hpp b/cpp/include/cuvs/neighbors/ivf_sq.hpp index 6ac765213c..09f0eddf3c 100644 --- a/cpp/include/cuvs/neighbors/ivf_sq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_sq.hpp @@ -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 */ @@ -118,7 +118,8 @@ struct CUVS_EXPORT list_spec { constexpr auto make_list_extents(SizeT n_rows) const -> list_extents { uint32_t padded = raft::round_up_safe(dim, kVecLen); - return raft::make_extents(n_rows, padded); + return raft::make_extents(raft::round_up_safe(n_rows, SizeT{kIndexGroupSize}), + padded); } }; diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh index fa7c0c36a2..1566056eb6 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh @@ -282,12 +282,8 @@ void extend(raft::resources const& handle, raft::resource::sync_stream(handle); auto& lists = index->lists(); for (uint32_t label = 0; label < n_lists; label++) { - ivf::resize_list(handle, - lists[label], - list_device_spec, - new_list_sizes[label], - old_list_sizes[label], - raft::Pow2::roundUp(old_list_sizes[label])); + ivf::resize_list( + handle, lists[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]); } } // Update the pointers and the sizes diff --git a/cpp/src/neighbors/ivf_list.cuh b/cpp/src/neighbors/ivf_list.cuh index a251b7befc..8ff0465e47 100644 --- a/cpp/src/neighbors/ivf_list.cuh +++ b/cpp/src/neighbors/ivf_list.cuh @@ -66,17 +66,13 @@ 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 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)) { - auto shared_list_size = old_logical_size; - if (new_used_size <= old_logical_size || + auto shared_list_size = old_used_size; + if (new_used_size <= old_used_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 @@ -108,16 +104,6 @@ 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 6f9268997c..9d8bb20764 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh @@ -354,12 +354,8 @@ void extend_inplace(raft::resources const& handle, raft::resource::sync_stream(handle); auto& lists = index->lists(); for (uint32_t label = 0; label < n_lists; label++) { - ivf::resize_list(handle, - lists[label], - list_device_spec, - new_list_sizes[label], - old_list_sizes[label], - raft::Pow2::roundUp(old_list_sizes[label])); + ivf::resize_list( + handle, lists[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]); } } ivf::detail::recompute_internal_state(handle, *index); From fba1eb4461912cc452ee63f987ef88536e3b76c4 Mon Sep 17 00:00:00 2001 From: achirkin Date: Thu, 9 Jul 2026 07:18:42 +0200 Subject: [PATCH 5/5] Fix serialize: store actual used size rather than the padded extents (don't round up) --- cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh | 11 +++-------- cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh | 11 +++-------- 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh index ef1898c17f..70b789ddc7 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.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 */ @@ -16,7 +16,6 @@ #include #include #include -#include #include @@ -27,7 +26,7 @@ namespace cuvs::neighbors::ivf_flat::detail { // backward compatibility. // TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward // compatible fashion. -constexpr int serialization_version = 4; +constexpr int serialization_version = 5; /** * Save the index to file. @@ -72,11 +71,7 @@ void serialize(raft::resources const& handle, std::ostream& os, const index list_store_spec{index_.dim(), true}; for (uint32_t label = 0; label < index_.n_lists(); label++) { - ivf::serialize_list(handle, - os, - index_.lists()[label], - list_store_spec, - raft::Pow2::roundUp(sizes_host(label))); + ivf::serialize_list(handle, os, index_.lists()[label], list_store_spec, sizes_host(label)); } raft::resource::sync_stream(handle); } diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh index 5e455302bb..22d89ba483 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.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 */ @@ -15,13 +15,12 @@ #include #include #include -#include #include namespace cuvs::neighbors::ivf_sq::detail { -constexpr int serialization_version = 1; +constexpr int serialization_version = 2; template void serialize(raft::resources const& handle, std::ostream& os, const index& index_) @@ -63,11 +62,7 @@ void serialize(raft::resources const& handle, std::ostream& os, const index list_store_spec{index_.dim(), true}; for (uint32_t label = 0; label < index_.n_lists(); label++) { - ivf::serialize_list(handle, - os, - index_.lists()[label], - list_store_spec, - raft::Pow2::roundUp(sizes_host(label))); + ivf::serialize_list(handle, os, index_.lists()[label], list_store_spec, sizes_host(label)); } raft::resource::sync_stream(handle); }