Skip to content
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.

Sign up for GitHub

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

[BUG] Fix search parameter check in CAGRA #1784

Merged
merged 5 commits into from
Aug 30, 2023
Merged

[BUG] Fix search parameter check in CAGRA #1784

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b87d4e0
Fix search papameter check in CAGRA
enp1s0 Aug 29, 2023
250e0e6
Merge branch 'branch-23.10' into fix-CAGRA-config-check
enp1s0 Aug 29, 2023
80400de
Update cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh
enp1s0 Aug 30, 2023
0368e2c
Fix comment
enp1s0 Aug 30, 2023
496689e
Merge branch 'branch-23.10' into fix-CAGRA-config-check
enp1s0 Aug 30, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,16 @@ struct search : public search_plan_impl<DATA_T, INDEX_T, DISTANCE_T> {
topk_workspace.resize(topk_workspace_size, resource::get_cuda_stream(res));
}

virtual void check(const uint32_t topk)
enp1s0 marked this conversation as resolved.
Show resolved Hide resolved
{
RAFT_EXPECTS(num_cta_per_query * 32 >= topk,
"`num_cta_per_query` (%u) * 32 must be equal to or greater than "
"`topk` (%u) when 'search_mode' is \"multi-cta\". "
"(`num_cta_per_query`=max(`search_width`, `itopk_size`/32))",
num_cta_per_query,
topk);
}

~search() {}

void operator()(raft::resources const& res,
Expand Down
11 changes: 2 additions & 9 deletions cpp/include/raft/neighbors/detail/cagra/search_plan.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,17 +250,10 @@ struct search_plan_impl : public search_plan_impl_base {
}
}

void check(uint32_t topk)
virtual void check(const uint32_t topk)
{
// For multi-CTA and multi kernel
enp1s0 marked this conversation as resolved.
Show resolved Hide resolved
RAFT_EXPECTS(topk <= itopk_size, "topk must be smaller than itopk_size = %lu", itopk_size);
if (algo == search_algo::MULTI_CTA) {
uint32_t mc_num_cta_per_query = max(search_width, itopk_size / 32);
RAFT_EXPECTS(mc_num_cta_per_query * 32 >= topk,
"`mc_num_cta_per_query` (%u) * 32 must be equal to or greater than "
"`topk` /%u) when 'search_mode' is \"multi-cta\"",
mc_num_cta_per_query,
topk);
}
}

inline void check_params()
Expand Down