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

Moving MG functions into unified API + raft::device_resources_snmg as device resource type for MG functions #454

Open
wants to merge 13 commits into
base: branch-25.02
Choose a base branch
from
Open

Moving MG functions into unified API + raft::device_resources_snmg as device resource type for MG functions #454

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
ee98593
Account for RAFT update
viclafargue Nov 8, 2024
3a99b40
use new device_resources_snmg
viclafargue Nov 15, 2024
e16b68e
improved device_resources_snmg
viclafargue Nov 15, 2024
cdd5cfb
Merge branch 'branch-24.12' into account-for-raft-update
viclafargue Nov 15, 2024
96e69fc
switch from RAFT_LOG_INFO to RAFT_LOG_DEBUG for mg logs
viclafargue Nov 15, 2024
657bf9e
clique as device_resource
viclafargue Nov 25, 2024
1fdccd4
updating MG tests
viclafargue Nov 26, 2024
d7fff4c
Merge branch 'branch-25.02' into account-for-raft-update
viclafargue Jan 13, 2025
45a41fa
API unification (removal of mg namespace)
viclafargue Jan 20, 2025
b0c7ab9
Merge branch 'branch-25.02' into account-for-raft-update
viclafargue Jan 20, 2025
db12ab9
Adding CUVS_BUILD_MG_ALGOS macro back
viclafargue Jan 20, 2025
d5f1500
Merge branch 'branch-25.02' into account-for-raft-update
cjnolet Jan 28, 2025
db4cf11
Main requested changes
viclafargue Jan 29, 2025
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
2 changes: 1 addition & 1 deletion cpp/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,7 +300,6 @@ if(BUILD_SHARED_LIBS)
src/neighbors/mg/mg_cagra_int8_t_uint32_t.cu
src/neighbors/mg/mg_cagra_uint8_t_uint32_t.cu
src/neighbors/mg/omp_checks.cpp
src/neighbors/mg/nccl_comm.cpp
)
endif()

Expand Down Expand Up @@ -593,6 +592,7 @@ if(BUILD_SHARED_LIBS)
if(BUILD_MG_ALGOS)
target_compile_definitions(cuvs PUBLIC CUVS_BUILD_MG_ALGOS)
target_compile_definitions(cuvs_objs PUBLIC CUVS_BUILD_MG_ALGOS)
target_compile_definitions(cuvs-cagra-search PUBLIC CUVS_BUILD_MG_ALGOS)
endif()

if(BUILD_CAGRA_HNSWLIB)
Expand Down
21 changes: 10 additions & 11 deletions cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
#include "cuvs_ann_bench_utils.h"
#include "cuvs_cagra_wrapper.h"
#include <cuvs/neighbors/mg.hpp>
#include <raft/core/resource/nccl_clique.hpp>
#include <raft/core/device_resources_snmg.hpp>

namespace cuvs::bench {
using namespace cuvs::neighbors;
Expand All @@ -41,13 +41,12 @@ class cuvs_mg_cagra : public algo<T>, public algo_gpu {
};

cuvs_mg_cagra(Metric metric, int dim, const build_param& param, int concurrent_searches = 1)
: algo<T>(metric, dim), index_params_(param)
: algo<T>(metric, dim), index_params_(param), clique_()
{
index_params_.cagra_params.metric = parse_metric_type(metric);
index_params_.ivf_pq_build_params->metric = parse_metric_type(metric);

// init nccl clique outside as to not affect benchmark
const raft::comms::nccl_clique& clique = raft::resource::get_nccl_clique(handle_);
clique_.set_memory_pool(80);
}

void build(const T* dataset, size_t nrow) final;
Expand All @@ -69,7 +68,7 @@ class cuvs_mg_cagra : public algo<T>, public algo_gpu {

[[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
{
auto stream = raft::resource::get_cuda_stream(handle_);
auto stream = raft::resource::get_cuda_stream(clique_);
return stream;
}

Expand All @@ -87,7 +86,7 @@ class cuvs_mg_cagra : public algo<T>, public algo_gpu {
std::unique_ptr<algo<T>> copy() override;

private:
raft::device_resources handle_;
raft::device_resources_snmg clique_;
Copy link
Member

Choose a reason for hiding this comment

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

device_resources_snmg is a type of raft::resource. The pattern here is to accept the raft::resources instance and call the free function inside raft::resource::nccl_resource underneath to get the clique out of the resource. This way the APIs don't need to look any different.

float refine_ratio_;
build_param index_params_;
cuvs::neighbors::mg::search_params<cagra::search_params> search_params_;
Expand All @@ -105,7 +104,7 @@ void cuvs_mg_cagra<T, IdxT>::build(const T* dataset, size_t nrow)

auto dataset_view =
raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, nrow, dim_);
auto idx = cuvs::neighbors::mg::build(handle_, build_params, dataset_view);
auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view);
index_ =
std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>(
std::move(idx));
Expand All @@ -132,15 +131,15 @@ void cuvs_mg_cagra<T, IdxT>::set_search_dataset(const T* dataset, size_t nrow)
template <typename T, typename IdxT>
void cuvs_mg_cagra<T, IdxT>::save(const std::string& file) const
{
cuvs::neighbors::mg::serialize(handle_, *index_, file);
cuvs::neighbors::cagra::serialize(clique_, *index_, file);
}

template <typename T, typename IdxT>
void cuvs_mg_cagra<T, IdxT>::load(const std::string& file)
{
index_ =
std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>(
std::move(cuvs::neighbors::mg::deserialize_cagra<T, IdxT>(handle_, file)));
std::move(cuvs::neighbors::cagra::deserialize<T, IdxT>(clique_, file)));
}

template <typename T, typename IdxT>
Expand All @@ -163,8 +162,8 @@ void cuvs_mg_cagra<T, IdxT>::search_base(
auto distances_view =
raft::make_host_matrix_view<float, int64_t, raft::row_major>(distances, batch_size, k);

cuvs::neighbors::mg::search(
handle_, *index_, search_params_, queries_view, neighbors_view, distances_view);
cuvs::neighbors::cagra::search(
clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
}

template <typename T, typename IdxT>
Expand Down
22 changes: 11 additions & 11 deletions cpp/bench/ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
#include "cuvs_ann_bench_utils.h"
#include "cuvs_ivf_flat_wrapper.h"
#include <cuvs/neighbors/mg.hpp>
#include <raft/core/resource/nccl_clique.hpp>
#include <raft/core/device_resources_snmg.hpp>

namespace cuvs::bench {
using namespace cuvs::neighbors;
Expand All @@ -37,11 +37,11 @@ class cuvs_mg_ivf_flat : public algo<T>, public algo_gpu {
};

cuvs_mg_ivf_flat(Metric metric, int dim, const build_param& param)
: algo<T>(metric, dim), index_params_(param)
: algo<T>(metric, dim), index_params_(param), clique_()
{
index_params_.metric = parse_metric_type(metric);
// init nccl clique outside as to not affect benchmark
const raft::comms::nccl_clique& clique = raft::resource::get_nccl_clique(handle_);

clique_.set_memory_pool(80);
}

void build(const T* dataset, size_t nrow) final;
Expand All @@ -62,7 +62,7 @@ class cuvs_mg_ivf_flat : public algo<T>, public algo_gpu {

[[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
{
auto stream = raft::resource::get_cuda_stream(handle_);
auto stream = raft::resource::get_cuda_stream(clique_);
return stream;
}

Expand All @@ -73,7 +73,7 @@ class cuvs_mg_ivf_flat : public algo<T>, public algo_gpu {
std::unique_ptr<algo<T>> copy() override;

private:
raft::device_resources handle_;
raft::device_resources_snmg clique_;
build_param index_params_;
cuvs::neighbors::mg::search_params<ivf_flat::search_params> search_params_;
std::shared_ptr<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>
Expand All @@ -85,7 +85,7 @@ void cuvs_mg_ivf_flat<T, IdxT>::build(const T* dataset, size_t nrow)
{
auto dataset_view =
raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, IdxT(nrow), IdxT(dim_));
auto idx = cuvs::neighbors::mg::build(handle_, index_params_, dataset_view);
auto idx = cuvs::neighbors::ivf_flat::build(clique_, index_params_, dataset_view);
index_ = std::make_shared<
cuvs::neighbors::mg::index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>(std::move(idx));
}
Expand All @@ -105,15 +105,15 @@ void cuvs_mg_ivf_flat<T, IdxT>::set_search_param(const search_param_base& param)
template <typename T, typename IdxT>
void cuvs_mg_ivf_flat<T, IdxT>::save(const std::string& file) const
{
cuvs::neighbors::mg::serialize(handle_, *index_, file);
cuvs::neighbors::ivf_flat::serialize(clique_, *index_, file);
}

template <typename T, typename IdxT>
void cuvs_mg_ivf_flat<T, IdxT>::load(const std::string& file)
{
index_ = std::make_shared<
cuvs::neighbors::mg::index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>(
std::move(cuvs::neighbors::mg::deserialize_flat<T, IdxT>(handle_, file)));
std::move(cuvs::neighbors::ivf_flat::deserialize<T, IdxT>(clique_, file)));
}

template <typename T, typename IdxT>
Expand All @@ -133,8 +133,8 @@ void cuvs_mg_ivf_flat<T, IdxT>::search(
auto distances_view = raft::make_host_matrix_view<float, int64_t, raft::row_major>(
distances, IdxT(batch_size), IdxT(k));

cuvs::neighbors::mg::search(
handle_, *index_, search_params_, queries_view, neighbors_view, distances_view);
cuvs::neighbors::ivf_flat::search(
clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
}

} // namespace cuvs::bench
22 changes: 11 additions & 11 deletions cpp/bench/ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
#include "cuvs_ann_bench_utils.h"
#include "cuvs_ivf_pq_wrapper.h"
#include <cuvs/neighbors/mg.hpp>
#include <raft/core/resource/nccl_clique.hpp>
#include <raft/core/device_resources_snmg.hpp>

namespace cuvs::bench {
using namespace cuvs::neighbors;
Expand All @@ -37,11 +37,11 @@ class cuvs_mg_ivf_pq : public algo<T>, public algo_gpu {
};

cuvs_mg_ivf_pq(Metric metric, int dim, const build_param& param)
: algo<T>(metric, dim), index_params_(param)
: algo<T>(metric, dim), index_params_(param), clique_()
{
index_params_.metric = parse_metric_type(metric);
// init nccl clique outside as to not affect benchmark
const raft::comms::nccl_clique& clique = raft::resource::get_nccl_clique(handle_);

clique_.set_memory_pool(80);
}

void build(const T* dataset, size_t nrow) final;
Expand All @@ -62,7 +62,7 @@ class cuvs_mg_ivf_pq : public algo<T>, public algo_gpu {

[[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
{
auto stream = raft::resource::get_cuda_stream(handle_);
auto stream = raft::resource::get_cuda_stream(clique_);
return stream;
}

Expand All @@ -73,7 +73,7 @@ class cuvs_mg_ivf_pq : public algo<T>, public algo_gpu {
std::unique_ptr<algo<T>> copy() override;

private:
raft::device_resources handle_;
raft::device_resources_snmg clique_;
build_param index_params_;
cuvs::neighbors::mg::search_params<ivf_pq::search_params> search_params_;
std::shared_ptr<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>> index_;
Expand All @@ -84,7 +84,7 @@ void cuvs_mg_ivf_pq<T, IdxT>::build(const T* dataset, size_t nrow)
{
auto dataset_view =
raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, IdxT(nrow), IdxT(dim_));
auto idx = cuvs::neighbors::mg::build(handle_, index_params_, dataset_view);
auto idx = cuvs::neighbors::ivf_pq::build(clique_, index_params_, dataset_view);
index_ =
std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>>(
std::move(idx));
Expand All @@ -104,15 +104,15 @@ void cuvs_mg_ivf_pq<T, IdxT>::set_search_param(const search_param_base& param)
template <typename T, typename IdxT>
void cuvs_mg_ivf_pq<T, IdxT>::save(const std::string& file) const
{
cuvs::neighbors::mg::serialize(handle_, *index_, file);
cuvs::neighbors::ivf_pq::serialize(clique_, *index_, file);
}

template <typename T, typename IdxT>
void cuvs_mg_ivf_pq<T, IdxT>::load(const std::string& file)
{
index_ =
std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>>(
std::move(cuvs::neighbors::mg::deserialize_pq<T, IdxT>(handle_, file)));
std::move(cuvs::neighbors::ivf_pq::deserialize<T, IdxT>(clique_, file)));
}

template <typename T, typename IdxT>
Expand All @@ -132,8 +132,8 @@ void cuvs_mg_ivf_pq<T, IdxT>::search(
auto distances_view = raft::make_host_matrix_view<float, int64_t, raft::row_major>(
distances, IdxT(batch_size), IdxT(k));

cuvs::neighbors::mg::search(
handle_, *index_, search_params_, queries_view, neighbors_view, distances_view);
cuvs::neighbors::ivf_pq::search(
clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
}

} // namespace cuvs::bench
Loading
Loading