diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 90d23d9aef..5a1c870691 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -231,6 +231,55 @@ function(ConfigureAnnBench) add_dependencies(CUVS_ANN_BENCH_ALL ${BENCH_NAME}) endfunction() +function(ConfigureNVBench) + set(oneValueArgs NAME) + set(multiValueArgs PATH LINKS CXXFLAGS INCLUDES) + + cmake_parse_arguments(ConfigureNVBench "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + set(BENCH_NAME ${ConfigureNVBench_NAME}_NVBENCH) + add_executable(${BENCH_NAME} ${ConfigureNVBench_PATH}) + target_link_libraries( + ${BENCH_NAME} + PRIVATE ${ConfigureNVBench_LINKS} nvbench::main CUDA::cudart_static + $ $ + ) + + set_target_properties( + ${BENCH_NAME} + PROPERTIES CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + POSITION_INDEPENDENT_CODE ON + BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" + ) + + set(NVBENCH_CXX_FLAGS ${CUVS_CXX_FLAGS} ${ConfigureNVBench_CXXFLAGS}) + target_compile_options( + ${BENCH_NAME} + PRIVATE "$<$:${NVBENCH_CXX_FLAGS}>" + "$<$:${CUVS_CUDA_FLAGS}>" + "$<$,$>:${CUVS_DEBUG_CUDA_FLAGS}>" + ) + + target_include_directories( + ${BENCH_NAME} + PUBLIC "$" + "$" + PRIVATE ${ConfigureNVBench_INCLUDES} + ) + + install( + TARGETS ${BENCH_NAME} + COMPONENT ann_bench + DESTINATION bin/ann + ) + + add_dependencies(CUVS_ANN_BENCH_ALL ${BENCH_NAME}) +endfunction() + # ################################################################################################## # * Configure benchmark targets ------------------------------------------------------------- @@ -272,7 +321,8 @@ if(CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE) endif() if(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE) - ConfigureAnnBench( + include(cmake/thirdparty/get_nvbench) + ConfigureNVBench( NAME CUVS_KNN_BRUTE_FORCE PATH $<$:src/cuvs/cuvs_brute_force_knn.cu> LINKS cuvs ) diff --git a/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu b/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu index 9aa491191a..f0e200654c 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu @@ -1,322 +1,107 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ + #include #include #include -#include -#include #include #include -#include -#include +#include + +#include + #include -#include -#include -namespace cuvs::neighbors::brute_force { +#include +#include +#include -struct print_metric { - cuvs::distance::DistanceType value; -}; +NVBENCH_DECLARE_TYPE_STRINGS(half, "F16", "half"); -struct RandomKNNInputs { - int num_queries; - int num_db_vecs; - int dim; - int k; - cuvs::distance::DistanceType metric; - bool row_major; -}; +namespace cuvs::neighbors::brute_force { -inline auto operator<<(std::ostream& os, const print_metric& p) -> std::ostream& +auto parse_metric(std::string const& metric) -> cuvs::distance::DistanceType { - switch (p.value) { - case cuvs::distance::DistanceType::L2Expanded: os << "L2Expanded"; break; - case cuvs::distance::DistanceType::L2SqrtExpanded: os << "L2SqrtExpanded"; break; - case cuvs::distance::DistanceType::CosineExpanded: os << "CosineExpanded"; break; - case cuvs::distance::DistanceType::L1: os << "L1"; break; - case cuvs::distance::DistanceType::L2Unexpanded: os << "L2Unexpanded"; break; - case cuvs::distance::DistanceType::L2SqrtUnexpanded: os << "L2SqrtUnexpanded"; break; - case cuvs::distance::DistanceType::InnerProduct: os << "InnerProduct"; break; - case cuvs::distance::DistanceType::Linf: os << "Linf"; break; - case cuvs::distance::DistanceType::Canberra: os << "Canberra"; break; - case cuvs::distance::DistanceType::LpUnexpanded: os << "LpUnexpanded"; break; - case cuvs::distance::DistanceType::CorrelationExpanded: os << "CorrelationExpanded"; break; - case cuvs::distance::DistanceType::JaccardExpanded: os << "JaccardExpanded"; break; - case cuvs::distance::DistanceType::HellingerExpanded: os << "HellingerExpanded"; break; - case cuvs::distance::DistanceType::Haversine: os << "Haversine"; break; - case cuvs::distance::DistanceType::BrayCurtis: os << "BrayCurtis"; break; - case cuvs::distance::DistanceType::JensenShannon: os << "JensenShannon"; break; - case cuvs::distance::DistanceType::HammingUnexpanded: os << "HammingUnexpanded"; break; - case cuvs::distance::DistanceType::KLDivergence: os << "KLDivergence"; break; - case cuvs::distance::DistanceType::RusselRaoExpanded: os << "RusselRaoExpanded"; break; - case cuvs::distance::DistanceType::DiceExpanded: os << "DiceExpanded"; break; - case cuvs::distance::DistanceType::Precomputed: os << "Precomputed"; break; - default: RAFT_FAIL("unreachable code"); - } - return os; + if (metric == "InnerProduct") { return cuvs::distance::DistanceType::InnerProduct; } + if (metric == "L2SqrtExpanded") { return cuvs::distance::DistanceType::L2SqrtExpanded; } + throw std::invalid_argument("Unsupported distance metric: " + metric); } -std::ostream& operator<<(std::ostream& os, const RandomKNNInputs& input) +template +void search_benchmark(nvbench::state& state, nvbench::type_list) { - return os << "num_queries:" << input.num_queries << " num_vecs:" << input.num_db_vecs - << " dim:" << input.dim << " k:" << input.k << " metric:" << print_metric{input.metric} - << " row_major:" << input.row_major; -} - -template -class BruteForceKNNBenchmark { - public: - BruteForceKNNBenchmark(const RandomKNNInputs& params, const std::string& type_str) - : stream_(raft::resource::get_cuda_stream(handle_)), - params_(params), - type_str_(type_str), - database(params_.num_db_vecs * params_.dim, stream_), - search_queries(params_.num_queries * params_.dim, stream_), - cuvs_indices_(params_.num_queries * params_.k, stream_), - cuvs_distances_(params_.num_queries * params_.k, stream_) - { - raft::matrix::fill( - handle_, - raft::make_device_matrix_view(database.data(), params_.num_db_vecs, params_.dim), - T{0.0}); - raft::matrix::fill( - handle_, - raft::make_device_matrix_view(search_queries.data(), params_.num_queries, params_.dim), - T{0.0}); - raft::matrix::fill( - handle_, - raft::make_device_matrix_view(cuvs_distances_.data(), params_.num_queries, params_.k), - DistT{0.0}); - } - - void runBenchmark() - { - DistT metric_arg = 3.0; - rmm::device_uvector workspace(0, stream_); - - std::chrono::duration build_dur; - std::chrono::duration search_dur; - - auto indices = raft::make_device_matrix_view( - cuvs_indices_.data(), params_.num_queries, params_.k); - auto distances = raft::make_device_matrix_view( - cuvs_distances_.data(), params_.num_queries, params_.k); - raft::resource::sync_stream(handle_, stream_); - - if (params_.row_major) { - { - auto idx_warm = - cuvs::neighbors::brute_force::build(handle_, - raft::make_device_matrix_view( - database.data(), params_.num_db_vecs, params_.dim), - params_.metric, - metric_arg); - cuvs::neighbors::brute_force::search( - handle_, - idx_warm, - raft::make_device_matrix_view( - search_queries.data(), params_.num_queries, params_.dim), - indices, - distances, - cuvs::neighbors::filtering::none_sample_filter{}); - flush_l2_cache(); - raft::resource::sync_stream(handle_, stream_); - } - - auto start = std::chrono::high_resolution_clock::now(); - auto idx = - cuvs::neighbors::brute_force::build(handle_, - raft::make_device_matrix_view( - database.data(), params_.num_db_vecs, params_.dim), - params_.metric, - metric_arg); - raft::resource::sync_stream(handle_, stream_); - auto end = std::chrono::high_resolution_clock::now(); - build_dur = end - start; - - start = std::chrono::high_resolution_clock::now(); - cuvs::neighbors::brute_force::search( - handle_, - idx, - raft::make_device_matrix_view( - search_queries.data(), params_.num_queries, params_.dim), - indices, - distances, - cuvs::neighbors::filtering::none_sample_filter{}); - raft::resource::sync_stream(handle_, stream_); - end = std::chrono::high_resolution_clock::now(); - search_dur = end - start; - - } else { - { - auto idx_warm = - cuvs::neighbors::brute_force::build(handle_, - raft::make_device_matrix_view( - database.data(), params_.num_db_vecs, params_.dim), - params_.metric, - metric_arg); - cuvs::neighbors::brute_force::search( - handle_, - idx_warm, - raft::make_device_matrix_view( - search_queries.data(), params_.num_queries, params_.dim), - indices, - distances, - cuvs::neighbors::filtering::none_sample_filter{}); - flush_l2_cache(); - raft::resource::sync_stream(handle_, stream_); - } - - auto start = std::chrono::high_resolution_clock::now(); - auto idx = cuvs::neighbors::brute_force::build( - handle_, - raft::make_device_matrix_view( - database.data(), params_.num_db_vecs, params_.dim), - params_.metric, - metric_arg); - raft::resource::sync_stream(handle_, stream_); - auto end = std::chrono::high_resolution_clock::now(); - build_dur = end - start; - - start = std::chrono::high_resolution_clock::now(); - cuvs::neighbors::brute_force::search( - handle_, - idx, - raft::make_device_matrix_view( - search_queries.data(), params_.num_queries, params_.dim), - indices, - distances, - cuvs::neighbors::filtering::none_sample_filter{}); - raft::resource::sync_stream(handle_, stream_); - end = std::chrono::high_resolution_clock::now(); - search_dur = end - start; - } - - double total_dur = build_dur.count() + search_dur.count(); - double throughput = static_cast(params_.num_queries) / (total_dur / 1000.0); - ; - printResult(params_, build_dur.count(), search_dur.count(), total_dur, throughput); + auto const num_queries = state.get_int64("num_queries"); + auto const num_db_vecs = state.get_int64("num_db_vecs"); + auto const dim = state.get_int64("dim"); + auto const k = state.get_int64("k"); + auto const metric = parse_metric(state.get_string("metric")); + auto const layout = state.get_string("layout"); + + if (layout != "row_major" && layout != "column_major") { + throw std::invalid_argument("Unsupported layout: " + layout); } - void setUp() - { - unsigned long long int seed = 1234ULL; - raft::random::RngState r(seed); - - // JensenShannon distance requires positive values - T min_val = params_.metric == cuvs::distance::DistanceType::JensenShannon ? T(0.0) : T(-1.0); - uniform(handle_, r, database.data(), params_.num_db_vecs * params_.dim, min_val, T(1.0)); - uniform(handle_, r, search_queries.data(), params_.num_queries * params_.dim, min_val, T(1.0)); - } - - private: - void flush_l2_cache() - { - int l2_cache_size = 0; - int device_id = 0; - RAFT_CUDA_TRY(cudaGetDevice(&device_id)); - RAFT_CUDA_TRY(cudaDeviceGetAttribute(&l2_cache_size, cudaDevAttrL2CacheSize, device_id)); - scratch_buf_ = rmm::device_buffer(l2_cache_size * 3, stream_); - RAFT_CUDA_TRY(cudaMemsetAsync(scratch_buf_.data(), 0, scratch_buf_.size(), stream_)); - }; - - void printResult(const RandomKNNInputs& params, - double build_time, - double search_time, - double total_time, - double throughput) - { - std::cout << std::left << std::setw(15) << type_str_ << std::setw(10) << params.num_queries - << std::setw(10) << params.num_db_vecs << std::setw(10) << params.dim << std::setw(10) - << params.k << std::setw(20) << print_metric{params.metric} << std::setw(15) - << (params.row_major ? "row" : "col") << std::right << std::setw(20) << std::fixed - << std::setprecision(3) << build_time << std::right << std::setw(20) << std::fixed - << std::setprecision(3) << search_time << std::right << std::setw(20) << std::fixed - << std::setprecision(3) << total_time << std::right << std::setw(20) << std::fixed - << std::setprecision(3) << throughput << "\n"; - } - raft::resources handle_; - cudaStream_t stream_ = 0; - RandomKNNInputs params_; - rmm::device_uvector database; - rmm::device_uvector search_queries; - rmm::device_uvector cuvs_indices_; - rmm::device_uvector cuvs_distances_; - rmm::device_buffer scratch_buf_; - std::string type_str_; -}; - -static std::vector getInputs() -{ - std::vector param_vec; - struct TestParams { - int num_queries; - int num_db_vecs; - int dim; - int k; - cuvs::distance::DistanceType metric; - bool row_major; - }; - - const std::vector params_group = raft::util::itertools::product( - {int(10), int(100), int(1024)}, - {int(1000000)}, - {int(32), int(256), int(1024)}, - {int(128), int(1024)}, - {cuvs::distance::DistanceType::InnerProduct, cuvs::distance::DistanceType::L2SqrtExpanded}, - {true, false}); - - param_vec.reserve(params_group.size()); - for (TestParams params : params_group) { - param_vec.push_back(RandomKNNInputs({params.num_queries, - params.num_db_vecs, - params.dim, - params.k, - params.metric, - params.row_major})); + raft::resources handle; + auto const stream = raft::resource::get_cuda_stream(handle); + state.set_cuda_stream(nvbench::make_cuda_stream_view(stream)); + + rmm::device_uvector database(num_db_vecs * dim, stream); + rmm::device_uvector queries(num_queries * dim, stream); + rmm::device_uvector neighbors(num_queries * k, stream); + rmm::device_uvector distances(num_queries * k, stream); + + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, database.data(), database.size(), T(-1.0f), T(1.0f)); + raft::random::uniform(handle, rng, queries.data(), queries.size(), T(-1.0f), T(1.0f)); + + auto const neighbors_view = + raft::make_device_matrix_view(neighbors.data(), num_queries, k); + auto const distances_view = + raft::make_device_matrix_view(distances.data(), num_queries, k); + + index_params index_params; + index_params.metric = metric; + index_params.metric_arg = 3.0; + search_params search_params; + + state.add_element_count(num_queries); + + if (layout == "row_major") { + auto const database_view = + raft::make_device_matrix_view(database.data(), num_db_vecs, dim); + auto const queries_view = + raft::make_device_matrix_view(queries.data(), num_queries, dim); + auto index = build(handle, index_params, database_view); + + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { + search(handle, search_params, index, queries_view, neighbors_view, distances_view); + }); + } else { + auto const database_view = raft::make_device_matrix_view( + database.data(), num_db_vecs, dim); + auto const queries_view = raft::make_device_matrix_view( + queries.data(), num_queries, dim); + auto index = build(handle, index_params, database_view); + + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { + search(handle, search_params, index, queries_view, neighbors_view, distances_view); + }); } - return param_vec; } -void printHeader() -{ - std::cout << std::left << std::setw(15) << "Type" << std::setw(10) << "Queries" << std::setw(10) - << "Vectors" << std::setw(10) << "Dim" << std::setw(10) << "K" << std::setw(20) - << "Metric" << std::setw(15) << "Layout" << std::right << std::setw(20) - << "Build Time (ms)" << std::right << std::setw(20) << "Search Time (ms)" << std::right - << std::setw(20) << "Total Time (ms)" << std::right << std::setw(20) - << "Throughput (q/s)" - << "\n"; - std::cout << std::string(165, '-') << "\n"; -} +using value_types = nvbench::type_list; -void runBenchmarkForType() -{ - auto selected_inputs = getInputs(); - for (const auto& input : selected_inputs) { - { - BruteForceKNNBenchmark benchmark(input, "float"); - benchmark.setUp(); - benchmark.runBenchmark(); - } - { - BruteForceKNNBenchmark benchmark(input, "half"); - benchmark.setUp(); - benchmark.runBenchmark(); - } - } -} +NVBENCH_BENCH_TYPES(search_benchmark, NVBENCH_TYPE_AXES(value_types)) + .set_name("brute_force_search") + .add_int64_axis("num_queries", {10, 100, 1024}) + .add_int64_axis("num_db_vecs", {1000000}) + .add_int64_axis("dim", {32, 256, 1024}) + .add_int64_axis("k", {128, 1024}) + .add_string_axis("metric", {"InnerProduct", "L2SqrtExpanded"}) + .add_string_axis("layout", {"row_major", "column_major"}); } // namespace cuvs::neighbors::brute_force - -int main() -{ - cuvs::neighbors::brute_force::printHeader(); - cuvs::neighbors::brute_force::runBenchmarkForType(); - return 0; -} diff --git a/cpp/cmake/thirdparty/get_nvbench.cmake b/cpp/cmake/thirdparty/get_nvbench.cmake new file mode 100644 index 0000000000..3752246f0a --- /dev/null +++ b/cpp/cmake/thirdparty/get_nvbench.cmake @@ -0,0 +1,14 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +# This function finds or builds NVBench. +function(find_and_configure_nvbench) + include(${rapids-cmake-dir}/cpm/nvbench.cmake) + rapids_cpm_nvbench(BUILD_STATIC) +endfunction() + +find_and_configure_nvbench()