Skip to content

Commit

Permalink
Merge branch 'branch-24.10' into vamana_build
Browse files Browse the repository at this point in the history
  • Loading branch information
@dantegd
dantegd authored Oct 2, 2024
2 parents 09ded4d + 8a3fac3 commit 0555bb0
Show file tree
Hide file tree
Showing 115 changed files with 10,581 additions and 1,011 deletions.
17 changes: 11 additions & 6 deletions cpp/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,18 +290,25 @@ 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/core/bitset.cu
src/distance/detail/pairwise_matrix/dispatch_canberra_float_float_float_int.cu
src/distance/detail/pairwise_matrix/dispatch_canberra_half_float_float_int.cu
src/distance/detail/pairwise_matrix/dispatch_canberra_double_double_double_int.cu
Expand Down Expand Up @@ -342,6 +349,8 @@ 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
Expand Down Expand Up @@ -398,15 +407,12 @@ add_library(
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 @@ -423,14 +429,13 @@ add_library(
src/neighbors/vamana_serialize_uint8.cu
src/neighbors/vamana_serialize_int8.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
src/stats/trustworthiness_score.cu
)

target_compile_definitions(cuvs PRIVATE "CUVS_EXPLICIT_INSTANTIATE_ONLY")

target_compile_options(
cuvs INTERFACE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--expt-extended-lambda
--expt-relaxed-constexpr>
Expand Down
2 changes: 0 additions & 2 deletions cpp/bench/ann/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,8 +174,6 @@ function(ConfigureAnnBench)
)
endif()

target_compile_definitions(${BENCH_NAME} PRIVATE "CUVS_EXPLICIT_INSTANTIATE_ONLY")

target_include_directories(
${BENCH_NAME}
PUBLIC "$<BUILD_INTERFACE:${CUVS_SOURCE_DIR}/include>"
Expand Down
20 changes: 17 additions & 3 deletions cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h
View file
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
8 changes: 4 additions & 4 deletions cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,7 @@ class BruteForceKNNBenchmark {
search_queries.data(), params_.num_queries, params_.dim),
indices,
distances,
std::nullopt);
cuvs::neighbors::filtering::none_sample_filter{});
flush_l2_cache();
raft::resource::sync_stream(handle_, stream_);
}
Expand All @@ -158,7 +158,7 @@ class BruteForceKNNBenchmark {
search_queries.data(), params_.num_queries, params_.dim),
indices,
distances,
std::nullopt);
cuvs::neighbors::filtering::none_sample_filter{});
raft::resource::sync_stream(handle_, stream_);
end = std::chrono::high_resolution_clock::now();
search_dur = end - start;
Expand All @@ -178,7 +178,7 @@ class BruteForceKNNBenchmark {
search_queries.data(), params_.num_queries, params_.dim),
indices,
distances,
std::nullopt);
cuvs::neighbors::filtering::none_sample_filter{});
flush_l2_cache();
raft::resource::sync_stream(handle_, stream_);
}
Expand All @@ -202,7 +202,7 @@ class BruteForceKNNBenchmark {
search_queries.data(), params_.num_queries, params_.dim),
indices,
distances,
std::nullopt);
cuvs::neighbors::filtering::none_sample_filter{});
raft::resource::sync_stream(handle_, stream_);
end = std::chrono::high_resolution_clock::now();
search_dur = end - start;
Expand Down
8 changes: 6 additions & 2 deletions cpp/bench/ann/src/cuvs/cuvs_wrapper.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,8 +155,12 @@ void cuvs_gpu<T>::search(
raft::make_device_matrix_view<algo_base::index_type, int64_t>(neighbors, batch_size, k);
auto distances_view = raft::make_device_matrix_view<float, int64_t>(distances, batch_size, k);

cuvs::neighbors::brute_force::search(
handle_, *index_, queries_view, neighbors_view, distances_view, std::nullopt);
cuvs::neighbors::brute_force::search(handle_,
*index_,
queries_view,
neighbors_view,
distances_view,
cuvs::neighbors::filtering::none_sample_filter{});
}

template <typename T>
Expand Down
Loading

0 comments on commit 0555bb0

Please sign in to comment.