Skip to content

Commit

Permalink
[BUG] Fix num_cta_per_query div (rapidsai#2107)
Browse files Browse the repository at this point in the history
This bug happened when trying to run a CAGRA index with `itopk=100` and `topk=100`. The `num_cta_per_query` variable was equal to 3 because 100 / 32 = 3.125 instead of ceildiv(100, 32) = 4. This resulted in the following error:
```
RuntimeError: RAFT failure at file=/opt/conda/conda-bld/work/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh line=183:
`num_cta_per_query` (3) * 32 must be equal to or greater than `topk` (100) when 'search_mode' is "multi-cta". (`num_cta_per_query`=max(`search_width`, `itopk_size`/32))
```

Authors:
  - Micka (https://github.com/lowener)

Approvers:
  - tsuki (https://github.com/enp1s0)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: rapidsai#2107
  • Loading branch information
lowener authored Jan 23, 2024
1 parent 558dc8f commit d89ab1b
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
7 changes: 4 additions & 3 deletions cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2023, NVIDIA CORPORATION.
* Copyright (c) 2023-2024, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -111,7 +111,8 @@ struct search : public search_plan_impl<DATA_T, INDEX_T, DISTANCE_T, SAMPLE_FILT
constexpr unsigned muti_cta_itopk_size = 32;
this->itopk_size = muti_cta_itopk_size;
search_width = 1;
num_cta_per_query = max(params.search_width, params.itopk_size / muti_cta_itopk_size);
num_cta_per_query =
max(params.search_width, raft::ceildiv(params.itopk_size, (size_t)muti_cta_itopk_size));
result_buffer_size = itopk_size + search_width * graph_degree;
typedef raft::Pow2<32> AlignBytes;
unsigned result_buffer_size_32 = AlignBytes::roundUp(result_buffer_size);
Expand Down Expand Up @@ -184,7 +185,7 @@ struct search : public search_plan_impl<DATA_T, INDEX_T, DISTANCE_T, SAMPLE_FILT
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`=max(`search_width`, ceildiv(`itopk_size`, 32)))",
num_cta_per_query,
topk);
}
Expand Down
2 changes: 1 addition & 1 deletion cpp/include/raft/neighbors/detail/cagra/search_plan.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ struct search_plan_impl : public search_plan_impl_base {
if (algo == search_algo::MULTI_CTA) {
mc_itopk_size = 32;
mc_search_width = 1;
mc_num_cta_per_query = max(search_width, itopk_size / 32);
mc_num_cta_per_query = max(search_width, raft::ceildiv(itopk_size, (size_t)32));
RAFT_LOG_DEBUG("# mc_itopk_size: %u", mc_itopk_size);
RAFT_LOG_DEBUG("# mc_search_width: %u", mc_search_width);
RAFT_LOG_DEBUG("# mc_num_cta_per_query: %u", mc_num_cta_per_query);
Expand Down

0 comments on commit d89ab1b

Please sign in to comment.