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

Fix max_queries for CAGRA #2081

Merged
merged 4 commits into from
Jan 9, 2024
Merged

Fix max_queries for CAGRA #2081

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d2f6055
Small fixes for CAGRA: max_queries auto selection
lowener Jan 5, 2024
fcd601f
Fix max_queries
lowener Jan 7, 2024
d91e152
Compute maxgridsize instead of hardcoding it
lowener Jan 9, 2024
ea1e551
Merge branch 'branch-24.02' into 24.02-cagra-smallfix
lowener Jan 9, 2024
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
4 changes: 2 additions & 2 deletions cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh
View file
Open in desktop
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 @@ -118,7 +118,7 @@ void search_main(raft::resources const& res,
RAFT_EXPECTS(queries.extent(1) == index.dim(), "Queries and index dim must match");
const uint32_t topk = neighbors.extent(1);

if (params.max_queries == 0) { params.max_queries = queries.extent(0); }
if (params.max_queries == 0) { params.max_queries = std::min<size_t>(queries.extent(0), 65535); }
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be more robust to compute the value, perhaps using std::limits for int 16, MAX_UNSIGNED_SHORT, or (1<<16)-1.


common::nvtx::range<common::nvtx::domain::raft> fun_scope(
"cagra::search(max_queries = %u, k = %u, dim = %zu)", params.max_queries, topk, index.dim());
Expand Down
4 changes: 2 additions & 2 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
@@ -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 @@ -147,7 +147,7 @@ struct search_plan_impl : public search_plan_impl_base {
// defines hash_bitlen, small_hash_bitlen, small_hash_reset interval, hash_size
inline void calc_hashmap_params(raft::resources const& res)
{
// for multipel CTA search
// for multiple CTA search
uint32_t mc_num_cta_per_query = 0;
uint32_t mc_search_width = 0;
uint32_t mc_itopk_size = 0;
Expand Down
5 changes: 2 additions & 3 deletions notebooks/utils.py
View file
Open in desktop
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 @@ -73,7 +73,7 @@ def benchmark_runs(self):
self.timings.append(t1 - t0)


def load_dataset(dataset_url, work_folder=None):
def load_dataset(dataset_url="http://ann-benchmarks.com/sift-128-euclidean.hdf5", work_folder=None):
"""Download dataset from url. It is expected that the dataset contains a hdf5 file in ann-benchmarks format

Parameters
Expand All @@ -82,7 +82,6 @@ def load_dataset(dataset_url, work_folder=None):
work_folder name of the local folder to store the dataset

"""
dataset_url = "http://ann-benchmarks.com/sift-128-euclidean.hdf5"
dataset_filename = dataset_url.split("/")[-1]

# We'll need to load store some data in this tutorial
Expand Down
12 changes: 6 additions & 6 deletions python/pylibraft/pylibraft/neighbors/cagra/cagra.pyx
View file
Open in desktop
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 @@ -91,7 +91,7 @@ from pylibraft.neighbors.common cimport _get_metric_string


cdef class IndexParams:
""""
"""
Parameters to build index for CAGRA nearest neighbor search

Parameters
Expand All @@ -107,10 +107,10 @@ cdef class IndexParams:
build_algo: string denoting the graph building algorithm to use,
default = "ivf_pq"
Valid values for algo: ["ivf_pq", "nn_descent"], where
- ivf_pq will use the IVF-PQ algorithm for building the knn graph
- nn_descent (experimental) will use the NN-Descent algorithm for
building the knn graph. It is expected to be generally
faster than ivf_pq.
- ivf_pq will use the IVF-PQ algorithm for building the knn graph
- nn_descent (experimental) will use the NN-Descent algorithm for
building the knn graph. It is expected to be generally
faster than ivf_pq.
"""
cdef c_cagra.index_params params

Expand Down