-
Notifications
You must be signed in to change notification settings - Fork 207
Fix IVF extend list resize bug #2297
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -1,12 +1,18 @@ | ||||||
| /* | ||||||
| * 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 | ||||||
| */ | ||||||
|
|
||||||
| #include <gtest/gtest.h> | ||||||
|
|
||||||
| #include "../ann_ivf_flat.cuh" | ||||||
|
|
||||||
| #include <raft/core/resource/cuda_stream.hpp> | ||||||
| #include <raft/linalg/init.cuh> | ||||||
| #include <raft/matrix/init.cuh> | ||||||
|
|
||||||
| #include <optional> | ||||||
|
|
||||||
| namespace cuvs::neighbors::ivf_flat { | ||||||
|
|
||||||
| typedef AnnIVFFlatTest<float, float, int64_t> 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<float, int64_t>(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<const float, int64_t>(data.data_handle(), rows, dim); | ||||||
| auto empty_index = build(handle, params, all_data_view); | ||||||
|
|
||||||
| auto base_data_view = | ||||||
| raft::make_device_matrix_view<const float, int64_t>(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<kIndexGroupSize>::roundUp(rows)); | ||||||
|
|
||||||
| auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows); | ||||||
| raft::linalg::range(indices.data_handle(), base_rows, rows, stream); | ||||||
|
|
||||||
| auto grow_data_view = raft::make_device_matrix_view<const float, int64_t>( | ||||||
| data.data_handle() + base_rows * dim, grow_rows, dim); | ||||||
| auto grow_indices_view = | ||||||
| raft::make_device_vector_view<const int64_t, int64_t>(indices.data_handle(), grow_rows); | ||||||
| auto first_grown_index = | ||||||
| extend(handle, | ||||||
| grow_data_view, | ||||||
| std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(grow_indices_view), | ||||||
| base_index); | ||||||
| raft::resource::sync_stream(handle); | ||||||
|
|
||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
I'd suggest to test that the first grown list reuses the base list without reallocation. |
||||||
| 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<raft::device_vector_view<const int64_t, int64_t>>(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 | ||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,17 +1,73 @@ | ||
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| #include <gtest/gtest.h> | ||
|
|
||
| #include "../ann_ivf_sq.cuh" | ||
|
|
||
| #include <raft/core/resource/cuda_stream.hpp> | ||
| #include <raft/linalg/init.cuh> | ||
| #include <raft/matrix/init.cuh> | ||
|
|
||
| #include <optional> | ||
|
|
||
| namespace cuvs::neighbors::ivf_sq { | ||
|
|
||
| typedef AnnIVFSQTest<float, float, int64_t> AnnIVFSQTestF_float; | ||
| TEST_P(AnnIVFSQTestF_float, AnnIVFSQ) { this->testAll(); } | ||
|
|
||
| INSTANTIATE_TEST_CASE_P(AnnIVFSQTest, AnnIVFSQTestF_float, ::testing::ValuesIn(inputs)); | ||
|
|
||
| TEST(AnnIVFSQTest, ExtendInPlaceUpdatesListSizeWithinCapacity) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'd have the same reviews for this test. |
||
| { | ||
| 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<float, int64_t>(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<const float, int64_t>(data.data_handle(), rows, dim); | ||
| auto index = build(handle, params, all_data_view); | ||
|
|
||
| auto base_data_view = | ||
| raft::make_device_matrix_view<const float, int64_t>(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<kIndexGroupSize>::roundUp(rows)); | ||
| auto* base_list = index.lists()[0].get(); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why do we need to get the base_list pointer here? Can't we later get it at line 69? |
||
|
|
||
| auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows); | ||
| raft::linalg::range(indices.data_handle(), base_rows, rows, stream); | ||
|
|
||
| auto grow_data_view = raft::make_device_matrix_view<const float, int64_t>( | ||
| data.data_handle() + base_rows * dim, grow_rows, dim); | ||
| auto grow_indices_view = | ||
| raft::make_device_vector_view<const int64_t, int64_t>(indices.data_handle(), grow_rows); | ||
| extend(handle, | ||
| grow_data_view, | ||
| std::make_optional<raft::device_vector_view<const int64_t, int64_t>>(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 | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I guess there are 3 cases:
If I understand it correctly, your current version will do both size update and reallocation for the 2nd case.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Only if
orig_list->size.compare_exchange_strong(shared_list_size, new_used_size)returns falseit returns false if orig_list->size was already mutated before this resize_list call ( if the current shared list size is not old_logical_size anymore)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's right. I misunderstood the logics.