API for CPU side output indices and distances of all_neighbors
#1905
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -80,33 +80,36 @@ static cuvs::neighbors::all_neighbors::all_neighbors_params convert_params( | |
| return out; | ||
| } | ||
|
|
||
| static void ensure_indices_dtype_compatibility(DLManagedTensor* indices) | ||
| { | ||
| auto dtype = indices->dl_tensor.dtype; | ||
| RAFT_EXPECTS(dtype.code == kDLInt && dtype.bits == 64, "indices must be int64 output tensor"); | ||
| RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(indices->dl_tensor) || | ||
| cuvs::core::is_dlpack_host_compatible(indices->dl_tensor), | ||
| "indices tensor must be either device-compatible or host-compatible"); | ||
| } | ||
|
|
||
| static void ensure_optional_distance_dtype_compatibility(DLManagedTensor* distances) | ||
| { | ||
| if (distances == nullptr) { return; } | ||
| auto dtype = distances->dl_tensor.dtype; | ||
| RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, | ||
| "distances must be float32 output tensor"); | ||
| RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances->dl_tensor) || | ||
| cuvs::core::is_dlpack_host_compatible(distances->dl_tensor), | ||
| "distances tensor must be either device-compatible or host-compatible"); | ||
| } | ||
|
|
||
| static void ensure_optional_core_distance_dtype_compatibility(DLManagedTensor* core_distances) | ||
| { | ||
| if (core_distances == nullptr) { return; } | ||
| auto dtype = core_distances->dl_tensor.dtype; | ||
| RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, | ||
| "core_distances must be float32 output tensor"); | ||
| RAFT_EXPECTS( | ||
| cuvs::core::is_dlpack_device_compatible(core_distances->dl_tensor) || | ||
| cuvs::core::is_dlpack_host_compatible(core_distances->dl_tensor), | ||
| "core_distances tensor must be either device-compatible or host-compatible"); | ||
| } | ||
|
|
||
| template <typename T> | ||
|
|
@@ -124,9 +127,9 @@ void _build_host(cuvsResources_t res, | |
| RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dlt), | ||
| "Host build expects host-compatible dataset tensor"); | ||
|
|
||
| ensure_indices_dtype_compatibility(indices_tensor); | ||
| ensure_optional_distance_dtype_compatibility(distances_tensor); | ||
| ensure_optional_core_distance_dtype_compatibility(core_distances_tensor); | ||
|
|
||
| // Check dependencies between parameters | ||
| if (core_distances_tensor != nullptr && distances_tensor == nullptr) { | ||
|
|
@@ -138,26 +141,63 @@ void _build_host(cuvsResources_t res, | |
|
|
||
| auto cpp_params = convert_params(params, n_rows, n_cols); | ||
|
|
||
| using dataset_mdspan_t = raft::host_matrix_view<const T, int64_t, raft::row_major>; | ||
|
|
||
| bool indices_is_host = cuvs::core::is_dlpack_host_compatible(indices_tensor->dl_tensor); | ||
| bool distances_is_host = distances_tensor ? cuvs::core::is_dlpack_host_compatible(distances_tensor->dl_tensor) : indices_is_host; | ||
|
|
||
| if (distances_tensor && distances_is_host != indices_is_host) { | ||
| RAFT_FAIL("distances and indices must be on the same memory location (both host or both device)"); | ||
| } | ||
|
|
||
| if (core_distances_tensor) { | ||
| bool core_distances_is_host = | ||
| cuvs::core::is_dlpack_host_compatible(core_distances_tensor->dl_tensor); | ||
| RAFT_EXPECTS(core_distances_is_host == indices_is_host, | ||
| "core_distances must be on the same memory location as indices and distances"); | ||
| } | ||
|
|
||
| auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor); | ||
|
|
||
| if (indices_is_host) { | ||
|
Contributor
There was a problem hiding this comment. The if (indices_is_host) { … } else { … } blocks here are essentially the same 30-ish lines duplicated four times across _build_host and _build_device (host branch and device branch in each). Would you be open to extracting a small templated helper that takes the dataset mdspan type + a host/device tag and does the optional unpacking + dispatch? Should cut this down a lot, or is this cod that'll significantly change in teh follow up and not worth the change? |
||
| using indices_mdspan_t = raft::host_matrix_view<int64_t, int64_t, raft::row_major>; | ||
| using distances_mdspan_t = raft::host_matrix_view<float, int64_t, raft::row_major>; | ||
| using core_mdspan_t = raft::host_vector_view<float, int64_t>; | ||
|
|
||
| auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor); | ||
|
|
||
| std::optional<distances_mdspan_t> distances = std::nullopt; | ||
| if (distances_tensor) { | ||
| distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor); | ||
| } | ||
|
|
||
| std::optional<core_mdspan_t> core_distances = std::nullopt; | ||
| if (core_distances_tensor) { | ||
| core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor); | ||
| } | ||
|
|
||
| cuvs::neighbors::all_neighbors::build( | ||
| cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); | ||
| } else { | ||
| using indices_mdspan_t = raft::device_matrix_view<int64_t, int64_t, raft::row_major>; | ||
| using distances_mdspan_t = raft::device_matrix_view<float, int64_t, raft::row_major>; | ||
| using core_mdspan_t = raft::device_vector_view<float, int64_t>; | ||
|
|
||
| auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor); | ||
|
|
||
| std::optional<distances_mdspan_t> distances = std::nullopt; | ||
| if (distances_tensor) { | ||
| distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor); | ||
| } | ||
|
|
||
| std::optional<core_mdspan_t> core_distances = std::nullopt; | ||
| if (core_distances_tensor) { | ||
| core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor); | ||
| } | ||
|
|
||
| cuvs::neighbors::all_neighbors::build( | ||
| cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); | ||
| } | ||
| } | ||
|
|
||
| template <typename T> | ||
|
|
@@ -175,9 +215,9 @@ void _build_device(cuvsResources_t device_res, | |
| RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(dlt), | ||
| "Device build expects device-compatible dataset tensor"); | ||
|
|
||
| ensure_indices_dtype_compatibility(indices_tensor); | ||
| ensure_optional_distance_dtype_compatibility(distances_tensor); | ||
| ensure_optional_core_distance_dtype_compatibility(core_distances_tensor); | ||
|
|
||
| // Check dependencies between parameters | ||
| if (core_distances_tensor != nullptr && distances_tensor == nullptr) { | ||
|
|
@@ -189,26 +229,62 @@ void _build_device(cuvsResources_t device_res, | |
|
|
||
| auto cpp_params = convert_params(params, n_rows, n_cols); | ||
|
|
||
| using dataset_mdspan_t = raft::device_matrix_view<const T, int64_t, raft::row_major>; | ||
| auto dataset = cuvs::core::from_dlpack<dataset_mdspan_t>(dataset_tensor); | ||
|
|
||
| bool indices_is_host = cuvs::core::is_dlpack_host_compatible(indices_tensor->dl_tensor); | ||
| bool distances_is_host = distances_tensor ? cuvs::core::is_dlpack_host_compatible(distances_tensor->dl_tensor) : indices_is_host; | ||
|
|
||
| if (distances_tensor && distances_is_host != indices_is_host) { | ||
| RAFT_FAIL("distances and indices must be on the same memory location (both host or both device)"); | ||
| } | ||
|
|
||
| if (core_distances_tensor) { | ||
| bool core_distances_is_host = | ||
| cuvs::core::is_dlpack_host_compatible(core_distances_tensor->dl_tensor); | ||
| RAFT_EXPECTS(core_distances_is_host == indices_is_host, | ||
| "core_distances must be on the same memory location as indices and distances"); | ||
| } | ||
|
|
||
| if (indices_is_host) { | ||
| using indices_mdspan_t = raft::host_matrix_view<int64_t, int64_t, raft::row_major>; | ||
| using distances_mdspan_t = raft::host_matrix_view<float, int64_t, raft::row_major>; | ||
| using core_mdspan_t = raft::host_vector_view<float, int64_t>; | ||
|
|
||
| auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor); | ||
|
|
||
| std::optional<distances_mdspan_t> distances = std::nullopt; | ||
| if (distances_tensor) { | ||
| distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor); | ||
| } | ||
|
|
||
| std::optional<core_mdspan_t> core_distances = std::nullopt; | ||
| if (core_distances_tensor) { | ||
| core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor); | ||
| } | ||
|
|
||
| cuvs::neighbors::all_neighbors::build( | ||
| cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); | ||
| } else { | ||
| using indices_mdspan_t = raft::device_matrix_view<int64_t, int64_t, raft::row_major>; | ||
| using distances_mdspan_t = raft::device_matrix_view<float, int64_t, raft::row_major>; | ||
| using core_mdspan_t = raft::device_vector_view<float, int64_t>; | ||
|
|
||
| auto indices = cuvs::core::from_dlpack<indices_mdspan_t>(indices_tensor); | ||
|
|
||
| std::optional<distances_mdspan_t> distances = std::nullopt; | ||
| if (distances_tensor) { | ||
| distances = cuvs::core::from_dlpack<distances_mdspan_t>(distances_tensor); | ||
| } | ||
|
|
||
| std::optional<core_mdspan_t> core_distances = std::nullopt; | ||
| if (core_distances_tensor) { | ||
| core_distances = cuvs::core::from_dlpack<core_mdspan_t>(core_distances_tensor); | ||
| } | ||
|
|
||
| cuvs::neighbors::all_neighbors::build( | ||
| cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); | ||
| } | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -119,11 +119,11 @@ struct all_neighbors_params { | |
| * to build all-neighbors knn graph | ||
| * @param[in] dataset raft::host_matrix_view input dataset expected to be located | ||
| * in host memory | ||
| * @param[out] indices nearest neighbor indices of shape [n_row x k] in device memory | ||
| * @param[out] distances nearest neighbor distances [n_row x k] in device memory | ||
| * @param[out] core_distances core distances of size [n_row] in device memory. Requires distances | ||
| * matrix to compute core_distances. If core_distances is given, the resulting indices and distances | ||
| * will be in mutual reachability space. | ||
| * @param[in] alpha distance scaling parameter as used in robust single linkage. | ||
| */ | ||
| void build( | ||
|
|
@@ -155,11 +155,11 @@ void build( | |
| * to build all-neighbors knn graph | ||
| * @param[in] dataset raft::device_matrix_view input dataset expected to be located | ||
| * in device memory | ||
| * @param[out] indices nearest neighbor indices of shape [n_row x k] in device memory | ||
| * @param[out] distances nearest neighbor distances [n_row x k] in device memory | ||
| * @param[out] core_distances core distances of size [n_row] in device memory. Requires distances | ||
| * matrix to compute core_distances. If core_distances is given, the resulting indices and distances | ||
| * will be in mutual reachability space. | ||
| * @param[in] alpha distance scaling parameter as used in robust single linkage. | ||
| */ | ||
| void build( | ||
|
|
@@ -171,6 +171,77 @@ void build( | |
| std::optional<raft::device_vector_view<float, int64_t, row_major>> core_distances = std::nullopt, | ||
| float alpha = 1.0); | ||
|
|
||
| /** | ||
| * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the training | ||
| * vectors) with host memory output buffers. | ||
| * | ||
| * Usage example: | ||
| * @code{.cpp} | ||
| * using namespace cuvs::neighbors; | ||
| * // use default index parameters | ||
| * all_neighbors::all_neighbors_params params; | ||
| * auto indices = raft::make_host_matrix<int64_t, int64_t>(handle, n_row, k); | ||
| * auto distances = raft::make_host_matrix<float, int64_t>(handle, n_row, k); | ||
| * all_neighbors::build(res, params, dataset, indices.view(), distances.view()); | ||
| * @endcode | ||
| * | ||
| * @param[in] handle raft::resources is an object managing resources | ||
| * @param[in] params an instance of all_neighbors::all_neighbors_params that are parameters | ||
| * to build all-neighbors knn graph | ||
| * @param[in] dataset raft::host_matrix_view input dataset expected to be located | ||
| * in host memory | ||
| * @param[out] indices nearest neighbor indices of shape [n_row x k] in host memory | ||
| * @param[out] distances nearest neighbor distances [n_row x k] in host memory | ||
| * @param[out] core_distances core distances of size [n_row] in host memory. Requires distances | ||
| * matrix to compute core_distances. If core_distances is given, the resulting indices and distances | ||
| * will be in mutual reachability space. | ||
| * @param[in] alpha distance scaling parameter as used in robust single linkage. | ||
| */ | ||
| void build( | ||
| const raft::resources& handle, | ||
| const all_neighbors_params& params, | ||
| raft::host_matrix_view<const float, int64_t, row_major> dataset, | ||
| raft::host_matrix_view<int64_t, int64_t, row_major> indices, | ||
| std::optional<raft::host_matrix_view<float, int64_t, row_major>> distances = std::nullopt, | ||
| std::optional<raft::host_vector_view<float, int64_t, row_major>> core_distances = std::nullopt, | ||
| float alpha = 1.0); | ||
|
|
||
| /** | ||
| * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the training | ||
| * vectors) with host memory output buffers. params.n_clusters should be 1 for data on device. To | ||
| * use a larger params.n_clusters for efficient device memory usage, put data on host RAM. | ||
| * | ||
| * Usage example: | ||
| * @code{.cpp} | ||
| * using namespace cuvs::neighbors; | ||
| * // use default index parameters | ||
| * all_neighbors::all_neighbors_params params; | ||
| * auto indices = raft::make_host_matrix<int64_t, int64_t>(handle, n_row, k); | ||
| * auto distances = raft::make_host_matrix<float, int64_t>(handle, n_row, k); | ||
| * all_neighbors::build(res, params, dataset, indices.view(), distances.view()); | ||
| * @endcode | ||
| * | ||
| * @param[in] handle raft::resources is an object managing resources | ||
| * @param[in] params an instance of all_neighbors::all_neighbors_params that are parameters | ||
| * to build all-neighbors knn graph | ||
| * @param[in] dataset raft::device_matrix_view input dataset expected to be located | ||
| * in device memory | ||
| * @param[out] indices nearest neighbor indices of shape [n_row x k] in host memory | ||
| * @param[out] distances nearest neighbor distances [n_row x k] in host memory | ||
| * @param[out] core_distances core distances of size [n_row] in host memory. Requires distances | ||
| * matrix to compute core_distances. If core_distances is given, the resulting indices and distances | ||
| * will be in mutual reachability space. | ||
| * @param[in] alpha distance scaling parameter as used in robust single linkage. | ||
| */ | ||
| void build( | ||
| const raft::resources& handle, | ||
| const all_neighbors_params& params, | ||
| raft::device_matrix_view<const float, int64_t, row_major> dataset, | ||
| raft::host_matrix_view<int64_t, int64_t, row_major> indices, | ||
| std::optional<raft::host_matrix_view<float, int64_t, row_major>> distances = std::nullopt, | ||
| std::optional<raft::host_vector_view<float, int64_t, row_major>> core_distances = std::nullopt, | ||
| float alpha = 1.0); | ||
|
|
||
| /** @} */ | ||
| } // namespace all_neighbors | ||
| } // namespace neighbors | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's kinda odd to see Python terminology in the C header, I guess one of the main consumers are the python apis, but still... not necessarily suggesting to change it, but was wondering ig it might be better to keep the docstring to C relevant concepts (like dlpack)