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

IVF-PQ: use device properties helper #1035

Merged
merged 1 commit into from
Nov 18, 2022
Merged

IVF-PQ: use device properties helper #1035

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
12 changes: 4 additions & 8 deletions cpp/include/raft/spatial/knn/detail/ivf_pq_search.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1019,7 +1019,8 @@ struct ivfpq_compute_similarity {
* beyond this limit do not consider increasing the number of active blocks per SM
* would improve locality anymore.
*/
static inline auto select(bool manage_local_topk,
static inline auto select(const cudaDeviceProp& dev_props,
bool manage_local_topk,
int locality_hint,
double preferred_shmem_carveout,
uint32_t pq_bits,
Expand All @@ -1029,12 +1030,6 @@ struct ivfpq_compute_similarity {
uint32_t n_probes,
uint32_t topk) -> selected
{
cudaDeviceProp dev_props;
{
int cur_dev;
RAFT_CUDA_TRY(cudaGetDevice(&cur_dev));
RAFT_CUDA_TRY(cudaGetDeviceProperties(&dev_props, cur_dev));
}
// Shared memory for storing the lookup table
size_t lut_mem = sizeof(LutT) * (pq_dim << pq_bits);
// Shared memory for storing pre-computed pieces to speedup the lookup table construction
Expand Down Expand Up @@ -1364,7 +1359,8 @@ void ivfpq_search_worker(const handle_t& handle,
}

auto search_instance =
ivfpq_compute_similarity<IdxT, ScoreT, LutT>::select(manage_local_topk,
ivfpq_compute_similarity<IdxT, ScoreT, LutT>::select(handle.get_device_properties(),
manage_local_topk,
coresidency,
preferred_shmem_carveout,
index.pq_bits(),
Expand Down