-
Notifications
You must be signed in to change notification settings - Fork 210
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 2 commits
7d5f827
f4cee4f
7cd89ac
4fad934
fba1eb4
a51ec2a
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,15 @@ | ||||||||||||||||||||||||
| /* | ||||||||||||||||||||||||
| * 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 <numeric> | ||||||||||||||||||||||||
| #include <vector> | ||||||||||||||||||||||||
|
|
||||||||||||||||||||||||
| namespace cuvs::neighbors::ivf_flat { | ||||||||||||||||||||||||
|
|
||||||||||||||||||||||||
| typedef AnnIVFFlatTest<float, float, int64_t> 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<float> 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<float>(row + col); | ||||||||||||||||||||||||
| } | ||||||||||||||||||||||||
| } | ||||||||||||||||||||||||
|
|
||||||||||||||||||||||||
| auto data = raft::make_device_matrix<float, int64_t>(handle, rows, dim); | ||||||||||||||||||||||||
| raft::copy(data.data_handle(), host_data.data(), host_data.size(), stream); | ||||||||||||||||||||||||
|
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 think we could initialize the data directly on device with zeros. We test the sizes instead of filled values. I assume the After the changes, we should |
||||||||||||||||||||||||
|
|
||||||||||||||||||||||||
| 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_GE(base_index.lists()[0]->indices_capacity(), rows); | ||||||||||||||||||||||||
|
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 think this makes the test more expressive |
||||||||||||||||||||||||
|
|
||||||||||||||||||||||||
| std::vector<int64_t> host_indices(grow_rows); | ||||||||||||||||||||||||
| std::iota(host_indices.begin(), host_indices.end(), base_rows); | ||||||||||||||||||||||||
| auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows); | ||||||||||||||||||||||||
| raft::copy(indices.data_handle(), host_indices.data(), host_indices.size(), stream); | ||||||||||||||||||||||||
|
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
raft has a device built-in method for it. |
||||||||||||||||||||||||
|
|
||||||||||||||||||||||||
| 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_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,77 @@ | ||
| /* | ||
| * 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 <numeric> | ||
| #include <vector> | ||
|
|
||
| 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; | ||
|
|
||
| std::vector<float> 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<float>(row + col); | ||
| } | ||
| } | ||
|
|
||
| auto data = raft::make_device_matrix<float, int64_t>(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<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_GE(index.lists()[0]->indices_capacity(), rows); | ||
|
|
||
| std::vector<int64_t> host_indices(grow_rows); | ||
| std::iota(host_indices.begin(), host_indices.end(), base_rows); | ||
| auto indices = raft::make_device_vector<int64_t, int64_t>(handle, grow_rows); | ||
| raft::copy(indices.data_handle(), host_indices.data(), host_indices.size(), 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_size(), rows); | ||
| } | ||
|
|
||
| } // namespace cuvs::neighbors::ivf_sq | ||
Uh oh!
There was an error while loading. Please reload this page.