Skip to content

Commit

Permalink
Merge branch 'doc-2410-index_docs' of github.com:rapidsai/cuvs into d…
Browse files Browse the repository at this point in the history
…oc-2410-index_docs
  • Loading branch information
cjnolet committed Oct 2, 2024
2 parents 397e5ff + 6143c2c commit 831c2a2
Show file tree
Hide file tree
Showing 65 changed files with 4,299 additions and 185 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -81,4 +81,4 @@ ivf_pq_index

# cuvs_bench
datasets/
/*.json
/*.json
21 changes: 17 additions & 4 deletions cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,7 @@ endif()
add_library(
cuvs-cagra-search STATIC
src/neighbors/cagra_search_float.cu
src/neighbors/cagra_search_half.cu
src/neighbors/cagra_search_int8.cu
src/neighbors/cagra_search_uint8.cu
src/neighbors/detail/cagra/compute_distance.cu
Expand Down Expand Up @@ -257,14 +258,10 @@ add_library(
src/neighbors/detail/cagra/search_multi_cta_half_uint32.cu
src/neighbors/detail/cagra/search_multi_cta_int8_uint32.cu
src/neighbors/detail/cagra/search_multi_cta_uint8_uint32.cu
src/neighbors/detail/cagra/search_multi_cta_float_uint64.cu
src/neighbors/detail/cagra/search_multi_cta_half_uint64.cu
src/neighbors/detail/cagra/search_single_cta_float_uint32.cu
src/neighbors/detail/cagra/search_single_cta_half_uint32.cu
src/neighbors/detail/cagra/search_single_cta_int8_uint32.cu
src/neighbors/detail/cagra/search_single_cta_uint8_uint32.cu
src/neighbors/detail/cagra/search_single_cta_float_uint64.cu
src/neighbors/detail/cagra/search_single_cta_half_uint64.cu
)

file(GLOB_RECURSE compute_distance_sources "src/neighbors/detail/cagra/compute_distance_*.cu")
Expand Down Expand Up @@ -293,16 +290,22 @@ target_compile_options(
add_library(
cuvs SHARED
src/cluster/kmeans_balanced_fit_float.cu
src/cluster/kmeans_fit_mg_float.cu
src/cluster/kmeans_fit_mg_double.cu
src/cluster/kmeans_fit_double.cu
src/cluster/kmeans_fit_float.cu
src/cluster/kmeans_auto_find_k_float.cu
src/cluster/kmeans_fit_predict_double.cu
src/cluster/kmeans_fit_predict_float.cu
src/cluster/kmeans_predict_double.cu
src/cluster/kmeans_predict_float.cu
src/cluster/kmeans_balanced_fit_float.cu
src/cluster/kmeans_balanced_fit_predict_float.cu
src/cluster/kmeans_balanced_predict_float.cu
src/cluster/kmeans_balanced_fit_int8.cu
src/cluster/kmeans_balanced_fit_predict_int8.cu
src/cluster/kmeans_balanced_predict_int8.cu
src/cluster/kmeans_transform_double.cu
src/cluster/kmeans_transform_float.cu
src/cluster/single_linkage_float.cu
src/distance/detail/pairwise_matrix/dispatch_canberra_float_float_float_int.cu
Expand Down Expand Up @@ -345,18 +348,22 @@ add_library(
src/distance/detail/pairwise_matrix/dispatch_russel_rao_half_float_float_int.cu
src/distance/detail/pairwise_matrix/dispatch_russel_rao_double_double_double_int.cu
src/distance/detail/pairwise_matrix/dispatch_rbf.cu
src/distance/detail/pairwise_matrix/dispatch_l2_expanded_double_double_double_int64_t.cu
src/distance/detail/pairwise_matrix/dispatch_l2_expanded_float_float_float_int64_t.cu
src/distance/detail/fused_distance_nn.cu
src/distance/distance.cu
src/distance/pairwise_distance.cu
src/neighbors/brute_force.cu
src/neighbors/cagra_build_float.cu
src/neighbors/cagra_build_half.cu
src/neighbors/cagra_build_int8.cu
src/neighbors/cagra_build_uint8.cu
src/neighbors/cagra_extend_float.cu
src/neighbors/cagra_extend_int8.cu
src/neighbors/cagra_extend_uint8.cu
src/neighbors/cagra_optimize.cu
src/neighbors/cagra_serialize_float.cu
src/neighbors/cagra_serialize_half.cu
src/neighbors/cagra_serialize_int8.cu
src/neighbors/cagra_serialize_uint8.cu
src/neighbors/detail/cagra/cagra_build.cpp
Expand All @@ -378,6 +385,7 @@ add_library(
src/neighbors/ivf_pq/ivf_pq_serialize.cu
src/neighbors/ivf_pq/ivf_pq_deserialize.cu
src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_build_extend_half_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false.cu
Expand All @@ -395,15 +403,19 @@ add_library(
src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu
src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_float_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_half_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_int8_t_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_uint8_t_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_float_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_half_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_int8_t_int64_t.cu
src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_uint8_t_int64_t.cu
src/neighbors/nn_descent.cu
src/neighbors/nn_descent_float.cu
src/neighbors/nn_descent_half.cu
src/neighbors/nn_descent_int8.cu
src/neighbors/nn_descent_uint8.cu
src/neighbors/reachability.cu
src/neighbors/refine/detail/refine_device_float_float.cu
src/neighbors/refine/detail/refine_device_half_float.cu
src/neighbors/refine/detail/refine_device_int8_t_float.cu
Expand All @@ -414,6 +426,7 @@ add_library(
src/neighbors/refine/detail/refine_host_uint8_t_float.cpp
src/neighbors/sample_filter.cu
src/selection/select_k_float_int64_t.cu
src/selection/select_k_float_int32_t.cu
src/selection/select_k_float_uint32_t.cu
src/selection/select_k_half_uint32_t.cu
src/stats/silhouette_score.cu
Expand Down
22 changes: 11 additions & 11 deletions cpp/bench/ann/src/common/benchmark.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -687,17 +687,17 @@ inline auto run_main(int argc, char** argv) -> int
override_kv,
metric_objective,
threads);
// } else if (dtype == "half") {
// dispatch_benchmark<half>(cmdline
// conf,
// force_overwrite,
// build_mode,
// search_mode,
// data_prefix,
// index_prefix,
// override_kv,
// metric_objective,
// threads);
} else if (dtype == "half") {
dispatch_benchmark<half>(cmdline,
conf,
force_overwrite,
build_mode,
search_mode,
data_prefix,
index_prefix,
override_kv,
metric_objective,
threads);
} else if (dtype == "uint8") {
dispatch_benchmark<std::uint8_t>(cmdline,
conf,
Expand Down
20 changes: 17 additions & 3 deletions cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
#include <rmm/device_uvector.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/failure_callback_resource_adaptor.hpp>
#include <rmm/mr/device/managed_memory_resource.hpp>
#include <rmm/mr/device/pool_memory_resource.hpp>

#include <memory>
Expand Down Expand Up @@ -75,13 +76,14 @@ inline auto rmm_oom_callback(std::size_t bytes, void*) -> bool
*/
class shared_raft_resources {
public:
using pool_mr_type = rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>;
using mr_type = rmm::mr::failure_callback_resource_adaptor<pool_mr_type>;
using pool_mr_type = rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>;
using mr_type = rmm::mr::failure_callback_resource_adaptor<pool_mr_type>;
using large_mr_type = rmm::mr::managed_memory_resource;

shared_raft_resources()
try : orig_resource_{rmm::mr::get_current_device_resource()},
pool_resource_(orig_resource_, 1024 * 1024 * 1024ull),
resource_(&pool_resource_, rmm_oom_callback, nullptr) {
resource_(&pool_resource_, rmm_oom_callback, nullptr), large_mr_() {
rmm::mr::set_current_device_resource(&resource_);
} catch (const std::exception& e) {
auto cuda_status = cudaGetLastError();
Expand All @@ -104,10 +106,16 @@ class shared_raft_resources {

~shared_raft_resources() noexcept { rmm::mr::set_current_device_resource(orig_resource_); }

auto get_large_memory_resource() noexcept
{
return static_cast<rmm::mr::device_memory_resource*>(&large_mr_);
}

private:
rmm::mr::device_memory_resource* orig_resource_;
pool_mr_type pool_resource_;
mr_type resource_;
large_mr_type large_mr_;
};

/**
Expand All @@ -130,6 +138,12 @@ class configured_raft_resources {
res_{std::make_unique<raft::device_resources>(
rmm::cuda_stream_view(get_stream_from_global_pool()))}
{
// set the large workspace resource to the raft handle, but without the deleter
// (this resource is managed by the shared_res).
raft::resource::set_large_workspace_resource(
*res_,
std::shared_ptr<rmm::mr::device_memory_resource>(shared_res_->get_large_memory_resource(),
raft::void_op{}));
}

/** Default constructor creates all resources anew. */
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/ann/src/cuvs/cuvs_benchmark.cu
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con
}; // namespace cuvs::bench

REGISTER_ALGO_INSTANCE(float);
// REGISTER_ALGO_INSTANCE(half);
REGISTER_ALGO_INSTANCE(half);
REGISTER_ALGO_INSTANCE(std::int8_t);
REGISTER_ALGO_INSTANCE(std::uint8_t);

Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/ann/src/cuvs/cuvs_cagra_half.cu
Original file line number Diff line number Diff line change
Expand Up @@ -16,5 +16,5 @@
#include "cuvs_cagra_wrapper.h"

namespace cuvs::bench {
// template class cuvs_cagra<half, uint32_t>;
template class cuvs_cagra<half, uint32_t>;
} // namespace cuvs::bench
6 changes: 5 additions & 1 deletion cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,11 @@ void cuvs_cagra<T, IdxT>::save(const std::string& file) const
template <typename T, typename IdxT>
void cuvs_cagra<T, IdxT>::save_to_hnswlib(const std::string& file) const
{
cuvs::neighbors::cagra::serialize_to_hnswlib(handle_, file, *index_);
if constexpr (!std::is_same_v<T, half>) {
cuvs::neighbors::cagra::serialize_to_hnswlib(handle_, file, *index_);
} else {
RAFT_FAIL("Cannot save fp16 index to hnswlib format");
}
}

template <typename T, typename IdxT>
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/ann/src/cuvs/cuvs_ivf_pq.cu
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

namespace cuvs::bench {
template class cuvs_ivf_pq<float, int64_t>;
// template class cuvs_ivf_pq<half, int64_t>;
template class cuvs_ivf_pq<half, int64_t>;
template class cuvs_ivf_pq<uint8_t, int64_t>;
template class cuvs_ivf_pq<int8_t, int64_t>;
} // namespace cuvs::bench
Loading

0 comments on commit 831c2a2

Please sign in to comment.