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

Reduce cagra build binary size #334

Open
wants to merge 4 commits into
base: branch-24.10
Choose a base branch
from
Open

Reduce cagra build binary size #334

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
48fe461
Use compiled IVF-PQ and NN descent in CAGRA build
tfeher Sep 18, 2024
2387ec3
Merge remote-tracking branch 'origin/branch-24.10' into cagra_build_b…
tfeher Sep 18, 2024
d4d3f20
Merge remote-tracking branch 'origin/branch-24.10' into cagra_build_b…
tfeher Sep 26, 2024
5409721
Merge branch 'branch-24.10' into cagra_build_binary_size
cjnolet Sep 27, 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
15 changes: 5 additions & 10 deletions cpp/src/neighbors/detail/cagra/cagra_build.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,9 @@

#include <cuvs/distance/distance.hpp>
#include <cuvs/neighbors/ivf_pq.hpp>
#include <cuvs/neighbors/nn_descent.hpp>
#include <cuvs/neighbors/refine.hpp>

// TODO: Fixme- this needs to be migrated
#include "../../ivf_pq/ivf_pq_build.cuh"
#include "../../ivf_pq/ivf_pq_search.cuh"
#include "../../nn_descent.cuh"

// TODO: This shouldn't be calling spatial/knn APIs
#include "../ann_utils.cuh"

Expand Down Expand Up @@ -156,8 +152,7 @@ void build_knn_graph(
}();

RAFT_LOG_DEBUG("# Building IVF-PQ index %s", model_name.c_str());
auto index =
cuvs::neighbors::ivf_pq::detail::build<DataT, int64_t>(res, pq.build_params, dataset);
auto index = cuvs::neighbors::ivf_pq::build(res, pq.build_params, dataset);

//
// search top (k + 1) neighbors
Expand All @@ -169,7 +164,8 @@ void build_knn_graph(
const auto num_queries = dataset.extent(0);

// Use the same maximum batch size as the ivf_pq::search to avoid allocating more than needed.
using cuvs::neighbors::ivf_pq::detail::kMaxQueries;
constexpr uint32_t kMaxQueries = 4096;

// Heuristic: the build_knn_graph code should use only a fraction of the workspace memory; the
// rest should be used by the ivf_pq::search. Here we say that the workspace size should be a good
// multiple of what is required for the I/O batching below.
Expand Down Expand Up @@ -357,8 +353,7 @@ void build_knn_graph(
raft::host_matrix_view<IdxT, int64_t, raft::row_major> knn_graph,
cuvs::neighbors::nn_descent::index_params build_params)
{
auto nn_descent_idx = cuvs::neighbors::nn_descent::index<IdxT>(res, knn_graph);
cuvs::neighbors::nn_descent::build<DataT, IdxT>(res, build_params, dataset, nn_descent_idx);
auto nn_descent_idx = cuvs::neighbors::nn_descent::build(res, build_params, dataset);
Copy link
Contributor

Choose a reason for hiding this comment

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

This is relevant: #264 (comment)


using internal_IdxT = typename std::make_unsigned<IdxT>::type;
using g_accessor = typename decltype(nn_descent_idx.graph())::accessor_type;
Expand Down
5 changes: 0 additions & 5 deletions cpp/src/neighbors/detail/cagra/cagra_search.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@

#include "factory.cuh"
#include "search_plan.cuh"
#include "search_single_cta_inst.cuh"

#include <raft/core/device_mdspan.hpp>
#include <raft/core/host_mdspan.hpp>
Expand All @@ -29,10 +28,6 @@
#include <cuvs/distance/distance.hpp>

#include <cuvs/neighbors/cagra.hpp>

// TODO: Fix these when ivf methods are moved over
#include "../../ivf_common.cuh"
#include "../../ivf_pq/ivf_pq_search.cuh"
#include <cuvs/neighbors/common.hpp>

// TODO: This shouldn't be calling spatial/knn apis
Expand Down
Loading