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

Separate CAGRA index type from internal idx type #1664

Merged
merged 7 commits into from
Jul 27, 2023
Merged

Separate CAGRA index type from internal idx type #1664

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
d1b53b1
Separate cagra index type from internal idx type
tfeher Jul 21, 2023
a617988
Merge branch 'branch-23.08' into cagra_matrix_idx
tfeher Jul 23, 2023
69effcd
Merge remote-tracking branch 'origin/branch-23.08' into cagra_matrix_idx
tfeher Jul 25, 2023
6436d96
Merge branch 'branch-23.08' into cagra_matrix_idx
tfeher Jul 25, 2023
f43d8e9
Merge branch 'branch-23.08' into cagra_matrix_idx
tfeher Jul 26, 2023
a00202a
Merge branch 'branch-23.08' into cagra_matrix_idx
tfeher Jul 26, 2023
d7d40bb
Merge branch 'branch-23.08' into cagra_matrix_idx
tfeher Jul 27, 2023
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
16 changes: 9 additions & 7 deletions cpp/bench/ann/src/raft/raft_cagra_wrapper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ class RaftCagra : public ANN<T> {
AlgoProperty get_property() const override
{
AlgoProperty property;
property.dataset_memory_type = MemoryType::Host;
property.dataset_memory_type = MemoryType::HostMmap;
property.query_memory_type = MemoryType::Device;
property.need_dataset_when_search = true;
return property;
Expand Down Expand Up @@ -106,12 +106,13 @@ RaftCagra<T, IdxT>::RaftCagra(Metric metric, int dim, const BuildParam& param)
template <typename T, typename IdxT>
void RaftCagra<T, IdxT>::build(const T* dataset, size_t nrow, cudaStream_t)
{
if (get_property().dataset_memory_type == MemoryType::Host) {
auto dataset_view = raft::make_host_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dimension_);
if (get_property().dataset_memory_type != MemoryType::Device) {
auto dataset_view =
raft::make_host_matrix_view<const T, int64_t>(dataset, IdxT(nrow), dimension_);
index_.emplace(raft::neighbors::cagra::build(handle_, index_params_, dataset_view));
} else {
auto dataset_view =
raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dimension_);
raft::make_device_matrix_view<const T, int64_t>(dataset, IdxT(nrow), dimension_);
index_.emplace(raft::neighbors::cagra::build(handle_, index_params_, dataset_view));
}
return;
Expand Down Expand Up @@ -152,9 +153,10 @@ void RaftCagra<T, IdxT>::search(
neighbors_IdxT = neighbors_storage.data();
}

auto queries_view = raft::make_device_matrix_view<const T, IdxT>(queries, batch_size, dimension_);
auto neighbors_view = raft::make_device_matrix_view<IdxT, IdxT>(neighbors_IdxT, batch_size, k);
auto distances_view = raft::make_device_matrix_view<float, IdxT>(distances, batch_size, k);
auto queries_view =
raft::make_device_matrix_view<const T, int64_t>(queries, batch_size, dimension_);
auto neighbors_view = raft::make_device_matrix_view<IdxT, int64_t>(neighbors_IdxT, batch_size, k);
auto distances_view = raft::make_device_matrix_view<float, int64_t>(distances, batch_size, k);

raft::neighbors::cagra::search(
handle_, search_params_, *index_, queries_view, neighbors_view, distances_view);
Expand Down
20 changes: 10 additions & 10 deletions cpp/bench/prims/neighbors/cagra_bench.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,9 @@ struct CagraBench : public fixture {
explicit CagraBench(const params& ps)
: fixture(true),
params_(ps),
queries_(make_device_matrix<T, IdxT>(handle, ps.n_queries, ps.n_dims)),
dataset_(make_device_matrix<T, IdxT>(handle, ps.n_samples, ps.n_dims)),
knn_graph_(make_device_matrix<IdxT, IdxT>(handle, ps.n_samples, ps.degree))
queries_(make_device_matrix<T, int64_t>(handle, ps.n_queries, ps.n_dims)),
dataset_(make_device_matrix<T, int64_t>(handle, ps.n_samples, ps.n_dims)),
knn_graph_(make_device_matrix<IdxT, int64_t>(handle, ps.n_samples, ps.degree))
{
// Generate random dataset and queriees
raft::random::RngState state{42};
Expand Down Expand Up @@ -87,11 +87,11 @@ struct CagraBench : public fixture {
search_params.thread_block_size = params_.block_size;
search_params.num_parents = params_.num_parents;

auto indices = make_device_matrix<IdxT, IdxT>(handle, params_.n_queries, params_.k);
auto distances = make_device_matrix<float, IdxT>(handle, params_.n_queries, params_.k);
auto ind_v = make_device_matrix_view<IdxT, IdxT, row_major>(
auto indices = make_device_matrix<IdxT, int64_t>(handle, params_.n_queries, params_.k);
auto distances = make_device_matrix<float, int64_t>(handle, params_.n_queries, params_.k);
auto ind_v = make_device_matrix_view<IdxT, int64_t, row_major>(
indices.data_handle(), params_.n_queries, params_.k);
auto dist_v = make_device_matrix_view<float, IdxT, row_major>(
auto dist_v = make_device_matrix_view<float, int64_t, row_major>(
distances.data_handle(), params_.n_queries, params_.k);

auto queries_v = make_const_mdspan(queries_.view());
Expand Down Expand Up @@ -125,9 +125,9 @@ struct CagraBench : public fixture {
private:
const params params_;
std::optional<const raft::neighbors::cagra::index<T, IdxT>> index_;
raft::device_matrix<T, IdxT, row_major> queries_;
raft::device_matrix<T, IdxT, row_major> dataset_;
raft::device_matrix<IdxT, IdxT, row_major> knn_graph_;
raft::device_matrix<T, int64_t, row_major> queries_;
raft::device_matrix<T, int64_t, row_major> dataset_;
raft::device_matrix<IdxT, int64_t, row_major> knn_graph_;
};

inline const std::vector<params> generate_inputs()
Expand Down
46 changes: 23 additions & 23 deletions cpp/include/raft/neighbors/cagra.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ namespace raft::neighbors::cagra {
* @endcode
*
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
* @tparam IdxT type of the dataset vector indices
*
* @param[in] res raft resources
* @param[in] dataset a matrix view (host or device) to a row-major matrix [n_rows, dim]
Expand All @@ -76,19 +76,19 @@ namespace raft::neighbors::cagra {
*/
template <typename DataT, typename IdxT, typename accessor>
void build_knn_graph(raft::resources const& res,
mdspan<const DataT, matrix_extent<IdxT>, row_major, accessor> dataset,
raft::host_matrix_view<IdxT, IdxT, row_major> knn_graph,
mdspan<const DataT, matrix_extent<int64_t>, row_major, accessor> dataset,
raft::host_matrix_view<IdxT, int64_t, row_major> knn_graph,
std::optional<float> refine_rate = std::nullopt,
std::optional<ivf_pq::index_params> build_params = std::nullopt,
std::optional<ivf_pq::search_params> search_params = std::nullopt)
{
using internal_IdxT = typename std::make_unsigned<IdxT>::type;

auto knn_graph_internal = make_host_matrix_view<internal_IdxT, internal_IdxT>(
auto knn_graph_internal = make_host_matrix_view<internal_IdxT, int64_t>(
reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),
knn_graph.extent(0),
knn_graph.extent(1));
auto dataset_internal = mdspan<const DataT, matrix_extent<internal_IdxT>, row_major, accessor>(
auto dataset_internal = mdspan<const DataT, matrix_extent<int64_t>, row_major, accessor>(
dataset.data_handle(), dataset.extent(0), dataset.extent(1));

cagra::detail::build_knn_graph(
Expand Down Expand Up @@ -119,7 +119,7 @@ void build_knn_graph(raft::resources const& res,
* @endcode
*
* @tparam DataT type of the data in the source dataset
* @tparam IdxT type of the indices in the source dataset
* @tparam IdxT type of the dataset vector indices
*
* @param[in] res raft resources
* @param[in] dataset a matrix view (host or device) to a row-major matrix [n_rows, dim]
Expand All @@ -133,20 +133,20 @@ template <typename DataT,
typename g_accessor =
host_device_accessor<std::experimental::default_accessor<IdxT>, memory_type::host>>
void sort_knn_graph(raft::resources const& res,
mdspan<const DataT, matrix_extent<IdxT>, row_major, d_accessor> dataset,
mdspan<IdxT, matrix_extent<IdxT>, row_major, g_accessor> knn_graph)
mdspan<const DataT, matrix_extent<int64_t>, row_major, d_accessor> dataset,
mdspan<IdxT, matrix_extent<int64_t>, row_major, g_accessor> knn_graph)
{
using internal_IdxT = typename std::make_unsigned<IdxT>::type;

using g_accessor_internal =
host_device_accessor<std::experimental::default_accessor<internal_IdxT>, g_accessor::mem_type>;
auto knn_graph_internal =
mdspan<internal_IdxT, matrix_extent<internal_IdxT>, row_major, g_accessor_internal>(
mdspan<internal_IdxT, matrix_extent<int64_t>, row_major, g_accessor_internal>(
reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),
knn_graph.extent(0),
knn_graph.extent(1));

auto dataset_internal = mdspan<const DataT, matrix_extent<internal_IdxT>, row_major, d_accessor>(
auto dataset_internal = mdspan<const DataT, matrix_extent<int64_t>, row_major, d_accessor>(
dataset.data_handle(), dataset.extent(0), dataset.extent(1));

cagra::detail::graph::sort_knn_graph(res, dataset_internal, knn_graph_internal);
Expand All @@ -170,20 +170,20 @@ template <typename IdxT = uint32_t,
typename g_accessor =
host_device_accessor<std::experimental::default_accessor<IdxT>, memory_type::host>>
void optimize(raft::resources const& res,
mdspan<IdxT, matrix_extent<IdxT>, row_major, g_accessor> knn_graph,
raft::host_matrix_view<IdxT, IdxT, row_major> new_graph)
mdspan<IdxT, matrix_extent<int64_t>, row_major, g_accessor> knn_graph,
raft::host_matrix_view<IdxT, int64_t, row_major> new_graph)
{
using internal_IdxT = typename std::make_unsigned<IdxT>::type;

auto new_graph_internal = raft::make_host_matrix_view<internal_IdxT, internal_IdxT>(
auto new_graph_internal = raft::make_host_matrix_view<internal_IdxT, int64_t>(
reinterpret_cast<internal_IdxT*>(new_graph.data_handle()),
new_graph.extent(0),
new_graph.extent(1));

using g_accessor_internal =
host_device_accessor<std::experimental::default_accessor<internal_IdxT>, memory_type::host>;
auto knn_graph_internal =
mdspan<internal_IdxT, matrix_extent<internal_IdxT>, row_major, g_accessor_internal>(
mdspan<internal_IdxT, matrix_extent<int64_t>, row_major, g_accessor_internal>(
reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),
knn_graph.extent(0),
knn_graph.extent(1));
Expand Down Expand Up @@ -237,7 +237,7 @@ template <typename T,
host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
index<T, IdxT> build(raft::resources const& res,
const index_params& params,
mdspan<const T, matrix_extent<IdxT>, row_major, Accessor> dataset)
mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
{
size_t intermediate_degree = params.intermediate_graph_degree;
size_t graph_degree = params.graph_degree;
Expand All @@ -256,11 +256,11 @@ index<T, IdxT> build(raft::resources const& res,
graph_degree = intermediate_degree;
}

auto knn_graph = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), intermediate_degree);
auto knn_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), intermediate_degree);

build_knn_graph(res, dataset, knn_graph.view());

auto cagra_graph = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), graph_degree);
auto cagra_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), graph_degree);

optimize<IdxT>(res, knn_graph.view(), cagra_graph.view());

Expand Down Expand Up @@ -289,9 +289,9 @@ template <typename T, typename IdxT>
void search(raft::resources const& res,
const search_params& params,
const index<T, IdxT>& idx,
raft::device_matrix_view<const T, IdxT, row_major> queries,
raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,
raft::device_matrix_view<float, IdxT, row_major> distances)
raft::device_matrix_view<const T, int64_t, row_major> queries,
raft::device_matrix_view<IdxT, int64_t, row_major> neighbors,
raft::device_matrix_view<float, int64_t, row_major> distances)
{
RAFT_EXPECTS(
queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0),
Expand All @@ -303,13 +303,13 @@ void search(raft::resources const& res,
"Number of query dimensions should equal number of dimensions in the index.");

using internal_IdxT = typename std::make_unsigned<IdxT>::type;
auto queries_internal = raft::make_device_matrix_view<const T, internal_IdxT, row_major>(
auto queries_internal = raft::make_device_matrix_view<const T, int64_t, row_major>(
queries.data_handle(), queries.extent(0), queries.extent(1));
auto neighbors_internal = raft::make_device_matrix_view<internal_IdxT, internal_IdxT, row_major>(
auto neighbors_internal = raft::make_device_matrix_view<internal_IdxT, int64_t, row_major>(
reinterpret_cast<internal_IdxT*>(neighbors.data_handle()),
neighbors.extent(0),
neighbors.extent(1));
auto distances_internal = raft::make_device_matrix_view<float, internal_IdxT, row_major>(
auto distances_internal = raft::make_device_matrix_view<float, int64_t, row_major>(
distances.data_handle(), distances.extent(0), distances.extent(1));

cagra::detail::search_main<T, internal_IdxT, IdxT>(
Expand Down
Loading