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 fffe5134ae..1566056eb6 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 */ @@ -282,11 +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], - 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_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_list.cuh b/cpp/src/neighbors/ivf_list.cuh index 24691463ff..8ff0465e47 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 */ diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh index 649a4f0720..9d8bb20764 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 */ @@ -354,11 +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], - 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); 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); } 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..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 @@ -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,12 @@ #include "../ann_ivf_flat.cuh" +#include +#include +#include + +#include + namespace cuvs::neighbors::ivf_flat { typedef AnnIVFFlatTest AnnIVFFlatTestF_float; @@ -19,4 +25,65 @@ 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; + + auto data = raft::make_device_matrix(handle, rows, dim); + raft::matrix::fill(handle, data.view(), 0.0f); + + 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_EQ(base_index.lists()[0]->indices_capacity(), raft::Pow2::roundUp(rows)); + + auto indices = raft::make_device_vector(handle, grow_rows); + 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); + 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(), base_index.lists()[0].get()); + 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..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 @@ -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,12 @@ #include "../ann_ivf_sq.cuh" +#include +#include +#include + +#include + namespace cuvs::neighbors::ivf_sq { typedef AnnIVFSQTest AnnIVFSQTestF_float; @@ -14,4 +20,54 @@ 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; + + auto data = raft::make_device_matrix(handle, rows, dim); + raft::matrix::fill(handle, data.view(), 0.0f); + + 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_EQ(index.lists()[0]->indices_capacity(), raft::Pow2::roundUp(rows)); + auto* base_list = index.lists()[0].get(); + + auto indices = raft::make_device_vector(handle, grow_rows); + 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); + 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(), base_list); + EXPECT_EQ(index.lists()[0]->get_size(), rows); +} + } // namespace cuvs::neighbors::ivf_sq