Skip to content

Commit

Permalink
addressing review
Browse files Browse the repository at this point in the history
  • Loading branch information
@viclafargue
viclafargue committed Feb 9, 2023
1 parent 126677c commit cbb2d2e
Show file tree
Hide file tree
Showing 2 changed files with 137 additions and 174 deletions.
275 changes: 137 additions & 138 deletions cpp/include/raft/neighbors/ivf_pq.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,143 @@ namespace raft::neighbors::ivf_pq {
* @{
*/

/**
* @brief Build the index from the dataset for efficient search.
*
* NB: Currently, the following distance metrics are supported:
* - L2Expanded
* - L2Unexpanded
* - InnerProduct
*
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
*
* @param handle
* @param params configure the index building
* @param[in] dataset a device matrix view to a row-major matrix [n_rows, dim]
*
* @return the constructed ivf-pq index
*/
template <typename T, typename IdxT = uint32_t>
inline auto build(raft::device_resources const& handle,
const index_params& params,
const raft::device_matrix_view<const T, IdxT, row_major>& dataset) -> index<IdxT>
{
IdxT n_rows = dataset.extent(0);
IdxT dim = dataset.extent(1);
return raft::spatial::knn::ivf_pq::detail::build(
handle, params, dataset.data_handle(), n_rows, dim);
}

/**
* @brief Build a new index containing the data of the original plus new extra vectors.
*
* Implementation note:
* The new data is clustered according to existing kmeans clusters, then the cluster
* centers are unchanged.
*
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
*
* @param handle
* @param orig_index original index
* @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()]
* @param[in] new_indices a device matrix view to a vector of indices [n_rows].
* If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr`
* here to imply a continuous range `[0...n_rows)`.
*
* @return the constructed extended ivf-pq index
*/
template <typename T, typename IdxT>
inline auto extend(raft::device_resources const& handle,
const index<IdxT>& orig_index,
const raft::device_matrix_view<const T, IdxT, row_major>& new_vectors,
const raft::device_matrix_view<const IdxT, IdxT, row_major>& new_indices)
-> index<IdxT>
{
IdxT n_rows = new_vectors.extent(0);
ASSERT(n_rows == new_indices.extent(0),
"new_vectors and new_indices have different number of rows");
ASSERT(new_vectors.extent(1) == orig_index.dim(),
"new_vectors should have the same dimension as the index");
return raft::spatial::knn::ivf_pq::detail::extend(
handle, orig_index, new_vectors.data_handle(), new_indices.data_handle(), n_rows);
}

/**
* @brief Extend the index with the new data.
* *
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
*
* @param handle
* @param[inout] index
* @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()]
* @param[in] new_indices a device matrix view to a vector of indices [n_rows].
* If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr`
* here to imply a continuous range `[0...n_rows)`.
*/
template <typename T, typename IdxT>
inline void extend(raft::device_resources const& handle,
index<IdxT>* index,
const raft::device_matrix_view<const T, IdxT, row_major>& new_vectors,
const raft::device_matrix_view<const IdxT, IdxT, row_major>& new_indices)
{
*index = extend(handle, *index, new_vectors, new_indices);
}

/**
* @brief Search ANN using the constructed index.
*
* See the [ivf_pq::build](#ivf_pq::build) documentation for a usage example.
*
* Note, this function requires a temporary buffer to store intermediate results between cuda kernel
* calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can
* pass a pool memory resource or a large enough pre-allocated memory resource to reduce or
* eliminate entirely allocations happening within `search`.
* The exact size of the temporary buffer depends on multiple factors and is an implementation
* detail. However, you can safely specify a small initial size for the memory pool, so that only a
* few allocations happen to grow it during the first invocations of the `search`.
*
* @tparam T data element type
* @tparam IdxT type of the indices
*
* @param handle
* @param params configure the search
* @param index ivf-pq constructed index
* @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()]
* @param k the number of neighbors to find for each query.
* @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset
* [n_queries, k]
* @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
* k]
*/
template <typename T, typename IdxT>
inline void search(raft::device_resources const& handle,
const search_params& params,
const index<IdxT>& index,
const raft::device_matrix_view<const T, IdxT, row_major>& queries,
uint32_t k,
const raft::device_matrix_view<IdxT, IdxT, row_major>& neighbors,
const raft::device_matrix_view<float, IdxT, row_major>& distances)
{
IdxT n_queries = queries.extent(0);
bool check_n_rows = (n_queries == neighbors.extent(0)) && (n_queries == distances.extent(0));
ASSERT(check_n_rows,
"queries, neighbors and distances parameters have inconsistent number of rows");
return raft::spatial::knn::ivf_pq::detail::search(handle,
params,
index,
queries.data_handle(),
n_queries,
k,
neighbors.data_handle(),
distances.data_handle(),
handle.get_workspace_resource());
}

/** @} */ // end group ivf_pq

/**
* @brief Build the index from the dataset for efficient search.
*
Expand Down Expand Up @@ -74,34 +211,6 @@ auto build(raft::device_resources const& handle,
return raft::spatial::knn::ivf_pq::detail::build(handle, params, dataset, n_rows, dim);
}

/**
* @brief Build the index from the dataset for efficient search.
*
* NB: Currently, the following distance metrics are supported:
* - L2Expanded
* - L2Unexpanded
* - InnerProduct
*
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
*
* @param handle
* @param params configure the index building
* @param[in] dataset a device matrix view to a row-major matrix [n_rows, dim]
*
* @return the constructed ivf-pq index
*/
template <typename T, typename IdxT = uint32_t>
inline auto build(raft::device_resources const& handle,
const index_params& params,
const raft::device_matrix_view<const T, IdxT, row_major>& dataset) -> index<IdxT>
{
IdxT n_rows = dataset.extent(0);
IdxT dim = dataset.extent(1);
return raft::spatial::knn::ivf_pq::detail::build(
handle, params, dataset.data_handle(), n_rows, dim);
}

/**
* @brief Build a new index containing the data of the original plus new extra vectors.
*
Expand Down Expand Up @@ -145,39 +254,6 @@ auto extend(raft::device_resources const& handle,
handle, orig_index, new_vectors, new_indices, n_rows);
}

/**
* @brief Build a new index containing the data of the original plus new extra vectors.
*
* Implementation note:
* The new data is clustered according to existing kmeans clusters, then the cluster
* centers are unchanged.
*
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
*
* @param handle
* @param orig_index original index
* @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()]
* @param[in] new_indices a device matrix view to a vector of indices [n_rows].
* If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr`
* here to imply a continuous range `[0...n_rows)`.
*
* @return the constructed extended ivf-pq index
*/
template <typename T, typename IdxT>
inline auto extend(raft::device_resources const& handle,
const index<IdxT>& orig_index,
const raft::device_matrix_view<const T, IdxT, row_major>& new_vectors,
const raft::device_matrix_view<const IdxT, IdxT, row_major>& new_indices)
-> index<IdxT>
{
IdxT n_rows = new_vectors.extent(0);
ASSERT(n_rows == new_indices.extent(0),
"new_vectors and new_indices have different number of rows");
return raft::spatial::knn::ivf_pq::detail::extend(
handle, orig_index, new_vectors.data_handle(), new_indices.data_handle(), n_rows);
}

/**
* @brief Extend the index with the new data.
* *
Expand All @@ -202,28 +278,6 @@ void extend(raft::device_resources const& handle,
*index = extend(handle, *index, new_vectors, new_indices, n_rows);
}

/**
* @brief Extend the index with the new data.
* *
* @tparam T data element type
* @tparam IdxT type of the indices in the source dataset
*
* @param handle
* @param[inout] index
* @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()]
* @param[in] new_indices a device matrix view to a vector of indices [n_rows].
* If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr`
* here to imply a continuous range `[0...n_rows)`.
*/
template <typename T, typename IdxT>
inline void extend(raft::device_resources const& handle,
index<IdxT>* index,
const raft::device_matrix_view<const T, IdxT, row_major>& new_vectors,
const raft::device_matrix_view<const IdxT, IdxT, row_major>& new_indices)
{
*index = extend(handle, *index, new_vectors, new_indices);
}

/**
* @brief Search ANN using the constructed index.
*
Expand Down Expand Up @@ -281,59 +335,4 @@ void search(raft::device_resources const& handle,
handle, params, index, queries, n_queries, k, neighbors, distances, mr);
}

/**
* @brief Search ANN using the constructed index.
*
* See the [ivf_pq::build](#ivf_pq::build) documentation for a usage example.
*
* Note, this function requires a temporary buffer to store intermediate results between cuda kernel
* calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can
* pass a pool memory resource or a large enough pre-allocated memory resource to reduce or
* eliminate entirely allocations happening within `search`.
* The exact size of the temporary buffer depends on multiple factors and is an implementation
* detail. However, you can safely specify a small initial size for the memory pool, so that only a
* few allocations happen to grow it during the first invocations of the `search`.
*
* @tparam T data element type
* @tparam IdxT type of the indices
*
* @param handle
* @param params configure the search
* @param index ivf-pq constructed index
* @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()]
* @param k the number of neighbors to find for each query.
* @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset
* [n_queries, k]
* @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
* k]
* @param mr an optional memory resource to use across the searches (you can provide a large enough
* memory pool here to avoid memory allocations within search).
*/
template <typename T, typename IdxT>
inline void search(raft::device_resources const& handle,
const search_params& params,
const index<IdxT>& index,
const raft::device_matrix_view<const T, IdxT, row_major>& queries,
uint32_t k,
const raft::device_matrix_view<IdxT, IdxT, row_major>& neighbors,
const raft::device_matrix_view<float, IdxT, row_major>& distances,
rmm::mr::device_memory_resource* mr = nullptr)
{
IdxT n_queries = queries.extent(0);
bool check_n_rows = (n_queries == neighbors.extent(0)) && (n_queries == distances.extent(0));
ASSERT(check_n_rows,
"queries, neighbors and distances parameters have inconsistent number of rows");
return raft::spatial::knn::ivf_pq::detail::search(handle,
params,
index,
queries.data_handle(),
n_queries,
k,
neighbors.data_handle(),
distances.data_handle(),
mr);
}

/** @} */ // end group ivf_pq

} // namespace raft::neighbors::ivf_pq
36 changes: 0 additions & 36 deletions cpp/include/raft_runtime/neighbors/ivf_pq.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,6 @@
namespace raft::runtime::neighbors::ivf_pq {

#define RAFT_INST_SEARCH(T, IdxT) \
void search(raft::device_resources const&, \
const raft::neighbors::ivf_pq::search_params&, \
const raft::neighbors::ivf_pq::index<IdxT>&, \
const T*, \
uint32_t, \
uint32_t, \
IdxT*, \
float*, \
rmm::mr::device_memory_resource*); \
\
void search(raft::device_resources const&, \
const raft::neighbors::ivf_pq::search_params&, \
const raft::neighbors::ivf_pq::index<IdxT>&, \
Expand All @@ -50,32 +40,6 @@ RAFT_INST_SEARCH(uint8_t, uint64_t);
// wrappers, where exception handling is not compatible with return type that has nontrivial
// constructor.
#define RAFT_INST_BUILD_EXTEND(T, IdxT) \
auto build(raft::device_resources const& handle, \
const raft::neighbors::ivf_pq::index_params& params, \
const T* dataset, \
IdxT n_rows, \
uint32_t dim) \
->raft::neighbors::ivf_pq::index<IdxT>; \
\
auto extend(raft::device_resources const& handle, \
const raft::neighbors::ivf_pq::index<IdxT>& orig_index, \
const T* new_vectors, \
const IdxT* new_indices, \
IdxT n_rows) \
->raft::neighbors::ivf_pq::index<IdxT>; \
\
void build(raft::device_resources const& handle, \
const raft::neighbors::ivf_pq::index_params& params, \
const T* dataset, \
IdxT n_rows, \
uint32_t dim, \
raft::neighbors::ivf_pq::index<IdxT>* idx); \
\
void extend(raft::device_resources const& handle, \
raft::neighbors::ivf_pq::index<IdxT>* idx, \
const T* new_vectors, \
const IdxT* new_indices, \
IdxT n_rows); \
auto build(raft::device_resources const& handle, \
const raft::neighbors::ivf_pq::index_params& params, \
const raft::device_matrix_view<const T, IdxT, row_major>& dataset) \
Expand Down

0 comments on commit cbb2d2e

Please sign in to comment.