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

Account for RAFT update to SNMG APIs #454

Open
wants to merge 7 commits into
base: branch-24.12
Choose a base branch
from
Open

Account for RAFT update to SNMG APIs #454

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 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
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
1 change: 0 additions & 1 deletion cpp/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,7 +298,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
19 changes: 9 additions & 10 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_;
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::mg::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::mg::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::mg::deserialize_cagra<T, IdxT>(clique_, file)));
}

template <typename T, typename IdxT>
Expand All @@ -164,7 +163,7 @@ void cuvs_mg_cagra<T, IdxT>::search_base(
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);
clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
}

template <typename T, typename IdxT>
Expand Down
20 changes: 10 additions & 10 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::mg::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::mg::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::mg::deserialize_flat<T, IdxT>(clique_, file)));
}

template <typename T, typename IdxT>
Expand All @@ -134,7 +134,7 @@ void cuvs_mg_ivf_flat<T, IdxT>::search(
distances, IdxT(batch_size), IdxT(k));

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

} // namespace cuvs::bench
20 changes: 10 additions & 10 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::mg::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::mg::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::mg::deserialize_pq<T, IdxT>(clique_, file)));
}

template <typename T, typename IdxT>
Expand All @@ -133,7 +133,7 @@ void cuvs_mg_ivf_pq<T, IdxT>::search(
distances, IdxT(batch_size), IdxT(k));

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

} // namespace cuvs::bench
Loading
Loading