Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/branch-25.02' into rhdong/cagra-…
Browse files Browse the repository at this point in the history 
…merge
  • Loading branch information
@rhdong
rhdong committed Jan 27, 2025
2 parents ba244fb + 5527cdf commit 1dd5140
Show file tree
Hide file tree
Showing 5 changed files with 208 additions and 28 deletions.
35 changes: 22 additions & 13 deletions cpp/test/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,19 +138,28 @@ if(BUILD_TESTS)
)

ConfigureTest(
NAME
NEIGHBORS_ANN_CAGRA_TEST
PATH
neighbors/ann_cagra/bug_extreme_inputs_oob.cu
neighbors/ann_cagra/bug_multi_cta_crash.cu
neighbors/ann_cagra/test_float_uint32_t.cu
neighbors/ann_cagra/test_half_uint32_t.cu
neighbors/ann_cagra/test_int8_t_uint32_t.cu
neighbors/ann_cagra/test_uint8_t_uint32_t.cu
GPUS
1
PERCENT
100
NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu
neighbors/ann_cagra/bug_multi_cta_crash.cu GPUS 1 PERCENT 100
)

ConfigureTest(
NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST PATH neighbors/ann_cagra/test_float_uint32_t.cu GPUS
1 PERCENT 100
)

ConfigureTest(
NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST PATH neighbors/ann_cagra/test_half_uint32_t.cu GPUS 1
PERCENT 100
)

ConfigureTest(
NAME NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST PATH neighbors/ann_cagra/test_int8_t_uint32_t.cu GPUS
1 PERCENT 100
)

ConfigureTest(
NAME NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST PATH neighbors/ann_cagra/test_uint8_t_uint32_t.cu
GPUS 1 PERCENT 100
)

ConfigureTest(
Expand Down
184 changes: 174 additions & 10 deletions cpp/test/neighbors/ann_cagra.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1060,14 +1060,15 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam<AnnCagraInputs> {
inline std::vector<AnnCagraInputs> generate_inputs()
{
// TODO(tfeher): test MULTI_CTA kernel with search_width > 1 to allow multiple CTA per queries
// Charge graph dim, search algo and max_query parameter
std::vector<AnnCagraInputs> inputs = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
{1, 8, 17},
{1, 16}, // k
{graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
{16}, // k
{graph_build_algo::NN_DESCENT},
{search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
{0, 1, 10, 100}, // query size
{0, 10}, // query size
{0},
{256},
{1},
Expand All @@ -1076,11 +1077,30 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{true},
{0.995});

// Fixed dim, and changing neighbors and query size (output matrix size)
auto inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{1, 100},
{1000},
{8},
{1, 16}, // k
{graph_build_algo::NN_DESCENT},
{search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
{0}, // query size
{0},
{256},
{1},
{cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
{false},
{true},
{0.995});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// Varying dim and build algo.
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
{1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024}, // dim
{16}, // k
{1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 1024}, // dim
{16}, // k
{graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
{search_algo::AUTO},
{10},
Expand All @@ -1092,6 +1112,8 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{true},
{0.995});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// Varying team_size, graph_build_algo
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
Expand All @@ -1109,6 +1131,7 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{0.995});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// Varying graph_build_algo, itopk_size
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
Expand All @@ -1126,9 +1149,10 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{0.995});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// Varying n_rows, host_dataset
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{10000, 20000},
{10000},
{32},
{10},
{graph_build_algo::AUTO},
Expand All @@ -1143,7 +1167,8 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{0.985});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// a few PQ configurations
// A few PQ configurations.
// Varying dim, vq_n_centers
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{10000},
Expand Down Expand Up @@ -1171,7 +1196,8 @@ inline std::vector<AnnCagraInputs> generate_inputs()
}
}

// refinement options
// Refinement options
// Varying host_dataset, ivf_pq_search_refine_ratio
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{5000},
Expand All @@ -1190,10 +1216,11 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{1.0f, 2.0f, 3.0f});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// Varying dim, adding non_owning_memory_buffer_flag
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
{1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 619, 1024}, // dim
{1, 5, 8, 64, 137, 256, 619, 1024}, // dim
{10},
{graph_build_algo::IVF_PQ},
{search_algo::AUTO},
Expand All @@ -1213,6 +1240,143 @@ inline std::vector<AnnCagraInputs> generate_inputs()
return inputs;
}

const std::vector<AnnCagraInputs> inputs = generate_inputs();
inline std::vector<AnnCagraInputs> generate_addnode_inputs()
{
// changing dim
std::vector<AnnCagraInputs> inputs = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
{1, 8, 17, 64, 128, 137, 512, 1024}, // dim
{16}, // k
{graph_build_algo::NN_DESCENT},
{search_algo::AUTO},
{10},
{0},
{64},
{1},
{cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
{false},
{true},
{0.995});

// testing host and device datasets
auto inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{10000},
{32},
{10},
{graph_build_algo::AUTO},
{search_algo::AUTO},
{10},
{0}, // team_size
{64},
{1},
{cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
{false, true},
{false},
{0.985});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// a few PQ configurations
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{10000},
{192, 1024}, // dim
{16}, // k
{graph_build_algo::IVF_PQ},
{search_algo::AUTO},
{10},
{0},
{64},
{1},
{cuvs::distance::DistanceType::L2Expanded},
{false},
{true},
{0.6}); // don't demand high recall without refinement
for (uint32_t pq_len : {2}) { // for now, only pq_len = 2 is supported, more options coming soon
for (uint32_t vq_n_centers : {100}) {
for (auto input : inputs2) {
vpq_params ps{};
ps.pq_dim = input.dim / pq_len;
ps.vq_n_centers = vq_n_centers;
input.compression.emplace(ps);
inputs.push_back(input);
}
}
}

return inputs;
}

inline std::vector<AnnCagraInputs> generate_filtering_inputs()
{
// Charge graph dim, search algo
std::vector<AnnCagraInputs> inputs = raft::util::itertools::product<AnnCagraInputs>(
{100},
{1000},
{1, 8, 17},
{16}, // k
{graph_build_algo::NN_DESCENT},
{search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
{0}, // query size
{0},
{256},
{1},
{cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
{false},
{true},
{0.995});

// Fixed dim, and changing neighbors and query size (output matrix size)
auto inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{1, 100},
{1000},
{8},
{1, 16}, // k
{graph_build_algo::NN_DESCENT},
{search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
{0}, // query size
{0},
{256},
{1},
{cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct},
{false},
{true},
{0.995});
inputs.insert(inputs.end(), inputs2.begin(), inputs2.end());

// a few PQ configurations
inputs2 = raft::util::itertools::product<AnnCagraInputs>(
{100},
{10000},
{256}, // dim
{16}, // k
{graph_build_algo::IVF_PQ},
{search_algo::AUTO},
{10},
{0},
{64},
{1},
{cuvs::distance::DistanceType::L2Expanded},
{false},
{true},
{0.6}); // don't demand high recall without refinement
for (uint32_t pq_len : {2}) { // for now, only pq_len = 2 is supported, more options coming soon
for (uint32_t vq_n_centers : {100}) {
for (auto input : inputs2) {
vpq_params ps{};
ps.pq_dim = input.dim / pq_len;
ps.vq_n_centers = vq_n_centers;
input.compression.emplace(ps);
inputs.push_back(input);
}
}
}

return inputs;
}
const std::vector<AnnCagraInputs> inputs = generate_inputs();
const std::vector<AnnCagraInputs> inputs_addnode = generate_addnode_inputs();
const std::vector<AnnCagraInputs> inputs_filtering = generate_filtering_inputs();

} // namespace cuvs::neighbors::cagra
6 changes: 4 additions & 2 deletions cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,10 @@ TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge) { this->testCagra(); }
INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF_U32, ::testing::ValuesIn(inputs));
INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest,
AnnCagraAddNodesTestF_U32,
::testing::ValuesIn(inputs));
INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestF_U32, ::testing::ValuesIn(inputs));
::testing::ValuesIn(inputs_addnode));
INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest,
AnnCagraFilterTestF_U32,
::testing::ValuesIn(inputs_filtering));
INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest,
AnnCagraIndexMergeTestF_U32,
::testing::ValuesIn(inputs));
Expand Down
5 changes: 4 additions & 1 deletion cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,10 @@ TEST_P(AnnCagraIndexMergeTestI8_U32, AnnCagra) { this->testCagra(); }
INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestI8_U32, ::testing::ValuesIn(inputs));
INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest,
AnnCagraAddNodesTestI8_U32,
::testing::ValuesIn(inputs));
::testing::ValuesIn(inputs_addnode));
INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest,
AnnCagraFilterTestI8_U32,
::testing::ValuesIn(inputs_filtering));
INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest,
AnnCagraIndexMergeTestI8_U32,
::testing::ValuesIn(inputs));
Expand Down
6 changes: 4 additions & 2 deletions cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,10 @@ TEST_P(AnnCagraIndexMergeTestU8_U32, AnnCagra) { this->testCagra(); }
INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestU8_U32, ::testing::ValuesIn(inputs));
INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest,
AnnCagraAddNodesTestU8_U32,
::testing::ValuesIn(inputs));
INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestU8_U32, ::testing::ValuesIn(inputs));
::testing::ValuesIn(inputs_addnode));
INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest,
AnnCagraFilterTestU8_U32,
::testing::ValuesIn(inputs_filtering));
INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest,
AnnCagraIndexMergeTestU8_U32,
::testing::ValuesIn(inputs));
Expand Down

0 comments on commit 1dd5140

Please sign in to comment.