From 8b9803717b5c0a99792c2ab55958deb1ef1e60db Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 27 Jan 2026 14:54:36 +0000 Subject: [PATCH 01/48] Forwarding all the previous work --- cpp/CMakeLists.txt | 3 +- cpp/bench/prims/CMakeLists.txt | 118 ++++++ cpp/bench/prims/src/distance/fused_l2_nn.cu | 379 ++++++++++++++++++++ cpp/src/cluster/detail/kmeans_balanced.cuh | 143 ++++++-- cpp/src/distance/unfused_distance_nn.cuh | 343 ++++++++++++++++++ cpp/tests/CMakeLists.txt | 2 +- cpp/tests/neighbors/1_nn.cu | 215 +++++++++++ cpp/tests/neighbors/1_nn_helper.cuh | 206 +++++++++++ 8 files changed, 1386 insertions(+), 23 deletions(-) create mode 100644 cpp/bench/prims/CMakeLists.txt create mode 100644 cpp/bench/prims/src/distance/fused_l2_nn.cu create mode 100644 cpp/src/distance/unfused_distance_nn.cuh create mode 100644 cpp/tests/neighbors/1_nn.cu create mode 100644 cpp/tests/neighbors/1_nn_helper.cuh diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index e80220bce0..0f24c9a455 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -50,7 +50,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) option(BUILD_SHARED_LIBS "Build cuvs shared libraries" ON) option(BUILD_TESTS "Build cuvs unit-tests" ON) option(BUILD_C_LIBRARY "Build cuVS C API library" ON) -option(BUILD_CUVS_BENCH "Build cuVS ann benchmarks" OFF) +option(BUILD_CUVS_BENCH "Build cuVS ann and prims benchmarks" OFF) option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON) option(BUILD_MG_ALGOS "Build with multi-GPU support" ON) option(CUDA_ENABLE_KERNELINFO "Enable kernel resource usage info" OFF) @@ -808,4 +808,5 @@ endif() if(BUILD_CUVS_BENCH) add_subdirectory(bench/ann/) + add_subdirectory(bench/prims) endif() diff --git a/cpp/bench/prims/CMakeLists.txt b/cpp/bench/prims/CMakeLists.txt new file mode 100644 index 0000000000..4852099873 --- /dev/null +++ b/cpp/bench/prims/CMakeLists.txt @@ -0,0 +1,118 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +list(APPEND CMAKE_MODULE_PATH "${CUVS_SOURCE_DIR}") + +# ################################################################################################## +# * benchmark options ------------------------------------------------------------------------------ + +option(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN "Include Fused L2 NN benchmark" ON) +option(CUVS_PRIMS_BENCH_USE_CUBLAS_SAMPLE "Include cuBLAS sample benchmark" ON) +option(CUVS_PRIMS_BENCH_USE_WMMA "Include WMMA sample benchmark" ON) + +# ################################################################################################## +# * Process options ---------------------------------------------------------- + +find_package(Threads REQUIRED) + +# ################################################################################################## +# * Fetch requirements ------------------------------------------------------------- + +include(cmake/thirdparty/get_nlohmann_json) + +# ################################################################################################## +# * Target function ------------------------------------------------------------- + +function(ConfigurePrimsBench) + + set(oneValueArgs NAME) + set(multiValueArgs PATH LINKS CXXFLAGS) + + #if(NOT BUILD_CPU_ONLY) + set(GPU_BUILD ON) + #endif() + + cmake_parse_arguments( + ConfigurePrimsBench "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} + ) + + set(BENCH_NAME ${ConfigurePrimsBench_NAME}_PRIMS_BENCH) + + add_executable(${BENCH_NAME} ${ConfigurePrimsBench_PATH}) + target_compile_definitions(${BENCH_NAME} PRIVATE PRIMS_BENCH_BUILD_MAIN) + target_link_libraries( + ${BENCH_NAME} PRIVATE benchmark::benchmark $<$:CUDA::nvtx3> + ) + + target_link_libraries( + ${BENCH_NAME} + PRIVATE ${ConfigurePrimsBench_LINKS} + nlohmann_json::nlohmann_json + Threads::Threads + $<$:CUDA::cudart_static> + $ + $ + ) + + set_target_properties( + ${BENCH_NAME} + PROPERTIES # set target compile options + CXX_STANDARD 17 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + POSITION_INDEPENDENT_CODE ON + INTERFACE_POSITION_INDEPENDENT_CODE ON + BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" + ) + + set(${ConfigurePrimsBench_CXXFLAGS} ${CUVS_CXX_FLAGS} ${ConfigurePrimsBench_CXXFLAGS}) + + target_compile_options( + ${BENCH_NAME} PRIVATE "$<$:${ConfigurePrimsBench_CXXFLAGS}>" + "$<$:${CUVS_CUDA_FLAGS}>" + ) + + if(CUVS_PRIMS_BENCH_USE_${ConfigurePrimsBench_NAME}) + target_compile_definitions( + ${BENCH_NAME} + PUBLIC + CUVS_PRIMS_BENCH_USE_${ConfigurePrimsBench_NAME}=CUVS_PRIMS_BENCH_USE_${ConfigurePrimsBench_NAME} + ) + endif() + + target_include_directories( + ${BENCH_NAME} + PUBLIC "$" + PRIVATE ${ConfigurePrimsBench_INCLUDES} + ) + + install( + TARGETS ${BENCH_NAME} + COMPONENT prims_bench + DESTINATION bin/prims + ) + + add_dependencies(CUVS_PRIMS_BENCH_ALL ${BENCH_NAME}) +endfunction() + +# ################################################################################################## +# * Configure benchmark targets ------------------------------------------------------------- + +if(NOT TARGET CUVS_PRIMS_BENCH_ALL) + add_custom_target(CUVS_PRIMS_BENCH_ALL) +endif() + +if(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN) + ConfigurePrimsBench( + NAME FUSED_L2_NN PATH src/distance/fused_l2_nn.cu + LINKS CUDA::cublas rmm::rmm + ) +endif() + +message("CUDAToolkit_LIBRARY_DIR: ${CUDAToolkit_LIBRARY_DIR}") diff --git a/cpp/bench/prims/src/distance/fused_l2_nn.cu b/cpp/bench/prims/src/distance/fused_l2_nn.cu new file mode 100644 index 0000000000..42905b8224 --- /dev/null +++ b/cpp/bench/prims/src/distance/fused_l2_nn.cu @@ -0,0 +1,379 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include "../../../../src/distance/fused_distance_nn.cuh" +#include "../../../../src/distance/unfused_distance_nn.cuh" + +#include +#include +#include +#include + +using cuvs::distance::DistanceType; +using cuvs::distance::fusedDistanceNNMinReduce; +using cuvs::distance::reduce_min; +using cuvs::distance::unfused_distance_nn; +using cuvs::distance::unfusedDistanceNNMinReduce; + +enum class AlgorithmType { gemm, unfused, fused }; + +__global__ void fill_int8(int8_t* buff, int len) +{ + int tid = threadIdx.x + blockIdx.x * blockDim.x; + // Fill the buffer with pseudo-random int8_t values using a simple LCG + if (tid < len) { + // Simple LCG: x_n+1 = (a * x_n + c) % m + // Use tid as seed, constants chosen for decent distribution + int seed = tid * 1103515245 + 12345; + buff[tid] = static_cast((seed >> 16) & 0xFF); + } +} + +template +void benchmark_fusedl2nn(benchmark::State& state) +{ + const int m = state.range(0); + const int n = state.range(1); + const int k = state.range(2); + // const bool sqrt = state.range(3); + // const DistanceType metric = DistanceType(state.range(4)); + + raft::device_resources handle; + rmm::cuda_stream_view stream; + + stream = raft::resource::get_cuda_stream(handle); + + auto x = raft::make_device_matrix(handle, m, k); + auto y = raft::make_device_matrix(handle, n, k); + auto x_norm = raft::make_device_vector(handle, m); + auto y_norm = raft::make_device_vector(handle, n); + auto out = raft::make_device_vector(handle, m); + + raft::random::RngState rng{1234}; + if constexpr (std::is_same_v) { + fill_int8<<<1000, 256, 0, stream>>>(x.data_handle(), m * k); + fill_int8<<<1000, 256, 0, stream>>>(y.data_handle(), n * k); + } else { + raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); + raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); + } + + // Pre-compute norms + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, m, stream); + raft::linalg::rowNorm( + y_norm.data_handle(), y.data_handle(), k, n, stream); + + // Calculate the workspace size + // for fused it is m * sizeof(IdxT) + // for unfused, gemm and tensor it is m * n * sizeof(AccT); + + size_t workspace_size = (algo == AlgorithmType::fused) ? m * sizeof(IdxT) : m * n * sizeof(AccT); + + raft::device_vector workspace = + raft::make_device_vector(handle, workspace_size); + + // Warm up + if constexpr (algo != AlgorithmType::fused) { + unfusedDistanceNNMinReduce(handle, + out.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + static_cast(m), + static_cast(n), + static_cast(k), + (AccT*)workspace.data_handle(), + sqrt, + true, + true, + metric, + float(0.0), + stream); + } + + RAFT_CUDA_TRY(cudaMemsetAsync(workspace.data_handle(), 0, workspace_size, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, m * sizeof(OutT), stream)); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + + for (auto _ : state) { + if constexpr (algo == AlgorithmType::fused) { + fusedDistanceNNMinReduce(out.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + static_cast(m), + static_cast(n), + static_cast(k), + (void*)workspace.data_handle(), + sqrt, + true, + true, + metric, + 0.0, + stream); + } + + if constexpr (algo == AlgorithmType::unfused) { + unfusedDistanceNNMinReduce(handle, + out.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + static_cast(m), + static_cast(n), + static_cast(k), + (AccT*)workspace.data_handle(), + sqrt, + true, + true, + metric, + float(0.0), + stream); + } + + if constexpr (algo == AlgorithmType::gemm) { + unfusedDistanceNNMinReduce(handle, + out.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + static_cast(m), + static_cast(n), + static_cast(k), + (AccT*)workspace.data_handle(), + sqrt, + true, + true, + metric, + float(0.0), + stream); + } + } + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + if constexpr (algo == AlgorithmType::gemm) { + if (metric == DistanceType::L2Expanded) { + reduce_min(out.data_handle(), + (AccT*)workspace.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + stream, + sqrt, + true); + } else if (metric == DistanceType::L2SqrtExpanded) { + reduce_min( + out.data_handle(), + (AccT*)workspace.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + stream, + sqrt, + true); + } else if (metric == DistanceType::CosineExpanded) { + reduce_min( + out.data_handle(), + (AccT*)workspace.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + stream, + sqrt, + true); + } + } + + state.counters["M"] = m; + state.counters["N"] = n; + state.counters["K"] = k; + int64_t total_ops = int64_t(2) * m * n * k; + state.counters["FLOP/s"] = + benchmark::Counter(total_ops, benchmark::Counter::kIsIterationInvariantRate); +} + +template +static void CustomArguments(benchmark::internal::Benchmark* b) +{ + /*constexpr int K = 1024; + std::vector m_list = {4 * K, 8 * K}; + std::vector n_list = {4 * K, 8 * K, 16 * K}; + // std::vector k_list = {128, 512, 1024, 1536}; + std::vector k_list = {128, 256, 512}; + for (auto k : k_list) { + for (auto m : m_list) { + for (auto n : n_list) { + b->Args({m, n, k}); + } + } + }*/ + // b->Args({65536, 256, 128, true, DistanceType::L2Expanded}); + // b->Args({65536, 256, 128, false, DistanceType::CosineExpanded}); + // b->Args({65536, 10000, 768, false}); +} + +int main(int argc, char** argv) +{ + using IdxT = int64_t; + IdxT M = 1024; + IdxT N = 1024; + IdxT K = 128; + for (int i = 0; i < argc; i++) { + if (strcmp(argv[i], "-M") == 0) { + M = std::stoi(argv[i + 1]); + for (int j = i; j < argc - 2; j++) { + argv[j] = argv[j + 2]; + } + argc -= 2; + i -= 2; + } else if (strcmp(argv[i], "-N") == 0) { + N = std::stoi(argv[i + 1]); + for (int j = i; j < argc - 2; j++) { + argv[j] = argv[j + 2]; + } + argc -= 2; + i -= 2; + } else if (strcmp(argv[i], "-K") == 0) { + K = std::stoi(argv[i + 1]); + for (int j = i; j < argc - 2; j++) { + argv[j] = argv[j + 2]; + } + argc -= 2; + i -= 2; + } + } + + constexpr bool sqrt = false; + constexpr DistanceType dist = DistanceType::L2Expanded; + + benchmark::internal::Benchmark* bench; + + // In the following instances IdxT is int64_t, it does not seem to have much impact on the + // performance OutT is always + + // Method: Fused, DataT: float, AccT: float + bench = benchmark::RegisterBenchmark("fused/float/float", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::fused, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: unfused, DataT: float, AccT: float + bench = benchmark::RegisterBenchmark("unfused/float/float", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: gemm, DataT: float, AccT: float + bench = benchmark::RegisterBenchmark("gemm/float/float", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: unfused, DataT: half, AccT: float + bench = benchmark::RegisterBenchmark("unfused/half/float", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: gemm, DataT: half, AccT: float + bench = benchmark::RegisterBenchmark("gemm/half/float", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: unfused, DataT: half, AccT: half + bench = benchmark::RegisterBenchmark("unfused/half/half", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: gemm, DataT: half, AccT: half + bench = benchmark::RegisterBenchmark("gemm/half/half", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: unfused, DataT: int8_t, AccT: int32_t + bench = benchmark::RegisterBenchmark("unfused/int8_t/int32_t", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Method: gemm, DataT: int8_t, AccT: int32_t + bench = benchmark::RegisterBenchmark("gemm/int8_t/int32_t", + benchmark_fusedl2nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K}); + + // Initialize benchmark + ::benchmark::Initialize(&argc, argv); + if (::benchmark::ReportUnrecognizedArguments(argc, argv)) return -1; + ::benchmark::RunSpecifiedBenchmarks(); + ::benchmark::Shutdown(); + return 0; +} diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 2637e81517..171786e4ad 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -1,11 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include "../../distance/fused_distance_nn.cuh" +#include "../../distance/unfused_distance_nn.cuh" + #include "kmeans_common.cuh" #include @@ -16,8 +18,10 @@ #include #include #include + #include #include +#include #include #include #include @@ -51,6 +55,24 @@ namespace cuvs::cluster::kmeans::detail { constexpr static inline float kAdjustCentersWeight = 7.0f; +template +bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) +{ + cudaDeviceProp prop; + prop = raft::resource::get_device_properties(handle); + if (prop.major <= 8) { + // Use fused for Ampere or before + return true; + } else if (prop.major == 9 && (m >= 4096 || n >= 4096)) { + // On Hopper if m, n are bigger than 4096, use fused + return true; + } else if (prop.major >= 12) { + // On Blackwell onwards, use unfused + return false; + } + return false; +} + /** * @brief Predict labels for the dataset; floating-point types only. * @@ -72,6 +94,38 @@ constexpr static inline float kAdjustCentersWeight = 7.0f; * @param[out] labels Output predictions [n_rows] * @param[inout] mr (optional) Memory resource to use for temporary allocations */ +#include +#include +#include +#include + +static void add(int64_t i, int64_t j, int64_t k) +{ + // Static vector of pairs: (tuple, counter) + static std::vector, int>> tupleList; + + std::tuple tuple = std::make_tuple(i, j, k); + + auto it = std::find_if( + tupleList.begin(), tupleList.end(), [&](const auto& p) { return p.first == tuple; }); + if (it != tupleList.end()) { + // Tuple exists, increment counter + it->second += 1; + } else { + // New tuple, add with counter 1 + tupleList.emplace_back(tuple, 1); + } + /*FILE* outFile; + + outFile = fopen("tuples.txt", "w"); + for (const auto& entry : tupleList) { + int64_t _i, _j, _k; + std::tie(_i, _j, _k) = entry.first; + fprintf(outFile, "%d,%lu,%lu,%lu\n", entry.second, _i, _j, _k); + } + fclose(outFile); */ +} + template inline std::enable_if_t> predict_core( const raft::resources& handle, @@ -85,42 +139,88 @@ inline std::enable_if_t> predict_core( LabelT* labels, rmm::device_async_resource_ref mr) { - auto stream = raft::resource::get_cuda_stream(handle); + std::string result = "predict_core_" + std::to_string(n_rows) + "_" + std::to_string(n_clusters) + + "_" + std::to_string(dim) + "_" + std::to_string(sizeof(MathT)); + nvtxRangePushA(result.c_str()); + auto stream = raft::resource::get_cuda_stream(handle); + bool should_use_fused = use_fused(handle, n_rows, n_clusters, dim); + switch (params.metric) { case cuvs::distance::DistanceType::L2Expanded: case cuvs::distance::DistanceType::L2SqrtExpanded: { - auto workspace = raft::make_device_mdarray( - handle, mr, raft::make_extents((sizeof(int)) * n_rows)); - + nvtxRangePushA("pc::workspace_alloc"); + size_t workspace_size = + should_use_fused ? sizeof(IdxT) * n_rows : sizeof(MathT) * n_rows * n_clusters; + auto workspace = + raft::make_device_mdarray(handle, mr, raft::make_extents(workspace_size)); + nvtxRangePop(); + nvtxRangePushA("pc::output_alloc"); auto minClusterAndDistance = raft::make_device_mdarray, IdxT>( handle, mr, raft::make_extents(n_rows)); + nvtxRangePop(); + /*auto unf_minClusterAndDistance = + raft::make_device_mdarray, IdxT>( + handle, mr, raft::make_extents(n_rows));*/ + + nvtxRangePushA("pc::output_init"); raft::KeyValuePair initial_value(0, std::numeric_limits::max()); thrust::fill(raft::resource::get_thrust_policy(handle), minClusterAndDistance.data_handle(), minClusterAndDistance.data_handle() + minClusterAndDistance.size(), initial_value); + nvtxRangePop(); + nvtxRangePushA("pc::row_norm"); auto centroidsNorm = raft::make_device_mdarray(handle, mr, raft::make_extents(n_clusters)); raft::linalg::rowNorm( centroidsNorm.data_handle(), centers, dim, n_clusters, stream); - cuvs::distance::fusedDistanceNNMinReduce, IdxT>( - minClusterAndDistance.data_handle(), - dataset, - centers, - dataset_norm, - centroidsNorm.data_handle(), - n_rows, - n_clusters, - dim, - (void*)workspace.data_handle(), - (params.metric == cuvs::distance::DistanceType::L2Expanded) ? false : true, - false, - true, - params.metric, - 0.0f, - stream); + add(int64_t(n_rows), int64_t(n_clusters), int64_t(dim)); + nvtxRangePop(); + static int print_guard = 0; + if (should_use_fused) { + if (print_guard < 5) printf("fused called\n"); + nvtxRangePushA("pc::fused_call"); + cuvs::distance::fusedDistanceNNMinReduce, IdxT>( + minClusterAndDistance.data_handle(), + dataset, + centers, + dataset_norm, + centroidsNorm.data_handle(), + n_rows, + n_clusters, + dim, + (void*)workspace.data_handle(), + (params.metric == cuvs::distance::DistanceType::L2Expanded) ? true : false, + false, + true, + params.metric, + 0.0f, + stream); + nvtxRangePop(); + } else { + if (print_guard < 5) printf("unfused called\n"); + cuvs::distance:: + unfusedDistanceNNMinReduce, IdxT>( + handle, + minClusterAndDistance.data_handle(), + dataset, + centers, + dataset_norm, + centroidsNorm.data_handle(), + n_rows, + n_clusters, + dim, + (void*)workspace.data_handle(), + (params.metric == cuvs::distance::DistanceType::L2Expanded) ? false : true, + false, + true, + params.metric, + 0.0f, + stream); + } + print_guard++; // todo(lsugy): use KVP + iterator in caller. // Copy keys to output labels @@ -205,6 +305,7 @@ inline std::enable_if_t> predict_core( RAFT_FAIL("The chosen distance metric is not supported (%d)", int(params.metric)); } } + nvtxRangePop(); } /** diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh new file mode 100644 index 0000000000..6c6e285289 --- /dev/null +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -0,0 +1,343 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#ifndef __UNFUSED_DISTANCE_NN_H +#define __UNFUSED_DISTANCE_NN_H + +#pragma once + +#include +#include + +namespace cuvs { +namespace distance { + +/** + * \ingroup unfused_distance_nn + * @{ + */ +/** + * @brief Unfused distance and 1-nearest-neighbor computation in a single call. + * + * Unlike the fused call, here we do explicit gemm call followed by a reduction call. + * This code path exists because the fused path is sometime not the optimal due to + * GEMM-like kernel not being optimized yet. + * + * + * @tparam DataT data type + * @tparam OutT output type to either store 1-NN indices and their minimum + * distances or store only the min distances. Accordingly, one + * has to pass an appropriate `ReduceOpT` + * @tparam IdxT indexing arithmetic type + * @tparam ReduceOpT A struct to perform the final needed reduction operation + * and also to initialize the output array elements with the + * appropriate initial value needed for reduction. + * @tparam KVPReduceOpT A struct providing functions for key-value pair comparison. + * + * @param[out] min will contain the reduced output (Length = `m`) + * (on device) + * @param[in] x first matrix. Row major. Dim = `m x k`. + * (on device). + * @param[in] y second matrix. Row major. Dim = `n x k`. + * (on device). + * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). + * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) + * @param[in] m gemm m + * @param[in] n gemm n + * @param[in] k gemm k + * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) + * @param[in] redOp reduction operator in the epilogue + * @param[in] pairRedOp reduction operation on key value pairs + * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt + * @param[in] initOutBuffer whether to initialize the output buffer before the + * main kernel launch + * @param[in] isRowMajor whether the input/output is row or column major. + * @param[in] metric Distance metric to be used (supports L2, cosine) + * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) + * @param[in] stream cuda stream + */ + +template +__host__ __device__ T max_val() +{ + if constexpr (std::is_same::value) { + return CUDART_MAX_NORMAL_FP16; + } else { + return cuda::std::numeric_limits::max(); + } +} + +template +T min_val() +{ + if constexpr (std::is_same::value) { + return CUDART_MIN_DENORM_FP16; + } else { + return cuda::std::numeric_limits::min(); + } +} + +template +struct Reducer { + typedef raft::KeyValuePair KVType; + + __device__ KVType operator()(const KVType& a, const KVType& b) + { + if ((a.value < b.value) || (a.value == b.value && a.key < b.key)) { + return a; + } else { + return b; + } + } + + __device__ AccT operator()(const AccT& a, const AccT& b) { return a < b ? a : b; } +}; + +template +__global__ void reduce_min_kernel(OutT* out, + const AccT* z, + const AccT* x_norm, + const AccT* y_norm, + IdxT m, + IdxT n, + bool is_sqrt, + bool initOutBuffer) +{ + IdxT tid = threadIdx.x + blockIdx.x * blockDim.x; + IdxT row = blockIdx.x; + + AccT x_norm_row = x_norm[row]; + + typedef raft::KeyValuePair KVType; + KVType thread_min; + + thread_min.value = max_val(); + thread_min.key = max_val(); + + for (IdxT col = threadIdx.x; col < n; col += TPB) { + AccT dist = 0.0; + + if constexpr (metric == DistanceType::L2SqrtExpanded || metric == DistanceType::L2Expanded) { + dist = x_norm_row + y_norm[col] - AccT(2) * z[row * n + col]; + } else if constexpr (metric == DistanceType::CosineExpanded) { + dist = AccT(1.0) - (z[row * n + col] / (x_norm_row * y_norm[col])); + } + if (dist < thread_min.value) { + thread_min.value = dist; + thread_min.key = col; + } + } + + typedef cub::BlockReduce BlockReduceT; + __shared__ typename BlockReduceT::TempStorage temp_storage; + auto block_result = BlockReduceT(temp_storage).Reduce(thread_min, Reducer{}); + + if (threadIdx.x == 0) { + if (is_sqrt) { block_result.value = raft::sqrt(block_result.value); } + if constexpr (std::is_same_v) { + if (initOutBuffer == true) { + out[row] = block_result; + } else { + out[row] = block_result.value < out[row].value ? block_result : out[row]; + } + } else { + if (initOutBuffer == true) { + out[row] = block_result.value; + } else { + out[row] = block_result.value < out[row] ? block_result.value : out[row]; + } + } + } +} + +template +void reduce_min(OutT* out, + const AccT* z, + const AccT* x_norm, + const AccT* y_norm, + IdxT m, + IdxT n, + cudaStream_t stream, + bool is_sqrt, + bool initOutBuffer) +{ + const int TPB = 128; + + int blocks = m; + reduce_min_kernel + <<>>(out, z, x_norm, y_norm, m, n, is_sqrt, initOutBuffer); +} + +template +void unfused_distance_nn(raft::resources const& handle, + OutT* out, + const DataT* x, + const DataT* y, + IdxT M, + IdxT N, + IdxT K, + const AccT* x_norm, + const AccT* y_norm, + AccT* workspace, + bool is_sqrt, + bool initOutBuffer, + DistanceType metric, + cudaStream_t stream) +{ + assert((metric == DistanceType::CosineExpanded) || (metric == DistanceType::L2Expanded) || + (metric == DistanceType::L2SqrtExpanded)); + + cudaDataType_t xyType, zType; + cublasComputeType_t computeType; + + const half h_alpha = 1.0, h_beta = 0.0; + const float f_alpha = 1.0, f_beta = 0.0; + const double d_alpha = 1.0, d_beta = 0.0; + const int32_t i32_alpha = 1, i32_beta = 0; + const int8_t i8_alpha = 1, i8_beta = 0; + + const void* alpha = nullptr; + const void* beta = nullptr; + if constexpr (std::is_same_v) { + xyType = CUDA_R_32F; + zType = CUDA_R_32F; + computeType = CUBLAS_COMPUTE_32F; + // computeType = CUBLAS_COMPUTE_32F_FAST_TF32; + // computeType = CUBLAS_COMPUTE_32F_FAST_16F; + // computeType = CUBLAS_COMPUTE_32F_FAST_16BF; + // computeType = CUBLAS_COMPUTE_32F_EMULATED_16BFX9; + alpha = reinterpret_cast(&f_alpha); + beta = reinterpret_cast(&f_beta); + } else if constexpr (std::is_same_v) { + xyType = CUDA_R_64F; + zType = CUDA_R_64F; + computeType = CUBLAS_COMPUTE_64F; + alpha = reinterpret_cast(&d_alpha); + beta = reinterpret_cast(&d_beta); + } else if constexpr (std::is_same_v) { + xyType = CUDA_R_16F; + if constexpr (std::is_same_v) { + zType = CUDA_R_16F; + computeType = CUBLAS_COMPUTE_16F; + // computeType = CUBLAS_COMPUTE_32F; + alpha = reinterpret_cast(&h_alpha); + beta = reinterpret_cast(&h_beta); + } else if constexpr (std::is_same_v) { + zType = CUDA_R_32F; + computeType = CUBLAS_COMPUTE_32F; + alpha = reinterpret_cast(&f_alpha); + beta = reinterpret_cast(&f_beta); + } + } else if constexpr (std::is_same_v) { + xyType = CUDA_R_8I; + if constexpr (std::is_same_v) { + zType = CUDA_R_32I; + computeType = CUBLAS_COMPUTE_32I; + alpha = reinterpret_cast(&i32_alpha); + beta = reinterpret_cast(&i32_beta); + } else if constexpr (std::is_same_v) { + zType = CUDA_R_32F; + computeType = CUBLAS_COMPUTE_32F; + alpha = reinterpret_cast(&i8_alpha); + beta = reinterpret_cast(&i8_beta); + } + } + auto cublas_h = raft::resource::get_cublas_handle(handle); + RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, + CUBLAS_OP_T, + CUBLAS_OP_N, + N, + M, + K, // Dimensions (swapped due to row/col-major difference) + alpha, // alpha + y, + xyType, + K, // B, its data type, and leading dimension + x, + xyType, + K, // A, its data type, and leading dimension + beta, // beta + workspace, + zType, + N, // C, its data type, and leading dimension + computeType, // Computation type + CUBLAS_GEMM_DEFAULT // Algorithm selection + )); + + if (reduce == true) { + if (metric == DistanceType::L2Expanded) { + reduce_min( + out, workspace, x_norm, y_norm, M, N, stream, is_sqrt, initOutBuffer); + } else if (metric == DistanceType::L2SqrtExpanded) { + reduce_min( + out, workspace, x_norm, y_norm, M, N, stream, is_sqrt, initOutBuffer); + } else if (metric == DistanceType::CosineExpanded) { + reduce_min( + out, workspace, x_norm, y_norm, M, N, stream, is_sqrt, initOutBuffer); + } + } +} +// + +/** + * @brief Wrapper around fusedDistanceNN with minimum reduction operators. + * + * fusedDistanceNN cannot be compiled in the distance library due to the lambda + * operators, so this wrapper covers the most common case (minimum). + * + * @tparam DataT data type + * @tparam OutT output type to either store 1-NN indices and their minimum + * distances (e.g. raft::KeyValuePair) or store only the min + * distances. + * @tparam IdxT indexing arithmetic type + * @param[out] min will contain the reduced output (Length = `m`) + * (on device) + * @param[in] x first matrix. Row major. Dim = `m x k`. + * (on device). + * @param[in] y second matrix. Row major. Dim = `n x k`. + * (on device). + * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). + * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) + * @param[in] m gemm m + * @param[in] n gemm n + * @param[in] k gemm k + * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) + * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt + * @param[in] initOutBuffer whether to initialize the output buffer before the + * main kernel launch + * @param[in] isRowMajor whether the input/output is row or column major. + * @param[in] metric Distance metric to be used (supports L2, cosine) + * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) + * @param[in] stream cuda stream + */ +template +void unfusedDistanceNNMinReduce(raft::resources const& handle, + OutT* min, + const DataT* x, + const DataT* y, + const AccT* xn, + const AccT* yn, + IdxT m, + IdxT n, + IdxT k, + void* workspace, + bool is_sqrt, + bool initOutBuffer, + bool isRowMajor, + DistanceType metric, + float metric_arg, + cudaStream_t stream) +{ + ASSERT(isRowMajor, "unfusedDistanceNN only supports row major inputs"); + unfused_distance_nn( + handle, min, x, y, m, n, k, xn, yn, (AccT*)workspace, is_sqrt, initOutBuffer, metric, stream); +} + +/** @} */ + +} // namespace distance +} // namespace cuvs + +#endif diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9fc620b4cb..e6d1a24b00 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -103,7 +103,7 @@ endfunction() ConfigureTest( NAME NEIGHBORS_TEST PATH neighbors/brute_force.cu neighbors/brute_force_prefiltered.cu neighbors/sparse_brute_force.cu - neighbors/refine.cu + neighbors/refine.cu neighbors/1_nn.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/1_nn.cu b/cpp/tests/neighbors/1_nn.cu new file mode 100644 index 0000000000..faa41a8413 --- /dev/null +++ b/cpp/tests/neighbors/1_nn.cu @@ -0,0 +1,215 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../test_utils.cuh" +#include "1_nn_helper.cuh" + +#include "../../src/distance/fused_distance_nn.cuh" +#include "../../src/distance/unfused_distance_nn.cuh" + +#include +#include + +namespace cuvs::neighbors { + +enum class ImplType { fused, unfused }; + +template +struct NNInputs { + IdxT m; + IdxT n; + IdxT k; + DistanceType metric; + bool sqrt; + uint64_t rng_seed; + double tol; +}; + +template +class NNTest : public ::testing::TestWithParam> { + public: + using OutT = raft::KeyValuePair; + NNTest() + : params_{::testing::TestWithParam>::GetParam()}, + m{params_.m}, + n{params_.n}, + k{params_.k}, + metric{params_.metric}, + sqrt{params_.sqrt}, + stream{raft::resource::get_cuda_stream(handle)}, + x{raft::make_device_matrix(handle, m, k)}, + y{raft::make_device_matrix(handle, n, k)}, + x_norm{raft::make_device_vector(handle, m)}, + y_norm{raft::make_device_vector(handle, n)}, + out{raft::make_device_vector(handle, m)}, + ref_out{raft::make_device_vector(handle, m)} + { + } + + protected: + void SetUp() override + { + raft::random::RngState rng{params_.rng_seed}; + raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); + raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); + + // Pre-compute norms + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, m, stream); + raft::linalg::rowNorm( + y_norm.data_handle(), y.data_handle(), k, m, stream); + + if constexpr (impl == ImplType::fused) { + workspace_size = n * sizeof(IdxT); + } else if constexpr (impl == ImplType::unfused) { + workspace_size = m * n * sizeof(AccT); + } + + // Reset buffers + RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, m * sizeof(OutT))); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + } + + void compute_1nn() + { + raft::device_vector workspace = + raft::make_device_vector(handle, workspace_size); + + ref_nn( + ref_out.data_handle(), x.data_handle(), y.data_handle(), m, n, k, sqrt, metric, stream); + + if constexpr (impl == ImplType::fused) { + if constexpr (std::is_same_v) { + cuvs::distance::fusedDistanceNNMinReduce(out.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + (void*)workspace.data_handle(), + sqrt, + true, + true, + metric, + 0.0, + stream); + } else { + static_assert(sizeof(DataT) == 0, + "fusedDistanceNNMinReduce is not implemented for datatype other than float"); + } + } else if constexpr (impl == ImplType::unfused) { + cuvs::distance::unfusedDistanceNNMinReduce( + handle, + out.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + (AccT*)workspace.data_handle(), + sqrt, + true, + true, + metric, + 0.0, + stream); + } + } + + void compare() + { + vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + ASSERT_TRUE(summary.max_diff < params_.tol) << summary; + } + + private: + raft::resources handle; + rmm::cuda_stream_view stream; + NNInputs params_; + ComparisonSummary summary; + IdxT m; + IdxT n; + IdxT k; + DistanceType metric; + bool sqrt; + raft::device_matrix x; + raft::device_matrix y; + raft::device_vector x_norm; + raft::device_vector y_norm; + raft::device_vector out; + raft::device_vector ref_out; + size_t workspace_size; +}; + +template +const std::vector> input_fp32 = { + {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, + {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, + {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, +}; + +// Test fused implementation with single-precision +typedef NNTest NNTest_fp32_fused; +TEST_P(NNTest_fp32_fused, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp32_fused, ::testing::ValuesIn(input_fp32)); + +// Test unfused implementation with single-precision +typedef NNTest NNTest_fp32_unfused; +TEST_P(NNTest_fp32_unfused, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp32_unfused, ::testing::ValuesIn(input_fp32)); + +template +const std::vector> input_fp16 = { + {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, + {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, + {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, +}; + +// Test unfused implementation with fp16, int8 +// Fused implementation has no support for fp16, int8 so no test for it +typedef NNTest NNTest_fp16_unfused; +TEST_P(NNTest_fp16_unfused, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp16_unfused, ::testing::ValuesIn(input_fp16)); + +template +const std::vector> input_int8 = { + {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, + {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, + {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, +}; + +// Test unfused implementation with fp16, int8 +// Fused implementation has no support for fp16, int8 so no test for it +typedef NNTest NNTest_int8_unfused; +TEST_P(NNTest_int8_unfused, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); +} // namespace cuvs::neighbors diff --git a/cpp/tests/neighbors/1_nn_helper.cuh b/cpp/tests/neighbors/1_nn_helper.cuh new file mode 100644 index 0000000000..52972dbaf5 --- /dev/null +++ b/cpp/tests/neighbors/1_nn_helper.cuh @@ -0,0 +1,206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include +using cuvs::distance::DistanceType; + +namespace cuvs::neighbors { + +template +_RAFT_HOST_DEVICE T max_val() +{ + if constexpr (std::is_same::value) { + return CUDART_MAX_NORMAL_FP16; + } else { + return cuda::std::numeric_limits::max(); + } +} + +template +_RAFT_HOST_DEVICE T min_val() +{ + if constexpr (std::is_same::value) { + return CUDART_MIN_DENORM_FP16; + } else { + return cuda::std::numeric_limits::min(); + } +} + +template +__device__ AccT l2_distance(const DataT* v1, const DataT* v2, IdxT K) +{ + AccT th_dist = AccT(0.0); + for (IdxT k = 0; k < K; k++) { + AccT diff = AccT(v1[k]) - AccT(v2[k]); + th_dist += (diff * diff); + } + return th_dist; +} + +template +__device__ AccT cosine_distance(const DataT* v1, const DataT* v2, IdxT K) +{ + AccT v1_norm = AccT(0.0); + AccT v2_norm = AccT(0.0); + AccT v1v2 = AccT(0.0); + + for (IdxT k = 0; k < K; k++) { + v1_norm += (AccT(v1[k]) * AccT(v1[k])); + v2_norm += (AccT(v2[k]) * AccT(v2[k])); + v1v2 += (AccT(v1[k]) * AccT(v2[k])); + } + + return AccT(1.0) - (v1v2 / (v1_norm * v2_norm)); +} + +// This is a naive implementation of 1-NN computation +template +RAFT_KERNEL ref_nn_kernel( + OutT* out, const DataT* A, const DataT* B, IdxT M, IdxT N, IdxT K, bool sqrt, DistanceType metric) +{ + IdxT tid = threadIdx.x + blockIdx.x * IdxT(blockDim.x); + + for (IdxT m = tid; m < M; m += (blockDim.x * gridDim.x)) { + IdxT min_index = N + 1; + AccT min_dist = max_val(); + + for (IdxT n = 0; n < N; n++) { + AccT dist; + if (metric == DistanceType::L2SqrtExpanded || metric == DistanceType::L2Expanded) { + dist = l2_distance(&A[m * K], &B[n * K], K); + } else if (metric == DistanceType::CosineExpanded) { + dist = cosine_distance(&A[m * K], &B[n * K], K); + } + if (dist < min_dist) { + min_dist = dist; + min_index = n; + } + } + + if constexpr (std::is_fundamental::value) { + static_assert(std::is_same::value, "OutT and AccT are not same type"); + out[m] = AccT(min_dist); + } else { + // output is a raft::KeyValuePair + static_assert(std::is_same>::value, + "OutT is not raft::KeyValuePair<> type"); + out[m].key = IdxT(min_index); + if (sqrt) { + out[m].value = raft::sqrt(AccT(min_dist)); + } else { + out[m].value = AccT(min_dist); + } + } + } +} + +template +void ref_nn(OutT* out, + const DataT* A, + const DataT* B, + IdxT m, + IdxT n, + IdxT k, + bool sqrt, + DistanceType metric, + cudaStream_t stream) +{ + ref_nn_kernel + <<<(m + 127) / 128, 128, 0, stream>>>(out, A, B, m, n, k, sqrt, metric); + return; +} + +// Structure to track comparison failures +class ComparisonSummary { + public: + double max_diff; // Maximum difference found + uint64_t max_diff_index; // where does the maximum difference occur + double max_diff_a; // What was the `a` value at max difference + double max_diff_b; // What was the `b` value at max difference + double acc_diff; // sum of all the diffs + uint64_t n; // How many items are compared + uint64_t n_misses; // How many were wrong + int mutex; // Simple mutex lock for thread synchronization + + void init() + { + max_diff = 0.0; + max_diff_index = 0; + max_diff_a = 0.0; + max_diff_b = 0.0; + acc_diff = 0.0; + n = 0; + n_misses = 0; + } + + void update(double diff, uint64_t index, double a_val, double b_val, bool missed) + { + if (max_diff < diff) { + max_diff = diff; + max_diff_index = index; + max_diff_a = a_val; + max_diff_b = b_val; + } + acc_diff += diff; + n++; + n_misses = missed ? n_misses + 1 : n_misses; + } + + friend std::ostream& operator<<(std::ostream& os, const ComparisonSummary& summary) + { + if (summary.max_diff > 0.0) { + os << "Total compared " << summary.n << std::endl; + os << "Total missed " << summary.n_misses << std::endl; + os << "Average diff: " << summary.acc_diff / summary.n << std::endl; + os << "max_diff: " << summary.max_diff << " (" << summary.max_diff_a << " - " + << summary.max_diff_b << ")" << std::endl; + os << "max_diff_index: " << summary.max_diff_index << std::endl; + } + return os; + } +}; + +template +void vector_compare( + raft::resources const& handle, const OutT* a, const OutT* b, IdxT n, ComparisonSummary& summary) +{ + auto a_h = raft::make_host_vector(n); + auto b_h = raft::make_host_vector(n); + + raft::copy(a_h.data_handle(), a, n, raft::resource::get_cuda_stream(handle)); + + raft::copy(b_h.data_handle(), b, n, raft::resource::get_cuda_stream(handle)); + + summary.init(); + + for (IdxT i = 0; i < n; i++) { + double diff, a_val, b_val; + bool missed = false; + if constexpr (std::is_fundamental_v || std::is_same_v) { + // OutT is float, half or an integer type + a_val = double(a_h(i)); + b_val = double(b_h(i)); + } else { + // OutT is a raft::KeyValuePair + a_val = double(a_h(i).value); + b_val = double(b_h(i).value); + + missed = a_h(i).key != b_h(i).key; + } + + diff = std::abs(a_val - b_val); + summary.update(diff, i, a_val, b_val, missed); + } +} + +} // namespace cuvs::neighbors From 316df9f2aee76b637124088cd5325ae0c82677cc Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 28 Jan 2026 13:53:54 +0000 Subject: [PATCH 02/48] Removing and updating comments --- cpp/bench/prims/src/distance/fused_l2_nn.cu | 5 ++--- cpp/src/distance/unfused_distance_nn.cuh | 14 +++++++------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/cpp/bench/prims/src/distance/fused_l2_nn.cu b/cpp/bench/prims/src/distance/fused_l2_nn.cu index 42905b8224..b150fdebca 100644 --- a/cpp/bench/prims/src/distance/fused_l2_nn.cu +++ b/cpp/bench/prims/src/distance/fused_l2_nn.cu @@ -48,8 +48,6 @@ void benchmark_fusedl2nn(benchmark::State& state) const int m = state.range(0); const int n = state.range(1); const int k = state.range(2); - // const bool sqrt = state.range(3); - // const DistanceType metric = DistanceType(state.range(4)); raft::device_resources handle; rmm::cuda_stream_view stream; @@ -79,7 +77,7 @@ void benchmark_fusedl2nn(benchmark::State& state) // Calculate the workspace size // for fused it is m * sizeof(IdxT) - // for unfused, gemm and tensor it is m * n * sizeof(AccT); + // for unfused and gemm it is m * n * sizeof(AccT); size_t workspace_size = (algo == AlgorithmType::fused) ? m * sizeof(IdxT) : m * n * sizeof(AccT); @@ -238,6 +236,7 @@ int main(int argc, char** argv) IdxT M = 1024; IdxT N = 1024; IdxT K = 128; + for (int i = 0; i < argc; i++) { if (strcmp(argv[i], "-M") == 0) { M = std::stoi(argv[i + 1]); diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 6c6e285289..2bff748ccb 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -248,20 +248,20 @@ void unfused_distance_nn(raft::resources const& handle, RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N, - N, + N, // Dimensions (swapped due to row/col-major difference) M, - K, // Dimensions (swapped due to row/col-major difference) - alpha, // alpha + K, + alpha, y, xyType, - K, // B, its data type, and leading dimension + K, x, xyType, - K, // A, its data type, and leading dimension - beta, // beta + K, + beta, workspace, zType, - N, // C, its data type, and leading dimension + N, computeType, // Computation type CUBLAS_GEMM_DEFAULT // Algorithm selection )); From 4276494ea7e07361af3771880c1b8abe23917cef Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 28 Jan 2026 14:13:37 +0000 Subject: [PATCH 03/48] RNG does not support int8_t, therefore commenting the test case for now --- cpp/tests/neighbors/1_nn.cu | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/cpp/tests/neighbors/1_nn.cu b/cpp/tests/neighbors/1_nn.cu index faa41a8413..bffc3e132b 100644 --- a/cpp/tests/neighbors/1_nn.cu +++ b/cpp/tests/neighbors/1_nn.cu @@ -204,12 +204,12 @@ const std::vector> input_int8 = { // Test unfused implementation with fp16, int8 // Fused implementation has no support for fp16, int8 so no test for it -typedef NNTest NNTest_int8_unfused; -TEST_P(NNTest_int8_unfused, test) -{ - this->compute_1nn(); - this->compare(); -} - -INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); -} // namespace cuvs::neighbors +// typedef NNTest NNTest_int8_unfused; +// TEST_P(NNTest_int8_unfused, test) +// { +// this->compute_1nn(); +// this->compare(); +// } +// +// INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); + } // namespace cuvs::neighbors From c36f25e4d3e871eb2de33fbf3a65e2398d8d0ea3 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 28 Jan 2026 15:12:43 +0000 Subject: [PATCH 04/48] Adding cutlass dependency to the benchmarking code --- cpp/bench/prims/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/bench/prims/CMakeLists.txt b/cpp/bench/prims/CMakeLists.txt index 4852099873..cb12e405e0 100644 --- a/cpp/bench/prims/CMakeLists.txt +++ b/cpp/bench/prims/CMakeLists.txt @@ -111,7 +111,7 @@ endif() if(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN) ConfigurePrimsBench( NAME FUSED_L2_NN PATH src/distance/fused_l2_nn.cu - LINKS CUDA::cublas rmm::rmm + LINKS CUDA::cublas rmm::rmm $ ) endif() From f5ed1c79af8fb0e737bd8a4194ffe0c893689fb2 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 29 Jan 2026 13:17:48 +0000 Subject: [PATCH 05/48] Removing unused API declaration --- cpp/bench/prims/src/distance/fused_l2_nn.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/bench/prims/src/distance/fused_l2_nn.cu b/cpp/bench/prims/src/distance/fused_l2_nn.cu index b150fdebca..714681964b 100644 --- a/cpp/bench/prims/src/distance/fused_l2_nn.cu +++ b/cpp/bench/prims/src/distance/fused_l2_nn.cu @@ -19,7 +19,6 @@ using cuvs::distance::DistanceType; using cuvs::distance::fusedDistanceNNMinReduce; using cuvs::distance::reduce_min; -using cuvs::distance::unfused_distance_nn; using cuvs::distance::unfusedDistanceNNMinReduce; enum class AlgorithmType { gemm, unfused, fused }; From 268a50552bfe6b14ea8b8d872a641bf51860cdea Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 29 Jan 2026 13:20:43 +0000 Subject: [PATCH 06/48] Fixing Blackwell's SM version --- cpp/src/cluster/detail/kmeans_balanced.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 171786e4ad..62ac7e30f2 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -66,7 +66,7 @@ bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) } else if (prop.major == 9 && (m >= 4096 || n >= 4096)) { // On Hopper if m, n are bigger than 4096, use fused return true; - } else if (prop.major >= 12) { + } else if (prop.major >= 10) { // On Blackwell onwards, use unfused return false; } From 381653f2c253c9886242ea2d7f923dd345d5017d Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 29 Jan 2026 13:22:02 +0000 Subject: [PATCH 07/48] Removing debugging and profiling code --- cpp/src/cluster/detail/kmeans_balanced.cuh | 49 ---------------------- 1 file changed, 49 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 62ac7e30f2..46d2542e26 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -99,33 +99,6 @@ bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) #include #include -static void add(int64_t i, int64_t j, int64_t k) -{ - // Static vector of pairs: (tuple, counter) - static std::vector, int>> tupleList; - - std::tuple tuple = std::make_tuple(i, j, k); - - auto it = std::find_if( - tupleList.begin(), tupleList.end(), [&](const auto& p) { return p.first == tuple; }); - if (it != tupleList.end()) { - // Tuple exists, increment counter - it->second += 1; - } else { - // New tuple, add with counter 1 - tupleList.emplace_back(tuple, 1); - } - /*FILE* outFile; - - outFile = fopen("tuples.txt", "w"); - for (const auto& entry : tupleList) { - int64_t _i, _j, _k; - std::tie(_i, _j, _k) = entry.first; - fprintf(outFile, "%d,%lu,%lu,%lu\n", entry.second, _i, _j, _k); - } - fclose(outFile); */ -} - template inline std::enable_if_t> predict_core( const raft::resources& handle, @@ -139,49 +112,31 @@ inline std::enable_if_t> predict_core( LabelT* labels, rmm::device_async_resource_ref mr) { - std::string result = "predict_core_" + std::to_string(n_rows) + "_" + std::to_string(n_clusters) + - "_" + std::to_string(dim) + "_" + std::to_string(sizeof(MathT)); - nvtxRangePushA(result.c_str()); auto stream = raft::resource::get_cuda_stream(handle); bool should_use_fused = use_fused(handle, n_rows, n_clusters, dim); switch (params.metric) { case cuvs::distance::DistanceType::L2Expanded: case cuvs::distance::DistanceType::L2SqrtExpanded: { - nvtxRangePushA("pc::workspace_alloc"); size_t workspace_size = should_use_fused ? sizeof(IdxT) * n_rows : sizeof(MathT) * n_rows * n_clusters; auto workspace = raft::make_device_mdarray(handle, mr, raft::make_extents(workspace_size)); - nvtxRangePop(); - nvtxRangePushA("pc::output_alloc"); auto minClusterAndDistance = raft::make_device_mdarray, IdxT>( handle, mr, raft::make_extents(n_rows)); - nvtxRangePop(); - /*auto unf_minClusterAndDistance = - raft::make_device_mdarray, IdxT>( - handle, mr, raft::make_extents(n_rows));*/ - nvtxRangePushA("pc::output_init"); raft::KeyValuePair initial_value(0, std::numeric_limits::max()); thrust::fill(raft::resource::get_thrust_policy(handle), minClusterAndDistance.data_handle(), minClusterAndDistance.data_handle() + minClusterAndDistance.size(), initial_value); - nvtxRangePop(); - nvtxRangePushA("pc::row_norm"); auto centroidsNorm = raft::make_device_mdarray(handle, mr, raft::make_extents(n_clusters)); raft::linalg::rowNorm( centroidsNorm.data_handle(), centers, dim, n_clusters, stream); - add(int64_t(n_rows), int64_t(n_clusters), int64_t(dim)); - nvtxRangePop(); - static int print_guard = 0; if (should_use_fused) { - if (print_guard < 5) printf("fused called\n"); - nvtxRangePushA("pc::fused_call"); cuvs::distance::fusedDistanceNNMinReduce, IdxT>( minClusterAndDistance.data_handle(), dataset, @@ -198,9 +153,7 @@ inline std::enable_if_t> predict_core( params.metric, 0.0f, stream); - nvtxRangePop(); } else { - if (print_guard < 5) printf("unfused called\n"); cuvs::distance:: unfusedDistanceNNMinReduce, IdxT>( handle, @@ -220,7 +173,6 @@ inline std::enable_if_t> predict_core( 0.0f, stream); } - print_guard++; // todo(lsugy): use KVP + iterator in caller. // Copy keys to output labels @@ -305,7 +257,6 @@ inline std::enable_if_t> predict_core( RAFT_FAIL("The chosen distance metric is not supported (%d)", int(params.metric)); } } - nvtxRangePop(); } /** From 7a49becde48d673d7e3dbc093bc4d4c67cc9c832 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 29 Jan 2026 13:24:48 +0000 Subject: [PATCH 08/48] Splitting the GEMM and reduction calls --- cpp/src/distance/unfused_distance_nn.cuh | 65 +++++++++++------------- cpp/tests/neighbors/1_nn.cu | 2 +- 2 files changed, 32 insertions(+), 35 deletions(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 2bff748ccb..fff950748e 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -171,24 +171,17 @@ void reduce_min(OutT* out, } template -void unfused_distance_nn(raft::resources const& handle, - OutT* out, - const DataT* x, - const DataT* y, - IdxT M, - IdxT N, - IdxT K, - const AccT* x_norm, - const AccT* y_norm, - AccT* workspace, - bool is_sqrt, - bool initOutBuffer, - DistanceType metric, - cudaStream_t stream) +void pairwise_distance_gemm(raft::resources const& handle, + AccT* z, + const DataT* x, + const DataT* y, + IdxT M, + IdxT N, + IdxT K, + const AccT* x_norm, + const AccT* y_norm, + cudaStream_t stream) { - assert((metric == DistanceType::CosineExpanded) || (metric == DistanceType::L2Expanded) || - (metric == DistanceType::L2SqrtExpanded)); - cudaDataType_t xyType, zType; cublasComputeType_t computeType; @@ -248,7 +241,7 @@ void unfused_distance_nn(raft::resources const& handle, RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N, - N, // Dimensions (swapped due to row/col-major difference) + N, // Dimensions (swapped due to row/col-major difference) M, K, alpha, @@ -259,25 +252,12 @@ void unfused_distance_nn(raft::resources const& handle, xyType, K, beta, - workspace, + z, zType, N, computeType, // Computation type CUBLAS_GEMM_DEFAULT // Algorithm selection )); - - if (reduce == true) { - if (metric == DistanceType::L2Expanded) { - reduce_min( - out, workspace, x_norm, y_norm, M, N, stream, is_sqrt, initOutBuffer); - } else if (metric == DistanceType::L2SqrtExpanded) { - reduce_min( - out, workspace, x_norm, y_norm, M, N, stream, is_sqrt, initOutBuffer); - } else if (metric == DistanceType::CosineExpanded) { - reduce_min( - out, workspace, x_norm, y_norm, M, N, stream, is_sqrt, initOutBuffer); - } - } } // @@ -331,8 +311,25 @@ void unfusedDistanceNNMinReduce(raft::resources const& handle, cudaStream_t stream) { ASSERT(isRowMajor, "unfusedDistanceNN only supports row major inputs"); - unfused_distance_nn( - handle, min, x, y, m, n, k, xn, yn, (AccT*)workspace, is_sqrt, initOutBuffer, metric, stream); + pairwise_distance_gemm( + handle, (AccT*)workspace, x, y, m, n, k, xn, yn, stream); + + ASSERT((metric == DistanceType::CosineExpanded) || (metric == DistanceType::L2Expanded) || + (metric == DistanceType::L2SqrtExpanded), + "Only cosine and L2 distance types are supported"); + + if (reduce == true) { + if (metric == DistanceType::L2Expanded) { + reduce_min( + min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + } else if (metric == DistanceType::L2SqrtExpanded) { + reduce_min( + min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + } else if (metric == DistanceType::CosineExpanded) { + reduce_min( + min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + } + } } /** @} */ diff --git a/cpp/tests/neighbors/1_nn.cu b/cpp/tests/neighbors/1_nn.cu index bffc3e132b..47b6788eb4 100644 --- a/cpp/tests/neighbors/1_nn.cu +++ b/cpp/tests/neighbors/1_nn.cu @@ -212,4 +212,4 @@ const std::vector> input_int8 = { // } // // INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); - } // namespace cuvs::neighbors +} // namespace cuvs::neighbors From f5a8a7915087216bdb7a6d0763bb86e021c200dc Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 29 Jan 2026 13:47:01 +0000 Subject: [PATCH 09/48] Updating the name of distance GEMM cal --- cpp/bench/prims/src/distance/fused_l2_nn.cu | 27 +++++++++------------ cpp/src/distance/unfused_distance_nn.cuh | 26 +++++++++----------- 2 files changed, 23 insertions(+), 30 deletions(-) diff --git a/cpp/bench/prims/src/distance/fused_l2_nn.cu b/cpp/bench/prims/src/distance/fused_l2_nn.cu index 714681964b..22e87c29d9 100644 --- a/cpp/bench/prims/src/distance/fused_l2_nn.cu +++ b/cpp/bench/prims/src/distance/fused_l2_nn.cu @@ -18,6 +18,7 @@ using cuvs::distance::DistanceType; using cuvs::distance::fusedDistanceNNMinReduce; +using cuvs::distance::pairwise_distance_gemm; using cuvs::distance::reduce_min; using cuvs::distance::unfusedDistanceNNMinReduce; @@ -146,22 +147,16 @@ void benchmark_fusedl2nn(benchmark::State& state) } if constexpr (algo == AlgorithmType::gemm) { - unfusedDistanceNNMinReduce(handle, - out.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - static_cast(m), - static_cast(n), - static_cast(k), - (AccT*)workspace.data_handle(), - sqrt, - true, - true, - metric, - float(0.0), - stream); + pairwise_distance_gemm(handle, + (AccT*)workspace.data_handle(), + x.data_handle(), + y.data_handle(), + static_cast(m), + static_cast(n), + static_cast(k), + x_norm.data_handle(), + y_norm.data_handle(), + stream); } } RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index fff950748e..1585b1a294 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -170,7 +170,7 @@ void reduce_min(OutT* out, <<>>(out, z, x_norm, y_norm, m, n, is_sqrt, initOutBuffer); } -template +template void pairwise_distance_gemm(raft::resources const& handle, AccT* z, const DataT* x, @@ -292,7 +292,7 @@ void pairwise_distance_gemm(raft::resources const& handle, * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) * @param[in] stream cuda stream */ -template +template void unfusedDistanceNNMinReduce(raft::resources const& handle, OutT* min, const DataT* x, @@ -311,24 +311,22 @@ void unfusedDistanceNNMinReduce(raft::resources const& handle, cudaStream_t stream) { ASSERT(isRowMajor, "unfusedDistanceNN only supports row major inputs"); - pairwise_distance_gemm( + pairwise_distance_gemm( handle, (AccT*)workspace, x, y, m, n, k, xn, yn, stream); ASSERT((metric == DistanceType::CosineExpanded) || (metric == DistanceType::L2Expanded) || (metric == DistanceType::L2SqrtExpanded), "Only cosine and L2 distance types are supported"); - if (reduce == true) { - if (metric == DistanceType::L2Expanded) { - reduce_min( - min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); - } else if (metric == DistanceType::L2SqrtExpanded) { - reduce_min( - min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); - } else if (metric == DistanceType::CosineExpanded) { - reduce_min( - min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); - } + if (metric == DistanceType::L2Expanded) { + reduce_min( + min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + } else if (metric == DistanceType::L2SqrtExpanded) { + reduce_min( + min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + } else if (metric == DistanceType::CosineExpanded) { + reduce_min( + min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); } } From 5a4a7957fd281fc7df57f540a1e7fd9ecb2dd47e Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 29 Jan 2026 13:53:52 +0000 Subject: [PATCH 10/48] Changing the file name to reflect the benchmark better --- .../prims/src/distance/{fused_l2_nn.cu => fused_distance_nn.cu} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename cpp/bench/prims/src/distance/{fused_l2_nn.cu => fused_distance_nn.cu} (100%) diff --git a/cpp/bench/prims/src/distance/fused_l2_nn.cu b/cpp/bench/prims/src/distance/fused_distance_nn.cu similarity index 100% rename from cpp/bench/prims/src/distance/fused_l2_nn.cu rename to cpp/bench/prims/src/distance/fused_distance_nn.cu From 1b372feb248c9dcf6a2ea16f05f6d0dd493bfac3 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 2 Feb 2026 13:29:05 +0000 Subject: [PATCH 11/48] Updating the name of benchmarking file --- cpp/bench/prims/src/distance/{fused_distance_nn.cu => 1_nn.cu} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename cpp/bench/prims/src/distance/{fused_distance_nn.cu => 1_nn.cu} (100%) diff --git a/cpp/bench/prims/src/distance/fused_distance_nn.cu b/cpp/bench/prims/src/distance/1_nn.cu similarity index 100% rename from cpp/bench/prims/src/distance/fused_distance_nn.cu rename to cpp/bench/prims/src/distance/1_nn.cu From 9e9e912fcb7443fc2a54154b55cb7eb8431d6355 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 2 Feb 2026 13:40:53 +0000 Subject: [PATCH 12/48] Updating test and bench file names --- cpp/bench/prims/CMakeLists.txt | 2 +- cpp/bench/prims/src/distance/{1_nn.cu => distance_nn.cu} | 0 cpp/tests/CMakeLists.txt | 2 +- cpp/tests/neighbors/{1_nn.cu => distance_nn.cu} | 0 cpp/tests/neighbors/{1_nn_helper.cuh => distance_nn_helper.cuh} | 0 5 files changed, 2 insertions(+), 2 deletions(-) rename cpp/bench/prims/src/distance/{1_nn.cu => distance_nn.cu} (100%) rename cpp/tests/neighbors/{1_nn.cu => distance_nn.cu} (100%) rename cpp/tests/neighbors/{1_nn_helper.cuh => distance_nn_helper.cuh} (100%) diff --git a/cpp/bench/prims/CMakeLists.txt b/cpp/bench/prims/CMakeLists.txt index cb12e405e0..cee82cd42a 100644 --- a/cpp/bench/prims/CMakeLists.txt +++ b/cpp/bench/prims/CMakeLists.txt @@ -110,7 +110,7 @@ endif() if(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN) ConfigurePrimsBench( - NAME FUSED_L2_NN PATH src/distance/fused_l2_nn.cu + NAME FUSED_L2_NN PATH src/distance/distance_nn.cu LINKS CUDA::cublas rmm::rmm $ ) endif() diff --git a/cpp/bench/prims/src/distance/1_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu similarity index 100% rename from cpp/bench/prims/src/distance/1_nn.cu rename to cpp/bench/prims/src/distance/distance_nn.cu diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index e6d1a24b00..e7bdaa1916 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -103,7 +103,7 @@ endfunction() ConfigureTest( NAME NEIGHBORS_TEST PATH neighbors/brute_force.cu neighbors/brute_force_prefiltered.cu neighbors/sparse_brute_force.cu - neighbors/refine.cu neighbors/1_nn.cu + neighbors/refine.cu neighbors/distance_nn.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/1_nn.cu b/cpp/tests/neighbors/distance_nn.cu similarity index 100% rename from cpp/tests/neighbors/1_nn.cu rename to cpp/tests/neighbors/distance_nn.cu diff --git a/cpp/tests/neighbors/1_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh similarity index 100% rename from cpp/tests/neighbors/1_nn_helper.cuh rename to cpp/tests/neighbors/distance_nn_helper.cuh From 8aff3dd568608e1469fa95ec9b0b27f231c29b99 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 3 Feb 2026 14:17:48 +0000 Subject: [PATCH 13/48] Changed to manual time to capture correct GPU time --- cpp/bench/prims/CMakeLists.txt | 2 +- cpp/bench/prims/src/distance/distance_nn.cu | 296 +++++++++++--------- cpp/tests/neighbors/distance_nn.cu | 2 +- 3 files changed, 162 insertions(+), 138 deletions(-) diff --git a/cpp/bench/prims/CMakeLists.txt b/cpp/bench/prims/CMakeLists.txt index cee82cd42a..38f35a9c8f 100644 --- a/cpp/bench/prims/CMakeLists.txt +++ b/cpp/bench/prims/CMakeLists.txt @@ -110,7 +110,7 @@ endif() if(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN) ConfigurePrimsBench( - NAME FUSED_L2_NN PATH src/distance/distance_nn.cu + NAME DISTANCE_NN PATH src/distance/distance_nn.cu LINKS CUDA::cublas rmm::rmm $ ) endif() diff --git a/cpp/bench/prims/src/distance/distance_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu index 22e87c29d9..9581bac946 100644 --- a/cpp/bench/prims/src/distance/distance_nn.cu +++ b/cpp/bench/prims/src/distance/distance_nn.cu @@ -108,7 +108,12 @@ void benchmark_fusedl2nn(benchmark::State& state) RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, m * sizeof(OutT), stream)); RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + cudaEvent_t start, stop; + cudaEventCreate(&start); + cudaEventCreate(&stop); + for (auto _ : state) { + cudaEventRecord(start, 0); if constexpr (algo == AlgorithmType::fused) { fusedDistanceNNMinReduce(out.data_handle(), x.data_handle(), @@ -158,44 +163,51 @@ void benchmark_fusedl2nn(benchmark::State& state) y_norm.data_handle(), stream); } - } - RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); - if constexpr (algo == AlgorithmType::gemm) { - if (metric == DistanceType::L2Expanded) { - reduce_min(out.data_handle(), - (AccT*)workspace.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - m, - n, - stream, - sqrt, - true); - } else if (metric == DistanceType::L2SqrtExpanded) { - reduce_min( - out.data_handle(), - (AccT*)workspace.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - m, - n, - stream, - sqrt, - true); - } else if (metric == DistanceType::CosineExpanded) { - reduce_min( - out.data_handle(), - (AccT*)workspace.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - m, - n, - stream, - sqrt, - true); - } + cudaEventRecord(stop, stream); + cudaEventSynchronize(stop); // wait until kernel is done + float ms = 0.0f; + cudaEventElapsedTime(&ms, start, stop); + state.SetIterationTime(ms / 1000.0); } + cudaEventDestroy(start); + cudaEventDestroy(stop); + // if constexpr (algo == AlgorithmType::gemm) { + // if (metric == DistanceType::L2Expanded) { + // reduce_min(out.data_handle(), + // (AccT*)workspace.data_handle(), + // x_norm.data_handle(), + // y_norm.data_handle(), + // m, + // n, + // stream, + // sqrt, + // true); + // } else if (metric == DistanceType::L2SqrtExpanded) { + // reduce_min( + // out.data_handle(), + // (AccT*)workspace.data_handle(), + // x_norm.data_handle(), + // y_norm.data_handle(), + // m, + // n, + // stream, + // sqrt, + // true); + // } else if (metric == DistanceType::CosineExpanded) { + // reduce_min( + // out.data_handle(), + // (AccT*)workspace.data_handle(), + // x_norm.data_handle(), + // y_norm.data_handle(), + // m, + // n, + // stream, + // sqrt, + // true); + // } + // } + state.counters["M"] = m; state.counters["N"] = n; state.counters["K"] = k; @@ -224,46 +236,16 @@ static void CustomArguments(benchmark::internal::Benchmark* b) // b->Args({65536, 10000, 768, false}); } -int main(int argc, char** argv) +template +static void add_all_configs(IdxT M, IdxT N, IdxT K) { - using IdxT = int64_t; - IdxT M = 1024; - IdxT N = 1024; - IdxT K = 128; - - for (int i = 0; i < argc; i++) { - if (strcmp(argv[i], "-M") == 0) { - M = std::stoi(argv[i + 1]); - for (int j = i; j < argc - 2; j++) { - argv[j] = argv[j + 2]; - } - argc -= 2; - i -= 2; - } else if (strcmp(argv[i], "-N") == 0) { - N = std::stoi(argv[i + 1]); - for (int j = i; j < argc - 2; j++) { - argv[j] = argv[j + 2]; - } - argc -= 2; - i -= 2; - } else if (strcmp(argv[i], "-K") == 0) { - K = std::stoi(argv[i + 1]); - for (int j = i; j < argc - 2; j++) { - argv[j] = argv[j + 2]; - } - argc -= 2; - i -= 2; - } - } - - constexpr bool sqrt = false; - constexpr DistanceType dist = DistanceType::L2Expanded; - benchmark::internal::Benchmark* bench; // In the following instances IdxT is int64_t, it does not seem to have much impact on the // performance OutT is always + constexpr bool sqrt = false; + constexpr DistanceType dist = DistanceType::L2Expanded; // Method: Fused, DataT: float, AccT: float bench = benchmark::RegisterBenchmark("fused/float/float", benchmark_fusedl2nn); - bench->Args({M, N, K}); + bench->Args({M, N, K})->UseManualTime(); // Method: unfused, DataT: float, AccT: float bench = benchmark::RegisterBenchmark("unfused/float/float", @@ -284,7 +266,7 @@ int main(int argc, char** argv) AlgorithmType::unfused, sqrt, dist>); - bench->Args({M, N, K}); + bench->Args({M, N, K})->UseManualTime(); // Method: gemm, DataT: float, AccT: float bench = benchmark::RegisterBenchmark("gemm/float/float", @@ -295,73 +277,115 @@ int main(int argc, char** argv) AlgorithmType::gemm, sqrt, dist>); - bench->Args({M, N, K}); - - // Method: unfused, DataT: half, AccT: float - bench = benchmark::RegisterBenchmark("unfused/half/float", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K}); - - // Method: gemm, DataT: half, AccT: float - bench = benchmark::RegisterBenchmark("gemm/half/float", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K}); - - // Method: unfused, DataT: half, AccT: half - bench = benchmark::RegisterBenchmark("unfused/half/half", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K}); + bench->Args({M, N, K})->UseManualTime(); + + // // Method: unfused, DataT: half, AccT: float + // bench = benchmark::RegisterBenchmark("unfused/half/float", + // benchmark_fusedl2nn, + // int64_t, + // AlgorithmType::unfused, + // sqrt, + // dist>); + // bench->Args({M, N, K}); + + // // Method: gemm, DataT: half, AccT: float + // bench = benchmark::RegisterBenchmark("gemm/half/float", + // benchmark_fusedl2nn, + // int64_t, + // AlgorithmType::gemm, + // sqrt, + // dist>); + // bench->Args({M, N, K}); + + // // Method: unfused, DataT: half, AccT: half + // bench = benchmark::RegisterBenchmark("unfused/half/half", + // benchmark_fusedl2nn, + // int64_t, + // AlgorithmType::unfused, + // sqrt, + // dist>); + // bench->Args({M, N, K}); + + // // Method: gemm, DataT: half, AccT: half + // bench = benchmark::RegisterBenchmark("gemm/half/half", + // benchmark_fusedl2nn, + // int64_t, + // AlgorithmType::gemm, + // sqrt, + // dist>); + // bench->Args({M, N, K}); + + // // Method: unfused, DataT: int8_t, AccT: int32_t + // bench = benchmark::RegisterBenchmark("unfused/int8_t/int32_t", + // benchmark_fusedl2nn, + // int64_t, + // AlgorithmType::unfused, + // sqrt, + // dist>); + // bench->Args({M, N, K}); + + // // Method: gemm, DataT: int8_t, AccT: int32_t + // bench = benchmark::RegisterBenchmark("gemm/int8_t/int32_t", + // benchmark_fusedl2nn, + // int64_t, + // AlgorithmType::gemm, + // sqrt, + // dist>); + // bench->Args({M, N, K}); +} - // Method: gemm, DataT: half, AccT: half - bench = benchmark::RegisterBenchmark("gemm/half/half", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K}); +int main(int argc, char** argv) +{ + using IdxT = int64_t; + IdxT M = 0; + IdxT N = 0; + IdxT K = 0; - // Method: unfused, DataT: int8_t, AccT: int32_t - bench = benchmark::RegisterBenchmark("unfused/int8_t/int32_t", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K}); + for (int i = 0; i < argc; i++) { + if (strcmp(argv[i], "-M") == 0) { + M = std::stoi(argv[i + 1]); + for (int j = i; j < argc - 2; j++) { + argv[j] = argv[j + 2]; + } + argc -= 2; + i -= 2; + } else if (strcmp(argv[i], "-N") == 0) { + N = std::stoi(argv[i + 1]); + for (int j = i; j < argc - 2; j++) { + argv[j] = argv[j + 2]; + } + argc -= 2; + i -= 2; + } else if (strcmp(argv[i], "-K") == 0) { + K = std::stoi(argv[i + 1]); + for (int j = i; j < argc - 2; j++) { + argv[j] = argv[j + 2]; + } + argc -= 2; + i -= 2; + } + } - // Method: gemm, DataT: int8_t, AccT: int32_t - bench = benchmark::RegisterBenchmark("gemm/int8_t/int32_t", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K}); + if (M == 0 && N == 0 && K == 0) { + M = 16 * 1024; + N = 4 * 1024; + K = 128; + add_all_configs(M, N, K); + } else { + add_all_configs(M, N, K); + } // Initialize benchmark ::benchmark::Initialize(&argc, argv); diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 47b6788eb4..25fef42f44 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -4,7 +4,7 @@ */ #include "../test_utils.cuh" -#include "1_nn_helper.cuh" +#include "distance_nn_helper.cuh" #include "../../src/distance/fused_distance_nn.cuh" #include "../../src/distance/unfused_distance_nn.cuh" From 0384953847aecf29b820d6d4efdb8b964161d51a Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 3 Feb 2026 15:13:00 +0000 Subject: [PATCH 14/48] Adding default benchmarking cases and refactoring --- cpp/bench/prims/src/distance/distance_nn.cu | 241 ++++++++------------ 1 file changed, 93 insertions(+), 148 deletions(-) diff --git a/cpp/bench/prims/src/distance/distance_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu index 9581bac946..0c18488804 100644 --- a/cpp/bench/prims/src/distance/distance_nn.cu +++ b/cpp/bench/prims/src/distance/distance_nn.cu @@ -43,7 +43,7 @@ template -void benchmark_fusedl2nn(benchmark::State& state) +void benchmark_distance_nn(benchmark::State& state) { const int m = state.range(0); const int n = state.range(1); @@ -172,41 +172,6 @@ void benchmark_fusedl2nn(benchmark::State& state) cudaEventDestroy(start); cudaEventDestroy(stop); - // if constexpr (algo == AlgorithmType::gemm) { - // if (metric == DistanceType::L2Expanded) { - // reduce_min(out.data_handle(), - // (AccT*)workspace.data_handle(), - // x_norm.data_handle(), - // y_norm.data_handle(), - // m, - // n, - // stream, - // sqrt, - // true); - // } else if (metric == DistanceType::L2SqrtExpanded) { - // reduce_min( - // out.data_handle(), - // (AccT*)workspace.data_handle(), - // x_norm.data_handle(), - // y_norm.data_handle(), - // m, - // n, - // stream, - // sqrt, - // true); - // } else if (metric == DistanceType::CosineExpanded) { - // reduce_min( - // out.data_handle(), - // (AccT*)workspace.data_handle(), - // x_norm.data_handle(), - // y_norm.data_handle(), - // m, - // n, - // stream, - // sqrt, - // true); - // } - // } state.counters["M"] = m; state.counters["N"] = n; @@ -217,27 +182,7 @@ void benchmark_fusedl2nn(benchmark::State& state) } template -static void CustomArguments(benchmark::internal::Benchmark* b) -{ - /*constexpr int K = 1024; - std::vector m_list = {4 * K, 8 * K}; - std::vector n_list = {4 * K, 8 * K, 16 * K}; - // std::vector k_list = {128, 512, 1024, 1536}; - std::vector k_list = {128, 256, 512}; - for (auto k : k_list) { - for (auto m : m_list) { - for (auto n : n_list) { - b->Args({m, n, k}); - } - } - }*/ - // b->Args({65536, 256, 128, true, DistanceType::L2Expanded}); - // b->Args({65536, 256, 128, false, DistanceType::CosineExpanded}); - // b->Args({65536, 10000, 768, false}); -} - -template -static void add_all_configs(IdxT M, IdxT N, IdxT K) +static void register_configs(IdxT M, IdxT N, IdxT K) { benchmark::internal::Benchmark* bench; @@ -248,102 +193,102 @@ static void add_all_configs(IdxT M, IdxT N, IdxT K) constexpr DistanceType dist = DistanceType::L2Expanded; // Method: Fused, DataT: float, AccT: float bench = benchmark::RegisterBenchmark("fused/float/float", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::fused, - sqrt, - dist>); + benchmark_distance_nn, + int64_t, + AlgorithmType::fused, + sqrt, + dist>); bench->Args({M, N, K})->UseManualTime(); // Method: unfused, DataT: float, AccT: float bench = benchmark::RegisterBenchmark("unfused/float/float", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); + benchmark_distance_nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); bench->Args({M, N, K})->UseManualTime(); // Method: gemm, DataT: float, AccT: float bench = benchmark::RegisterBenchmark("gemm/float/float", - benchmark_fusedl2nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); + benchmark_distance_nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K})->UseManualTime(); + + // Method: unfused, DataT: half, AccT: float + bench = benchmark::RegisterBenchmark("unfused/half/float", + benchmark_distance_nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); bench->Args({M, N, K})->UseManualTime(); - // // Method: unfused, DataT: half, AccT: float - // bench = benchmark::RegisterBenchmark("unfused/half/float", - // benchmark_fusedl2nn, - // int64_t, - // AlgorithmType::unfused, - // sqrt, - // dist>); - // bench->Args({M, N, K}); - - // // Method: gemm, DataT: half, AccT: float - // bench = benchmark::RegisterBenchmark("gemm/half/float", - // benchmark_fusedl2nn, - // int64_t, - // AlgorithmType::gemm, - // sqrt, - // dist>); - // bench->Args({M, N, K}); - - // // Method: unfused, DataT: half, AccT: half - // bench = benchmark::RegisterBenchmark("unfused/half/half", - // benchmark_fusedl2nn, - // int64_t, - // AlgorithmType::unfused, - // sqrt, - // dist>); - // bench->Args({M, N, K}); - - // // Method: gemm, DataT: half, AccT: half - // bench = benchmark::RegisterBenchmark("gemm/half/half", - // benchmark_fusedl2nn, - // int64_t, - // AlgorithmType::gemm, - // sqrt, - // dist>); - // bench->Args({M, N, K}); - - // // Method: unfused, DataT: int8_t, AccT: int32_t - // bench = benchmark::RegisterBenchmark("unfused/int8_t/int32_t", - // benchmark_fusedl2nn, - // int64_t, - // AlgorithmType::unfused, - // sqrt, - // dist>); - // bench->Args({M, N, K}); - - // // Method: gemm, DataT: int8_t, AccT: int32_t - // bench = benchmark::RegisterBenchmark("gemm/int8_t/int32_t", - // benchmark_fusedl2nn, - // int64_t, - // AlgorithmType::gemm, - // sqrt, - // dist>); - // bench->Args({M, N, K}); + // Method: gemm, DataT: half, AccT: float + bench = benchmark::RegisterBenchmark("gemm/half/float", + benchmark_distance_nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K})->UseManualTime(); + + // Method: unfused, DataT: half, AccT: half + bench = benchmark::RegisterBenchmark("unfused/half/half", + benchmark_distance_nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); + bench->Args({M, N, K})->UseManualTime(); + + // Method: gemm, DataT: half, AccT: half + bench = benchmark::RegisterBenchmark("gemm/half/half", + benchmark_distance_nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K})->UseManualTime(); + + // Method: unfused, DataT: int8_t, AccT: int32_t + bench = benchmark::RegisterBenchmark("unfused/int8_t/int32_t", + benchmark_distance_nn, + int64_t, + AlgorithmType::unfused, + sqrt, + dist>); + bench->Args({M, N, K})->UseManualTime(); + + // Method: gemm, DataT: int8_t, AccT: int32_t + bench = benchmark::RegisterBenchmark("gemm/int8_t/int32_t", + benchmark_distance_nn, + int64_t, + AlgorithmType::gemm, + sqrt, + dist>); + bench->Args({M, N, K})->UseManualTime(); } int main(int argc, char** argv) @@ -379,12 +324,12 @@ int main(int argc, char** argv) } if (M == 0 && N == 0 && K == 0) { - M = 16 * 1024; - N = 4 * 1024; - K = 128; - add_all_configs(M, N, K); + register_configs(16 * 1024, 4 * 1024, 128); + register_configs(16 * 1024, 4 * 1024, 64); + register_configs(8 * 1024, 2 * 1024, 64); + register_configs(4 * 1024, 1024, 64); } else { - add_all_configs(M, N, K); + register_configs(M, N, K); } // Initialize benchmark From 6aa6b4ad70dce7e2c0fc549b2f9be82833117392 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 3 Feb 2026 15:33:58 +0000 Subject: [PATCH 15/48] Refactoring --- cpp/src/cluster/detail/kmeans_balanced.cuh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 46d2542e26..09f6340693 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -18,7 +18,6 @@ #include #include #include - #include #include #include @@ -55,6 +54,9 @@ namespace cuvs::cluster::kmeans::detail { constexpr static inline float kAdjustCentersWeight = 7.0f; +/** + * @brief Returns true is fused implementation should used. + */ template bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) { @@ -124,7 +126,6 @@ inline std::enable_if_t> predict_core( raft::make_device_mdarray(handle, mr, raft::make_extents(workspace_size)); auto minClusterAndDistance = raft::make_device_mdarray, IdxT>( handle, mr, raft::make_extents(n_rows)); - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); thrust::fill(raft::resource::get_thrust_policy(handle), minClusterAndDistance.data_handle(), From e65cd9a73d5a6ca6d4d3cc84be448a5170839cf5 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 3 Feb 2026 15:56:43 +0000 Subject: [PATCH 16/48] Improving comments and refactoring --- cpp/src/cluster/detail/kmeans_balanced.cuh | 1 - cpp/src/distance/unfused_distance_nn.cuh | 70 +++++----------------- 2 files changed, 14 insertions(+), 57 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 09f6340693..4dce955f4d 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -7,7 +7,6 @@ #include "../../distance/fused_distance_nn.cuh" #include "../../distance/unfused_distance_nn.cuh" - #include "kmeans_common.cuh" #include diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 1585b1a294..e57cce83ce 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -14,51 +14,6 @@ namespace cuvs { namespace distance { -/** - * \ingroup unfused_distance_nn - * @{ - */ -/** - * @brief Unfused distance and 1-nearest-neighbor computation in a single call. - * - * Unlike the fused call, here we do explicit gemm call followed by a reduction call. - * This code path exists because the fused path is sometime not the optimal due to - * GEMM-like kernel not being optimized yet. - * - * - * @tparam DataT data type - * @tparam OutT output type to either store 1-NN indices and their minimum - * distances or store only the min distances. Accordingly, one - * has to pass an appropriate `ReduceOpT` - * @tparam IdxT indexing arithmetic type - * @tparam ReduceOpT A struct to perform the final needed reduction operation - * and also to initialize the output array elements with the - * appropriate initial value needed for reduction. - * @tparam KVPReduceOpT A struct providing functions for key-value pair comparison. - * - * @param[out] min will contain the reduced output (Length = `m`) - * (on device) - * @param[in] x first matrix. Row major. Dim = `m x k`. - * (on device). - * @param[in] y second matrix. Row major. Dim = `n x k`. - * (on device). - * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). - * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) - * @param[in] m gemm m - * @param[in] n gemm n - * @param[in] k gemm k - * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) - * @param[in] redOp reduction operator in the epilogue - * @param[in] pairRedOp reduction operation on key value pairs - * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt - * @param[in] initOutBuffer whether to initialize the output buffer before the - * main kernel launch - * @param[in] isRowMajor whether the input/output is row or column major. - * @param[in] metric Distance metric to be used (supports L2, cosine) - * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) - * @param[in] stream cuda stream - */ - template __host__ __device__ T max_val() { @@ -259,19 +214,22 @@ void pairwise_distance_gemm(raft::resources const& handle, CUBLAS_GEMM_DEFAULT // Algorithm selection )); } -// /** - * @brief Wrapper around fusedDistanceNN with minimum reduction operators. + * \ingroup unfused_distance_nn + * @{ + */ +/** + * @brief Unfused distance and 1-nearest-neighbor computation in a single call. + * + * Unlike the fused implementation, GEMM and reduction kernel are called separately. + * This code path exists because the fused path is sometimes not the optimal choice. * - * fusedDistanceNN cannot be compiled in the distance library due to the lambda - * operators, so this wrapper covers the most common case (minimum). + * @tparam DataT data type + * @tparam OutT output type to either store 1-NN indices and their minimum + * distances or store only the min distances. + * @tparam IdxT indexing arithmetic type * - * @tparam DataT data type - * @tparam OutT output type to either store 1-NN indices and their minimum - * distances (e.g. raft::KeyValuePair) or store only the min - * distances. - * @tparam IdxT indexing arithmetic type * @param[out] min will contain the reduced output (Length = `m`) * (on device) * @param[in] x first matrix. Row major. Dim = `m x k`. @@ -283,8 +241,8 @@ void pairwise_distance_gemm(raft::resources const& handle, * @param[in] m gemm m * @param[in] n gemm n * @param[in] k gemm k - * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) - * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt + * @param[in] workspace temp workspace. Size = sizeof(DataT) * m. (on device) + * @param[in] is_sqrt Whether the output `min` should contain L2-sqrt * @param[in] initOutBuffer whether to initialize the output buffer before the * main kernel launch * @param[in] isRowMajor whether the input/output is row or column major. From d653bce04edba0411e290c9c4657b9ea1ddeb5fc Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 3 Feb 2026 16:38:14 +0000 Subject: [PATCH 17/48] Removing unused headers --- cpp/src/cluster/detail/kmeans_balanced.cuh | 5 ----- 1 file changed, 5 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 4dce955f4d..253bd64c07 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -95,11 +95,6 @@ bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) * @param[out] labels Output predictions [n_rows] * @param[inout] mr (optional) Memory resource to use for temporary allocations */ -#include -#include -#include -#include - template inline std::enable_if_t> predict_core( const raft::resources& handle, From 381ac801b17c63dbbb34ab8ec87f99105aee1306 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 12 Feb 2026 14:50:23 +0000 Subject: [PATCH 18/48] Correcting the constant type for int8_t gemm --- cpp/src/distance/unfused_distance_nn.cuh | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index e57cce83ce..f21647b0db 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -144,7 +144,6 @@ void pairwise_distance_gemm(raft::resources const& handle, const float f_alpha = 1.0, f_beta = 0.0; const double d_alpha = 1.0, d_beta = 0.0; const int32_t i32_alpha = 1, i32_beta = 0; - const int8_t i8_alpha = 1, i8_beta = 0; const void* alpha = nullptr; const void* beta = nullptr; @@ -188,8 +187,8 @@ void pairwise_distance_gemm(raft::resources const& handle, } else if constexpr (std::is_same_v) { zType = CUDA_R_32F; computeType = CUBLAS_COMPUTE_32F; - alpha = reinterpret_cast(&i8_alpha); - beta = reinterpret_cast(&i8_beta); + alpha = reinterpret_cast(&f_alpha); + beta = reinterpret_cast(&f_beta); } } auto cublas_h = raft::resource::get_cublas_handle(handle); From 2874b414e1835b57e5c50969f837d33d9c8550e8 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 12 Feb 2026 14:54:28 +0000 Subject: [PATCH 19/48] Enable int8_t test cases --- cpp/tests/neighbors/distance_nn.cu | 46 +++++++++++++++++++++++------- 1 file changed, 36 insertions(+), 10 deletions(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 25fef42f44..aa6d75142f 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -27,6 +27,18 @@ struct NNInputs { double tol; }; +__global__ void fill_int8(int8_t* buff, int len) +{ + int tid = threadIdx.x + blockIdx.x * blockDim.x; + // Fill the buffer with pseudo-random int8_t values using a simple LCG + if (tid < len) { + // Simple LCG: x_n+1 = (a * x_n + c) % m + // Use tid as seed, constants chosen for decent distribution + int seed = tid * 1103515245 + 12345; + buff[tid] = static_cast((seed >> 16) & 0xFF); + } +} + template class NNTest : public ::testing::TestWithParam> { public: @@ -52,8 +64,13 @@ class NNTest : public ::testing::TestWithParam> { void SetUp() override { raft::random::RngState rng{params_.rng_seed}; - raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); - raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); + if constexpr (std::is_same_v) { + fill_int8<<<1000, 256, 0, stream>>>(x.data_handle(), m * k); + fill_int8<<<1000, 256, 0, stream>>>(y.data_handle(), n * k); + } else { + raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); + raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); + } // Pre-compute norms raft::linalg::rowNorm( @@ -204,12 +221,21 @@ const std::vector> input_int8 = { // Test unfused implementation with fp16, int8 // Fused implementation has no support for fp16, int8 so no test for it -// typedef NNTest NNTest_int8_unfused; -// TEST_P(NNTest_int8_unfused, test) -// { -// this->compute_1nn(); -// this->compare(); -// } -// -// INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); +typedef NNTest NNTest_int8_unfused; +TEST_P(NNTest_int8_unfused, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); + +typedef NNTest NNTest_int8_unfused2; +TEST_P(NNTest_int8_unfused2, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused2, ::testing::ValuesIn(input_int8)); } // namespace cuvs::neighbors From 84dbca1e261cf7e202557569e9a8862dacd99e49 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 13 Feb 2026 12:00:20 +0000 Subject: [PATCH 20/48] Removing redundant constant declarations --- cpp/src/distance/unfused_distance_nn.cuh | 72 ++++++++++-------------- 1 file changed, 29 insertions(+), 43 deletions(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index f21647b0db..1000c12382 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -140,57 +140,43 @@ void pairwise_distance_gemm(raft::resources const& handle, cudaDataType_t xyType, zType; cublasComputeType_t computeType; - const half h_alpha = 1.0, h_beta = 0.0; - const float f_alpha = 1.0, f_beta = 0.0; - const double d_alpha = 1.0, d_beta = 0.0; - const int32_t i32_alpha = 1, i32_beta = 0; - - const void* alpha = nullptr; - const void* beta = nullptr; - if constexpr (std::is_same_v) { - xyType = CUDA_R_32F; - zType = CUDA_R_32F; - computeType = CUBLAS_COMPUTE_32F; - // computeType = CUBLAS_COMPUTE_32F_FAST_TF32; - // computeType = CUBLAS_COMPUTE_32F_FAST_16F; - // computeType = CUBLAS_COMPUTE_32F_FAST_16BF; - // computeType = CUBLAS_COMPUTE_32F_EMULATED_16BFX9; - alpha = reinterpret_cast(&f_alpha); - beta = reinterpret_cast(&f_beta); - } else if constexpr (std::is_same_v) { - xyType = CUDA_R_64F; - zType = CUDA_R_64F; - computeType = CUBLAS_COMPUTE_64F; - alpha = reinterpret_cast(&d_alpha); - beta = reinterpret_cast(&d_beta); + if constexpr (std::is_same_v) { + xyType = CUDA_R_8I; + if constexpr (std::is_same_v) { + zType = CUDA_R_32I; + computeType = CUBLAS_COMPUTE_32I; + } else if constexpr (std::is_same_v) { + zType = CUDA_R_32F; + computeType = CUBLAS_COMPUTE_32F; + } } else if constexpr (std::is_same_v) { xyType = CUDA_R_16F; if constexpr (std::is_same_v) { zType = CUDA_R_16F; computeType = CUBLAS_COMPUTE_16F; - // computeType = CUBLAS_COMPUTE_32F; - alpha = reinterpret_cast(&h_alpha); - beta = reinterpret_cast(&h_beta); + // Alternative: CUBLAS_COMPUTE_32F can be used for higher precision, + // but requires changing alpha and beta to float type instead of half. + // See: https://docs.nvidia.com/cuda/cublas/index.html#cublasgemmex } else if constexpr (std::is_same_v) { zType = CUDA_R_32F; computeType = CUBLAS_COMPUTE_32F; - alpha = reinterpret_cast(&f_alpha); - beta = reinterpret_cast(&f_beta); - } - } else if constexpr (std::is_same_v) { - xyType = CUDA_R_8I; - if constexpr (std::is_same_v) { - zType = CUDA_R_32I; - computeType = CUBLAS_COMPUTE_32I; - alpha = reinterpret_cast(&i32_alpha); - beta = reinterpret_cast(&i32_beta); - } else if constexpr (std::is_same_v) { - zType = CUDA_R_32F; - computeType = CUBLAS_COMPUTE_32F; - alpha = reinterpret_cast(&f_alpha); - beta = reinterpret_cast(&f_beta); } + } else if constexpr (std::is_same_v) { + xyType = CUDA_R_32F; + zType = CUDA_R_32F; + computeType = CUBLAS_COMPUTE_32F_FAST_TF32; + // Note: Alternative compute types that can be used include: + // CUBLAS_COMPUTE_32F, CUBLAS_COMPUTE_32F_FAST_16F, + // CUBLAS_COMPUTE_32F_FAST_16BF, CUBLAS_COMPUTE_32F_EMULATED_16BFX9 + } else if constexpr (std::is_same_v) { + xyType = CUDA_R_64F; + zType = CUDA_R_64F; + computeType = CUBLAS_COMPUTE_64F; } + + const AccT alpha = static_cast(1); + const AccT beta = static_cast(0); + auto cublas_h = raft::resource::get_cublas_handle(handle); RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, CUBLAS_OP_T, @@ -198,14 +184,14 @@ void pairwise_distance_gemm(raft::resources const& handle, N, // Dimensions (swapped due to row/col-major difference) M, K, - alpha, + reinterpret_cast(&alpha), y, xyType, K, x, xyType, K, - beta, + reinterpret_cast(&beta), z, zType, N, From 36f346ba1ea3fdb252660d3ee8834445b8f9035e Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 13 Feb 2026 15:39:13 +0000 Subject: [PATCH 21/48] Replacing the stream sync call --- cpp/tests/neighbors/distance_nn.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index aa6d75142f..4674dd59c3 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -11,6 +11,7 @@ #include #include +#include namespace cuvs::neighbors { @@ -86,7 +87,7 @@ class NNTest : public ::testing::TestWithParam> { // Reset buffers RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, m * sizeof(OutT))); - RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + raft::resource::sync_stream(handle, stream); } void compute_1nn() From 6754290ed49ad1501b7b15f7aa17a9ba2f0a499a Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 13 Feb 2026 15:39:30 +0000 Subject: [PATCH 22/48] Formatting --- cpp/src/distance/unfused_distance_nn.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 1000c12382..a6c978971f 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -175,7 +175,7 @@ void pairwise_distance_gemm(raft::resources const& handle, } const AccT alpha = static_cast(1); - const AccT beta = static_cast(0); + const AccT beta = static_cast(0); auto cublas_h = raft::resource::get_cublas_handle(handle); RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, From aefe91bd77b7db7b5013b4284aae572f6c6aeaee Mon Sep 17 00:00:00 2001 From: Vinay Deshpande Date: Thu, 19 Feb 2026 15:15:50 +0000 Subject: [PATCH 23/48] Apply suggestions from code review Co-authored-by: Dante Gama Dessavre --- cpp/src/cluster/detail/kmeans_balanced.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 408693dcf8..722e420b23 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -142,7 +142,7 @@ inline std::enable_if_t> predict_core( n_clusters, dim, (void*)workspace.data_handle(), - (params.metric == cuvs::distance::DistanceType::L2Expanded) ? true : false, + (params.metric == cuvs::distance::DistanceType::L2Expanded) ? false : true, false, true, params.metric, From 5c31c28bb8c1f6a8bcaa37eb27ccb99d97cbb131 Mon Sep 17 00:00:00 2001 From: Vinay Deshpande Date: Thu, 19 Feb 2026 15:17:53 +0000 Subject: [PATCH 24/48] Update cpp/bench/prims/src/distance/distance_nn.cu Co-authored-by: Dante Gama Dessavre --- cpp/bench/prims/src/distance/distance_nn.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/bench/prims/src/distance/distance_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu index 0c18488804..ec215f889f 100644 --- a/cpp/bench/prims/src/distance/distance_nn.cu +++ b/cpp/bench/prims/src/distance/distance_nn.cu @@ -113,7 +113,7 @@ void benchmark_distance_nn(benchmark::State& state) cudaEventCreate(&stop); for (auto _ : state) { - cudaEventRecord(start, 0); + cudaEventRecord(start, stream); if constexpr (algo == AlgorithmType::fused) { fusedDistanceNNMinReduce(out.data_handle(), x.data_handle(), From 2c2196a35adf00906be726260f4468d8e4c14080 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 19 Feb 2026 16:14:14 +0000 Subject: [PATCH 25/48] Always assume GPU build --- cpp/bench/prims/CMakeLists.txt | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/cpp/bench/prims/CMakeLists.txt b/cpp/bench/prims/CMakeLists.txt index 38f35a9c8f..bf94096ae5 100644 --- a/cpp/bench/prims/CMakeLists.txt +++ b/cpp/bench/prims/CMakeLists.txt @@ -32,10 +32,6 @@ function(ConfigurePrimsBench) set(oneValueArgs NAME) set(multiValueArgs PATH LINKS CXXFLAGS) - #if(NOT BUILD_CPU_ONLY) - set(GPU_BUILD ON) - #endif() - cmake_parse_arguments( ConfigurePrimsBench "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} ) @@ -53,7 +49,7 @@ function(ConfigurePrimsBench) PRIVATE ${ConfigurePrimsBench_LINKS} nlohmann_json::nlohmann_json Threads::Threads - $<$:CUDA::cudart_static> + CUDA::cudart_static $ $ ) From 7e8afc0770bd41df768a8c9245b974dc74c6cc42 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 19 Feb 2026 16:14:44 +0000 Subject: [PATCH 26/48] Add test cases where m != n --- cpp/tests/neighbors/distance_nn.cu | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 4674dd59c3..f375b27fd9 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -168,9 +168,13 @@ class NNTest : public ::testing::TestWithParam> { template const std::vector> input_fp32 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, + {16384, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, + {8192, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, + {4096, 8192, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; // Test fused implementation with single-precision From 1bd407f0bf3966bffcc852b1c3edbe6a086d5c08 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 23 Feb 2026 14:56:14 +0000 Subject: [PATCH 27/48] Replacing with matrix::fill --- cpp/bench/prims/src/distance/distance_nn.cu | 20 +++++++++++++++----- cpp/tests/neighbors/distance_nn.cu | 11 +++++++++-- 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/cpp/bench/prims/src/distance/distance_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu index ec215f889f..ef3be9560d 100644 --- a/cpp/bench/prims/src/distance/distance_nn.cu +++ b/cpp/bench/prims/src/distance/distance_nn.cu @@ -14,6 +14,7 @@ #include #include #include +#include #include using cuvs::distance::DistanceType; @@ -81,8 +82,8 @@ void benchmark_distance_nn(benchmark::State& state) size_t workspace_size = (algo == AlgorithmType::fused) ? m * sizeof(IdxT) : m * n * sizeof(AccT); - raft::device_vector workspace = - raft::make_device_vector(handle, workspace_size); + raft::device_vector workspace = + raft::make_device_vector(handle, workspace_size); // Warm up if constexpr (algo != AlgorithmType::fused) { @@ -104,9 +105,18 @@ void benchmark_distance_nn(benchmark::State& state) stream); } - RAFT_CUDA_TRY(cudaMemsetAsync(workspace.data_handle(), 0, workspace_size, stream)); - RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, m * sizeof(OutT), stream)); - RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + raft::matrix::fill( + handle, + raft::make_device_matrix_view(workspace.data_handle(), workspace_size, size_t(1)), + char(0)); + if constexpr (std::is_same_v>) { + // OutT is a RAFT KeyValuePair + raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0, 0}); + } else { + // OutT is a scalar type + raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0}); + } + raft::resource::sync_stream(handle, stream); cudaEvent_t start, stop; cudaEventCreate(&start); diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f375b27fd9..5fe28c7f21 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -85,8 +85,15 @@ class NNTest : public ::testing::TestWithParam> { workspace_size = m * n * sizeof(AccT); } - // Reset buffers - RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, m * sizeof(OutT))); + // Reset buffer + if constexpr (std::is_same_v>) { + // OutT is a RAFT KeyValuePair + raft::matrix::fill( + handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0, 0}); + } else { + // OutT is a scalar type + raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0}); + } raft::resource::sync_stream(handle, stream); } From bef16b59daaa2e0f4c7578b2b36b7f76addd66f6 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 23 Feb 2026 15:06:46 +0000 Subject: [PATCH 28/48] Using RAFT host device function macro --- cpp/src/distance/unfused_distance_nn.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index a6c978971f..1ee919caa4 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -15,7 +15,7 @@ namespace cuvs { namespace distance { template -__host__ __device__ T max_val() +_RAFT_HOST_DEVICE T max_val() { if constexpr (std::is_same::value) { return CUDART_MAX_NORMAL_FP16; From e6f861b51ad75e066fc95c1e143aae71692d3775 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 23 Feb 2026 15:59:02 +0000 Subject: [PATCH 29/48] Varying some more sizes --- cpp/tests/neighbors/distance_nn.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 5fe28c7f21..4a08799e0f 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -207,9 +207,9 @@ INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp32_unfused, ::testing::ValuesIn(input_f template const std::vector> input_fp16 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, - {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, - {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; // Test unfused implementation with fp16, int8 @@ -226,9 +226,9 @@ INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp16_unfused, ::testing::ValuesIn(input_f template const std::vector> input_int8 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, - {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, - {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; // Test unfused implementation with fp16, int8 From 4cdf49e3d3c4dea7cce5c6841ab9c99ed5cb6cbc Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 23 Feb 2026 19:51:44 +0000 Subject: [PATCH 30/48] Fixing the dimension of vectors --- cpp/tests/neighbors/distance_nn.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 4a08799e0f..bee701844d 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -77,10 +77,10 @@ class NNTest : public ::testing::TestWithParam> { raft::linalg::rowNorm( x_norm.data_handle(), x.data_handle(), k, m, stream); raft::linalg::rowNorm( - y_norm.data_handle(), y.data_handle(), k, m, stream); + y_norm.data_handle(), y.data_handle(), k, n, stream); if constexpr (impl == ImplType::fused) { - workspace_size = n * sizeof(IdxT); + workspace_size = m * sizeof(IdxT); } else if constexpr (impl == ImplType::unfused) { workspace_size = m * n * sizeof(AccT); } From 31d1def27cddfb921692d27839dbcaa606a04a33 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 20 Mar 2026 14:02:38 +0000 Subject: [PATCH 31/48] Deleting warm-up as warm-up runs can be done using google benchmark cli --- cpp/bench/prims/src/distance/distance_nn.cu | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/cpp/bench/prims/src/distance/distance_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu index ef3be9560d..3c4f2860a2 100644 --- a/cpp/bench/prims/src/distance/distance_nn.cu +++ b/cpp/bench/prims/src/distance/distance_nn.cu @@ -85,26 +85,6 @@ void benchmark_distance_nn(benchmark::State& state) raft::device_vector workspace = raft::make_device_vector(handle, workspace_size); - // Warm up - if constexpr (algo != AlgorithmType::fused) { - unfusedDistanceNNMinReduce(handle, - out.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - static_cast(m), - static_cast(n), - static_cast(k), - (AccT*)workspace.data_handle(), - sqrt, - true, - true, - metric, - float(0.0), - stream); - } - raft::matrix::fill( handle, raft::make_device_matrix_view(workspace.data_handle(), workspace_size, size_t(1)), From 0790c565868746c24afb85b6c07a97cb8b71f669 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 20 Mar 2026 14:03:48 +0000 Subject: [PATCH 32/48] CUB block reduce header needs to be exclusively included now --- cpp/src/distance/unfused_distance_nn.cuh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 1ee919caa4..db1d197efe 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -8,6 +8,8 @@ #pragma once +#include + #include #include From 1324a1d095a66bcfbc2310ba1daab8f7dc7cc5cb Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 20 Mar 2026 14:05:23 +0000 Subject: [PATCH 33/48] Adding docstring to explain the reduction function --- cpp/src/distance/unfused_distance_nn.cuh | 26 ++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index db1d197efe..e75571c891 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -109,6 +109,32 @@ __global__ void reduce_min_kernel(OutT* out, } } +/** + * @brief Reduces pairwise distances to find the minimum distance per row. + * + * This function launches a kernel that computes the final distance values from + * the GEMM output and norms, then reduces across columns to find the minimum + * distance (and optionally the index) for each row. Distance computation and + * reduction is fused together in a single kernel thats why we can not use + * RAFT reduction here. + * + * @tparam DataT Input data type + * @tparam AccT Accumulator/distance type + * @tparam OutT Output type (either AccT for distances only, or KeyValuePair for index+distance) + * @tparam IdxT Index type + * @tparam metric Distance metric type (L2Expanded, L2SqrtExpanded, or CosineExpanded) + * + * @param[out] out Output array containing minimum distances (and optionally indices) + * per row. Length = `m`. (on device) + * @param[in] z GEMM output matrix (x * y^T). Dim = `m x n`. (on device) + * @param[in] x_norm Norms of rows in x. Length = `m`. (on device) + * @param[in] y_norm Norms of rows in y. Length = `n`. (on device) + * @param[in] m Number of rows in x (and output) + * @param[in] n Number of rows in y (columns to reduce over) + * @param[in] stream CUDA stream for kernel launch + * @param[in] is_sqrt Whether to apply square root to the final distance + * @param[in] initOutBuffer Whether to initialize the output buffer or merge with existing values + */ template void reduce_min(OutT* out, const AccT* z, From 5e355fa8aea22634117d42e6a2eb74a1cfbf2049 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 20 Mar 2026 14:07:05 +0000 Subject: [PATCH 34/48] Explicity setting stream for cublas gemm --- cpp/src/distance/unfused_distance_nn.cuh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index e75571c891..086ac76054 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -206,6 +206,8 @@ void pairwise_distance_gemm(raft::resources const& handle, const AccT beta = static_cast(0); auto cublas_h = raft::resource::get_cublas_handle(handle); + RAFT_CUBLAS_TRY(cublasSetStream(cublas_h, stream)); + RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N, From 331383493870345cdff08e434fb22e0a1bdddd5f Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 20 Mar 2026 15:03:53 +0000 Subject: [PATCH 35/48] Removing the prims benchmarks for now --- cpp/CMakeLists.txt | 3 +- cpp/bench/prims/CMakeLists.txt | 114 ------- cpp/bench/prims/src/distance/distance_nn.cu | 331 -------------------- 3 files changed, 1 insertion(+), 447 deletions(-) delete mode 100644 cpp/bench/prims/CMakeLists.txt delete mode 100644 cpp/bench/prims/src/distance/distance_nn.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5080a1ee74..3dcad74825 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -50,7 +50,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) option(BUILD_SHARED_LIBS "Build cuvs shared libraries" ON) option(BUILD_TESTS "Build cuvs unit-tests" ON) option(BUILD_C_LIBRARY "Build cuVS C API library" ON) -option(BUILD_CUVS_BENCH "Build cuVS ann and prims benchmarks" OFF) +option(BUILD_CUVS_BENCH "Build cuVS ann benchmarks" OFF) option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON) option(BUILD_MG_ALGOS "Build with multi-GPU support" ON) option(CUDA_ENABLE_KERNELINFO "Enable kernel resource usage info" OFF) @@ -939,5 +939,4 @@ endif() if(BUILD_CUVS_BENCH) add_subdirectory(bench/ann/) - add_subdirectory(bench/prims) endif() diff --git a/cpp/bench/prims/CMakeLists.txt b/cpp/bench/prims/CMakeLists.txt deleted file mode 100644 index bf94096ae5..0000000000 --- a/cpp/bench/prims/CMakeLists.txt +++ /dev/null @@ -1,114 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -list(APPEND CMAKE_MODULE_PATH "${CUVS_SOURCE_DIR}") - -# ################################################################################################## -# * benchmark options ------------------------------------------------------------------------------ - -option(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN "Include Fused L2 NN benchmark" ON) -option(CUVS_PRIMS_BENCH_USE_CUBLAS_SAMPLE "Include cuBLAS sample benchmark" ON) -option(CUVS_PRIMS_BENCH_USE_WMMA "Include WMMA sample benchmark" ON) - -# ################################################################################################## -# * Process options ---------------------------------------------------------- - -find_package(Threads REQUIRED) - -# ################################################################################################## -# * Fetch requirements ------------------------------------------------------------- - -include(cmake/thirdparty/get_nlohmann_json) - -# ################################################################################################## -# * Target function ------------------------------------------------------------- - -function(ConfigurePrimsBench) - - set(oneValueArgs NAME) - set(multiValueArgs PATH LINKS CXXFLAGS) - - cmake_parse_arguments( - ConfigurePrimsBench "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} - ) - - set(BENCH_NAME ${ConfigurePrimsBench_NAME}_PRIMS_BENCH) - - add_executable(${BENCH_NAME} ${ConfigurePrimsBench_PATH}) - target_compile_definitions(${BENCH_NAME} PRIVATE PRIMS_BENCH_BUILD_MAIN) - target_link_libraries( - ${BENCH_NAME} PRIVATE benchmark::benchmark $<$:CUDA::nvtx3> - ) - - target_link_libraries( - ${BENCH_NAME} - PRIVATE ${ConfigurePrimsBench_LINKS} - nlohmann_json::nlohmann_json - Threads::Threads - CUDA::cudart_static - $ - $ - ) - - set_target_properties( - ${BENCH_NAME} - PROPERTIES # set target compile options - CXX_STANDARD 17 - CXX_STANDARD_REQUIRED ON - CUDA_STANDARD 17 - CUDA_STANDARD_REQUIRED ON - POSITION_INDEPENDENT_CODE ON - INTERFACE_POSITION_INDEPENDENT_CODE ON - BUILD_RPATH "\$ORIGIN" - INSTALL_RPATH "\$ORIGIN" - ) - - set(${ConfigurePrimsBench_CXXFLAGS} ${CUVS_CXX_FLAGS} ${ConfigurePrimsBench_CXXFLAGS}) - - target_compile_options( - ${BENCH_NAME} PRIVATE "$<$:${ConfigurePrimsBench_CXXFLAGS}>" - "$<$:${CUVS_CUDA_FLAGS}>" - ) - - if(CUVS_PRIMS_BENCH_USE_${ConfigurePrimsBench_NAME}) - target_compile_definitions( - ${BENCH_NAME} - PUBLIC - CUVS_PRIMS_BENCH_USE_${ConfigurePrimsBench_NAME}=CUVS_PRIMS_BENCH_USE_${ConfigurePrimsBench_NAME} - ) - endif() - - target_include_directories( - ${BENCH_NAME} - PUBLIC "$" - PRIVATE ${ConfigurePrimsBench_INCLUDES} - ) - - install( - TARGETS ${BENCH_NAME} - COMPONENT prims_bench - DESTINATION bin/prims - ) - - add_dependencies(CUVS_PRIMS_BENCH_ALL ${BENCH_NAME}) -endfunction() - -# ################################################################################################## -# * Configure benchmark targets ------------------------------------------------------------- - -if(NOT TARGET CUVS_PRIMS_BENCH_ALL) - add_custom_target(CUVS_PRIMS_BENCH_ALL) -endif() - -if(CUVS_PRIMS_BENCH_USE_FUSED_L2_NN) - ConfigurePrimsBench( - NAME DISTANCE_NN PATH src/distance/distance_nn.cu - LINKS CUDA::cublas rmm::rmm $ - ) -endif() - -message("CUDAToolkit_LIBRARY_DIR: ${CUDAToolkit_LIBRARY_DIR}") diff --git a/cpp/bench/prims/src/distance/distance_nn.cu b/cpp/bench/prims/src/distance/distance_nn.cu deleted file mode 100644 index 3c4f2860a2..0000000000 --- a/cpp/bench/prims/src/distance/distance_nn.cu +++ /dev/null @@ -1,331 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include - -#include - -#include "../../../../src/distance/fused_distance_nn.cuh" -#include "../../../../src/distance/unfused_distance_nn.cuh" - -#include -#include -#include -#include -#include - -using cuvs::distance::DistanceType; -using cuvs::distance::fusedDistanceNNMinReduce; -using cuvs::distance::pairwise_distance_gemm; -using cuvs::distance::reduce_min; -using cuvs::distance::unfusedDistanceNNMinReduce; - -enum class AlgorithmType { gemm, unfused, fused }; - -__global__ void fill_int8(int8_t* buff, int len) -{ - int tid = threadIdx.x + blockIdx.x * blockDim.x; - // Fill the buffer with pseudo-random int8_t values using a simple LCG - if (tid < len) { - // Simple LCG: x_n+1 = (a * x_n + c) % m - // Use tid as seed, constants chosen for decent distribution - int seed = tid * 1103515245 + 12345; - buff[tid] = static_cast((seed >> 16) & 0xFF); - } -} - -template -void benchmark_distance_nn(benchmark::State& state) -{ - const int m = state.range(0); - const int n = state.range(1); - const int k = state.range(2); - - raft::device_resources handle; - rmm::cuda_stream_view stream; - - stream = raft::resource::get_cuda_stream(handle); - - auto x = raft::make_device_matrix(handle, m, k); - auto y = raft::make_device_matrix(handle, n, k); - auto x_norm = raft::make_device_vector(handle, m); - auto y_norm = raft::make_device_vector(handle, n); - auto out = raft::make_device_vector(handle, m); - - raft::random::RngState rng{1234}; - if constexpr (std::is_same_v) { - fill_int8<<<1000, 256, 0, stream>>>(x.data_handle(), m * k); - fill_int8<<<1000, 256, 0, stream>>>(y.data_handle(), n * k); - } else { - raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); - raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); - } - - // Pre-compute norms - raft::linalg::rowNorm( - x_norm.data_handle(), x.data_handle(), k, m, stream); - raft::linalg::rowNorm( - y_norm.data_handle(), y.data_handle(), k, n, stream); - - // Calculate the workspace size - // for fused it is m * sizeof(IdxT) - // for unfused and gemm it is m * n * sizeof(AccT); - - size_t workspace_size = (algo == AlgorithmType::fused) ? m * sizeof(IdxT) : m * n * sizeof(AccT); - - raft::device_vector workspace = - raft::make_device_vector(handle, workspace_size); - - raft::matrix::fill( - handle, - raft::make_device_matrix_view(workspace.data_handle(), workspace_size, size_t(1)), - char(0)); - if constexpr (std::is_same_v>) { - // OutT is a RAFT KeyValuePair - raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0, 0}); - } else { - // OutT is a scalar type - raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0}); - } - raft::resource::sync_stream(handle, stream); - - cudaEvent_t start, stop; - cudaEventCreate(&start); - cudaEventCreate(&stop); - - for (auto _ : state) { - cudaEventRecord(start, stream); - if constexpr (algo == AlgorithmType::fused) { - fusedDistanceNNMinReduce(out.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - static_cast(m), - static_cast(n), - static_cast(k), - (void*)workspace.data_handle(), - sqrt, - true, - true, - metric, - 0.0, - stream); - } - - if constexpr (algo == AlgorithmType::unfused) { - unfusedDistanceNNMinReduce(handle, - out.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - static_cast(m), - static_cast(n), - static_cast(k), - (AccT*)workspace.data_handle(), - sqrt, - true, - true, - metric, - float(0.0), - stream); - } - - if constexpr (algo == AlgorithmType::gemm) { - pairwise_distance_gemm(handle, - (AccT*)workspace.data_handle(), - x.data_handle(), - y.data_handle(), - static_cast(m), - static_cast(n), - static_cast(k), - x_norm.data_handle(), - y_norm.data_handle(), - stream); - } - cudaEventRecord(stop, stream); - cudaEventSynchronize(stop); // wait until kernel is done - float ms = 0.0f; - cudaEventElapsedTime(&ms, start, stop); - state.SetIterationTime(ms / 1000.0); - } - - cudaEventDestroy(start); - cudaEventDestroy(stop); - - state.counters["M"] = m; - state.counters["N"] = n; - state.counters["K"] = k; - int64_t total_ops = int64_t(2) * m * n * k; - state.counters["FLOP/s"] = - benchmark::Counter(total_ops, benchmark::Counter::kIsIterationInvariantRate); -} - -template -static void register_configs(IdxT M, IdxT N, IdxT K) -{ - benchmark::internal::Benchmark* bench; - - // In the following instances IdxT is int64_t, it does not seem to have much impact on the - // performance OutT is always - - constexpr bool sqrt = false; - constexpr DistanceType dist = DistanceType::L2Expanded; - // Method: Fused, DataT: float, AccT: float - bench = benchmark::RegisterBenchmark("fused/float/float", - benchmark_distance_nn, - int64_t, - AlgorithmType::fused, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: unfused, DataT: float, AccT: float - bench = benchmark::RegisterBenchmark("unfused/float/float", - benchmark_distance_nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: gemm, DataT: float, AccT: float - bench = benchmark::RegisterBenchmark("gemm/float/float", - benchmark_distance_nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: unfused, DataT: half, AccT: float - bench = benchmark::RegisterBenchmark("unfused/half/float", - benchmark_distance_nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: gemm, DataT: half, AccT: float - bench = benchmark::RegisterBenchmark("gemm/half/float", - benchmark_distance_nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: unfused, DataT: half, AccT: half - bench = benchmark::RegisterBenchmark("unfused/half/half", - benchmark_distance_nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: gemm, DataT: half, AccT: half - bench = benchmark::RegisterBenchmark("gemm/half/half", - benchmark_distance_nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: unfused, DataT: int8_t, AccT: int32_t - bench = benchmark::RegisterBenchmark("unfused/int8_t/int32_t", - benchmark_distance_nn, - int64_t, - AlgorithmType::unfused, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); - - // Method: gemm, DataT: int8_t, AccT: int32_t - bench = benchmark::RegisterBenchmark("gemm/int8_t/int32_t", - benchmark_distance_nn, - int64_t, - AlgorithmType::gemm, - sqrt, - dist>); - bench->Args({M, N, K})->UseManualTime(); -} - -int main(int argc, char** argv) -{ - using IdxT = int64_t; - IdxT M = 0; - IdxT N = 0; - IdxT K = 0; - - for (int i = 0; i < argc; i++) { - if (strcmp(argv[i], "-M") == 0) { - M = std::stoi(argv[i + 1]); - for (int j = i; j < argc - 2; j++) { - argv[j] = argv[j + 2]; - } - argc -= 2; - i -= 2; - } else if (strcmp(argv[i], "-N") == 0) { - N = std::stoi(argv[i + 1]); - for (int j = i; j < argc - 2; j++) { - argv[j] = argv[j + 2]; - } - argc -= 2; - i -= 2; - } else if (strcmp(argv[i], "-K") == 0) { - K = std::stoi(argv[i + 1]); - for (int j = i; j < argc - 2; j++) { - argv[j] = argv[j + 2]; - } - argc -= 2; - i -= 2; - } - } - - if (M == 0 && N == 0 && K == 0) { - register_configs(16 * 1024, 4 * 1024, 128); - register_configs(16 * 1024, 4 * 1024, 64); - register_configs(8 * 1024, 2 * 1024, 64); - register_configs(4 * 1024, 1024, 64); - } else { - register_configs(M, N, K); - } - - // Initialize benchmark - ::benchmark::Initialize(&argc, argv); - if (::benchmark::ReportUnrecognizedArguments(argc, argv)) return -1; - ::benchmark::RunSpecifiedBenchmarks(); - ::benchmark::Shutdown(); - return 0; -} From 3a4dfafa3ed86c794e6c0c64ea3cc258cd8e2d07 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Fri, 20 Mar 2026 15:12:46 +0000 Subject: [PATCH 36/48] Fixing a spelling error --- cpp/src/distance/unfused_distance_nn.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 086ac76054..d7f4d98321 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -115,7 +115,7 @@ __global__ void reduce_min_kernel(OutT* out, * This function launches a kernel that computes the final distance values from * the GEMM output and norms, then reduces across columns to find the minimum * distance (and optionally the index) for each row. Distance computation and - * reduction is fused together in a single kernel thats why we can not use + * reduction is fused together in a single kernel that is why we can not use * RAFT reduction here. * * @tparam DataT Input data type From 15b242d683ad1a9510ec0ed10d21ca10bb4e73ba Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 09:30:26 +0100 Subject: [PATCH 37/48] Removing unintentional whitespace change --- cpp/src/cluster/detail/kmeans_balanced.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index cbc5092a3b..319c8b82da 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -86,7 +86,6 @@ inline std::enable_if_t> predict_core( rmm::device_async_resource_ref mr) { auto stream = raft::resource::get_cuda_stream(handle); - switch (params.metric) { case cuvs::distance::DistanceType::L2Expanded: case cuvs::distance::DistanceType::L2SqrtExpanded: From 43e7e934e4d579fdd0971047fc4ee4b0c8adb59b Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 10:14:53 +0100 Subject: [PATCH 38/48] Adding a grid-strided loop to cover full array length --- cpp/tests/neighbors/distance_nn.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index bee701844d..64f5370c0f 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -32,7 +32,7 @@ __global__ void fill_int8(int8_t* buff, int len) { int tid = threadIdx.x + blockIdx.x * blockDim.x; // Fill the buffer with pseudo-random int8_t values using a simple LCG - if (tid < len) { + for (int i = tid; i < len; i += blockIdx.x * gridDim.x) { // Simple LCG: x_n+1 = (a * x_n + c) % m // Use tid as seed, constants chosen for decent distribution int seed = tid * 1103515245 + 12345; From d622448928d12d2f39ed1924d372e6a47c239093 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 10:25:29 +0100 Subject: [PATCH 39/48] Fixing the grid-strided loop --- cpp/tests/neighbors/distance_nn.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 64f5370c0f..4c39f88f38 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -32,11 +32,11 @@ __global__ void fill_int8(int8_t* buff, int len) { int tid = threadIdx.x + blockIdx.x * blockDim.x; // Fill the buffer with pseudo-random int8_t values using a simple LCG - for (int i = tid; i < len; i += blockIdx.x * gridDim.x) { + for (int i = tid; i < len; i += blockDim.x * gridDim.x) { // Simple LCG: x_n+1 = (a * x_n + c) % m - // Use tid as seed, constants chosen for decent distribution - int seed = tid * 1103515245 + 12345; - buff[tid] = static_cast((seed >> 16) & 0xFF); + // Use i as seed, constants chosen for decent distribution + int seed = i * 1103515245 + 12345; + buff[i] = static_cast((seed >> 16) & 0xFF); } } From 39b0259e799d9c027866d120a0f45672c9b83c15 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 11:51:31 +0100 Subject: [PATCH 40/48] Adding a seed_offset to generate different x an y vectors --- cpp/tests/neighbors/distance_nn.cu | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 4c39f88f38..de6ccf4c02 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -28,15 +28,13 @@ struct NNInputs { double tol; }; -__global__ void fill_int8(int8_t* buff, int len) +__global__ void fill_int8(int8_t* buff, int len, int seed_offset) { int tid = threadIdx.x + blockIdx.x * blockDim.x; // Fill the buffer with pseudo-random int8_t values using a simple LCG for (int i = tid; i < len; i += blockDim.x * gridDim.x) { - // Simple LCG: x_n+1 = (a * x_n + c) % m - // Use i as seed, constants chosen for decent distribution - int seed = i * 1103515245 + 12345; - buff[i] = static_cast((seed >> 16) & 0xFF); + int hash = (i + seed_offset) * 1103515245 + 12345; + buff[i] = static_cast((hash >> 16) & 0xFF); } } @@ -66,8 +64,8 @@ class NNTest : public ::testing::TestWithParam> { { raft::random::RngState rng{params_.rng_seed}; if constexpr (std::is_same_v) { - fill_int8<<<1000, 256, 0, stream>>>(x.data_handle(), m * k); - fill_int8<<<1000, 256, 0, stream>>>(y.data_handle(), n * k); + fill_int8<<<1000, 256, 0, stream>>>(x.data_handle(), m * k, 0); + fill_int8<<<1000, 256, 0, stream>>>(y.data_handle(), n * k, m * k); } else { raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); From 0e3cb111149c89e112cb7f7df1f13ca09ef44e58 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 11:53:34 +0100 Subject: [PATCH 41/48] Fixing the issue with cosine norm computation --- cpp/tests/neighbors/distance_nn.cu | 17 +++++++++++++---- cpp/tests/neighbors/distance_nn_helper.cuh | 2 ++ 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index de6ccf4c02..76a6cda7db 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -11,6 +11,7 @@ #include #include +#include #include namespace cuvs::neighbors { @@ -77,6 +78,12 @@ class NNTest : public ::testing::TestWithParam> { raft::linalg::rowNorm( y_norm.data_handle(), y.data_handle(), k, n, stream); + // CosineExpanded expects ||x|| not ||x||^2 + if (metric == DistanceType::CosineExpanded) { + raft::linalg::unaryOp(x_norm.data_handle(), x_norm.data_handle(), m, raft::sqrt_op{}, stream); + raft::linalg::unaryOp(y_norm.data_handle(), y_norm.data_handle(), n, raft::sqrt_op{}, stream); + } + if constexpr (impl == ImplType::fused) { workspace_size = m * sizeof(IdxT); } else if constexpr (impl == ImplType::unfused) { @@ -178,8 +185,10 @@ const std::vector> input_fp32 = { {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {8192, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, - {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, - {4096, 8192, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, + // Fused implementation for cosine distance ignores the sqrt paramter, therefore + // commenting the following two tests + // {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, + // {4096, 8192, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; // Test fused implementation with single-precision @@ -229,8 +238,7 @@ const std::vector> input_int8 = { {4096, 16384, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; -// Test unfused implementation with fp16, int8 -// Fused implementation has no support for fp16, int8 so no test for it +// DataT = int8_t, AccT = int32_t typedef NNTest NNTest_int8_unfused; TEST_P(NNTest_int8_unfused, test) { @@ -240,6 +248,7 @@ TEST_P(NNTest_int8_unfused, test) INSTANTIATE_TEST_CASE_P(NNTest, NNTest_int8_unfused, ::testing::ValuesIn(input_int8)); +// DataT = int8_t, AccT = float typedef NNTest NNTest_int8_unfused2; TEST_P(NNTest_int8_unfused2, test) { diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index 52972dbaf5..740c916fcb 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -60,6 +60,8 @@ __device__ AccT cosine_distance(const DataT* v1, const DataT* v2, IdxT K) v1v2 += (AccT(v1[k]) * AccT(v2[k])); } + v1_norm = raft::sqrt(v1_norm); + v2_norm = raft::sqrt(v2_norm); return AccT(1.0) - (v1v2 / (v1_norm * v2_norm)); } From 8472fe1d96a012db91e885f0d632a73910c7c4aa Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 14:55:15 +0100 Subject: [PATCH 42/48] Syncing before comparing --- cpp/tests/neighbors/distance_nn_helper.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index 740c916fcb..edebbe66d8 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -180,8 +180,8 @@ void vector_compare( auto b_h = raft::make_host_vector(n); raft::copy(a_h.data_handle(), a, n, raft::resource::get_cuda_stream(handle)); - raft::copy(b_h.data_handle(), b, n, raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle, raft::resource::get_cuda_stream(handle)); summary.init(); From 6f35fa78204fd30776d787d2beebbd08e6aeb848 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 14:56:02 +0100 Subject: [PATCH 43/48] Adding tests for L2SqrtExpanded --- cpp/tests/neighbors/distance_nn.cu | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 76a6cda7db..f51d26546a 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -183,6 +183,8 @@ const std::vector> input_fp32 = { {16384, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 4096, 64, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {8192, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, // Fused implementation for cosine distance ignores the sqrt paramter, therefore @@ -215,6 +217,8 @@ template const std::vector> input_fp16 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 4096, 64, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {4096, 16384, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; @@ -234,6 +238,8 @@ template const std::vector> input_int8 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, + {4096, 4096, 64, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, + {4096, 16384, 128, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {4096, 16384, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; From c80051ef1475dd0b8746964b3df76d972a3a8792 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 14:57:02 +0100 Subject: [PATCH 44/48] Update batch size calculation --- cpp/src/cluster/detail/kmeans_balanced.cuh | 30 +++++++++++++--------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 319c8b82da..26ce110158 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -172,22 +172,28 @@ inline std::enable_if_t> predict_core( * @return A suggested minibatch size and the expected memory cost per-row (in bytes) */ template -constexpr auto calc_minibatch_size(IdxT n_clusters, - IdxT n_rows, - IdxT dim, - cuvs::distance::DistanceType metric, - bool needs_conversion) -> std::tuple +auto calc_minibatch_size(const raft::resources& handle, + IdxT n_clusters, + IdxT n_rows, + IdxT dim, + cuvs::distance::DistanceType metric, + bool needs_conversion) -> std::tuple { n_clusters = std::max(1, n_clusters); // Estimate memory needs per row (i.e element of the batch). size_t mem_per_row = 0; switch (metric) { - // fusedL2NN needs a mutex and a key-value pair for each row. case distance::DistanceType::L2Expanded: case distance::DistanceType::L2SqrtExpanded: { - mem_per_row += sizeof(int); - mem_per_row += sizeof(raft::KeyValuePair); + if (use_fused(handle, n_rows, n_clusters, dim)) { + // fusedL2NN needs a mutex and a key-value pair for each row. + mem_per_row += sizeof(int); + mem_per_row += sizeof(raft::KeyValuePair); + } else { + // unfused path needs a full GEMM output (distance matrix row). + mem_per_row += sizeof(MathT) * n_clusters; + } } break; // Other metrics require storing a distance matrix. default: { @@ -378,8 +384,8 @@ void predict(const raft::resources& handle, raft::common::nvtx::range fun_scope( "predict(%zu, %u)", static_cast(n_rows), n_clusters); auto mem_res = mr.value_or(raft::resource::get_workspace_resource_ref(handle)); - auto [max_minibatch_size, _mem_per_row] = - calc_minibatch_size(n_clusters, n_rows, dim, params.metric, std::is_same_v); + auto [max_minibatch_size, _mem_per_row] = calc_minibatch_size( + handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); rmm::device_uvector cur_dataset( std::is_same_v ? 0 : max_minibatch_size * dim, stream, mem_res); bool need_compute_norm = @@ -990,8 +996,8 @@ void build_hierarchical(const raft::resources& handle, // TODO: Remove the explicit managed memory- we shouldn't be creating this on the user's behalf. rmm::mr::managed_memory_resource managed_memory; rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(handle); - auto [max_minibatch_size, mem_per_row] = - calc_minibatch_size(n_clusters, n_rows, dim, params.metric, std::is_same_v); + auto [max_minibatch_size, mem_per_row] = calc_minibatch_size( + handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); // Precompute the L2 norm of the dataset if relevant and not yet computed. rmm::device_uvector dataset_norm_buf(0, stream, device_memory); From 73e19dad8d519c50857ddebdcd080a1ad8a9b5c7 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 15:23:58 +0100 Subject: [PATCH 45/48] Adding error check for CUDA kernel launch failures --- cpp/src/distance/unfused_distance_nn.cuh | 1 + cpp/tests/neighbors/distance_nn.cu | 2 ++ cpp/tests/neighbors/distance_nn_helper.cuh | 2 ++ 3 files changed, 5 insertions(+) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index d7f4d98321..ecd2abbfc3 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -151,6 +151,7 @@ void reduce_min(OutT* out, int blocks = m; reduce_min_kernel <<>>(out, z, x_norm, y_norm, m, n, is_sqrt, initOutBuffer); + RAFT_CUDA_TRY(cudaGetLastError()); } template diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f51d26546a..d8f7625ff7 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -66,7 +66,9 @@ class NNTest : public ::testing::TestWithParam> { raft::random::RngState rng{params_.rng_seed}; if constexpr (std::is_same_v) { fill_int8<<<1000, 256, 0, stream>>>(x.data_handle(), m * k, 0); + RAFT_CUDA_TRY(cudaGetLastError()); fill_int8<<<1000, 256, 0, stream>>>(y.data_handle(), n * k, m * k); + RAFT_CUDA_TRY(cudaGetLastError()); } else { raft::random::uniform(handle, rng, x.data_handle(), m * k, DataT(-1.0), DataT(1.0)); raft::random::uniform(handle, rng, y.data_handle(), n * k, DataT(-1.0), DataT(1.0)); diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index edebbe66d8..fda7b76573 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -119,6 +119,8 @@ void ref_nn(OutT* out, { ref_nn_kernel <<<(m + 127) / 128, 128, 0, stream>>>(out, A, B, m, n, k, sqrt, metric); + + RAFT_CUDA_TRY(cudaGetLastError()); return; } From c4cfc69730b46338733987a9fb94a5330883691b Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 15:40:43 +0100 Subject: [PATCH 46/48] Making sure that all dimensions are non-zero before starting the distance computation --- cpp/src/distance/unfused_distance_nn.cuh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index ecd2abbfc3..6eb296f7f6 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -285,6 +285,8 @@ void unfusedDistanceNNMinReduce(raft::resources const& handle, cudaStream_t stream) { ASSERT(isRowMajor, "unfusedDistanceNN only supports row major inputs"); + + ASSERT(m > 0 && n > 0 && k > 0, "unfusedDistanceNN requires non-zero m, n, and k"); pairwise_distance_gemm( handle, (AccT*)workspace, x, y, m, n, k, xn, yn, stream); From b67dda62a707118b729e07ffbb5dadd782d1df20 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 15:46:47 +0100 Subject: [PATCH 47/48] Minor and formatting changes --- cpp/src/distance/unfused_distance_nn.cuh | 2 +- cpp/tests/neighbors/distance_nn.cu | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 6eb296f7f6..14c0c92b14 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -151,7 +151,7 @@ void reduce_min(OutT* out, int blocks = m; reduce_min_kernel <<>>(out, z, x_norm, y_norm, m, n, is_sqrt, initOutBuffer); - RAFT_CUDA_TRY(cudaGetLastError()); + RAFT_CUDA_TRY(cudaGetLastError()); } template diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index d8f7625ff7..f31f3ebacf 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -189,7 +189,7 @@ const std::vector> input_fp32 = { {4096, 16384, 128, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {8192, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, - // Fused implementation for cosine distance ignores the sqrt paramter, therefore + // Fused implementation for cosine distance ignores the sqrt parameter, therefore // commenting the following two tests // {4096, 4096, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, // {4096, 8192, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, From e8d6bd12070b48a662e066c1e62e59696e593abf Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 6 May 2026 15:54:18 +0100 Subject: [PATCH 48/48] Adding guards for zero-norm vectors --- cpp/src/distance/unfused_distance_nn.cuh | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index 14c0c92b14..f85de31937 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -78,8 +78,13 @@ __global__ void reduce_min_kernel(OutT* out, if constexpr (metric == DistanceType::L2SqrtExpanded || metric == DistanceType::L2Expanded) { dist = x_norm_row + y_norm[col] - AccT(2) * z[row * n + col]; + // GEMM round-off can produce slightly negative expanded distances; clamp to zero. + dist = (dist > AccT(0)) ? dist : AccT(0); } else if constexpr (metric == DistanceType::CosineExpanded) { - dist = AccT(1.0) - (z[row * n + col] / (x_norm_row * y_norm[col])); + // Guard against zero-norm vectors to avoid inf/NaN from division by zero. + AccT denom = x_norm_row * y_norm[col]; + denom = (denom > AccT(0)) ? denom : AccT(1); + dist = AccT(1.0) - (z[row * n + col] / denom); } if (dist < thread_min.value) { thread_min.value = dist; @@ -92,7 +97,9 @@ __global__ void reduce_min_kernel(OutT* out, auto block_result = BlockReduceT(temp_storage).Reduce(thread_min, Reducer{}); if (threadIdx.x == 0) { - if (is_sqrt) { block_result.value = raft::sqrt(block_result.value); } + if (is_sqrt) { + block_result.value = raft::sqrt(block_result.value > AccT(0) ? block_result.value : AccT(0)); + } if constexpr (std::is_same_v) { if (initOutBuffer == true) { out[row] = block_result;