diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index b40c6fbcf0..26ce110158 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -171,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: { @@ -377,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 = @@ -989,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); diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 348a0ec0e0..7a9f97aa74 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -56,6 +57,31 @@ namespace cuvs::cluster::kmeans::detail { +/** + * @brief Returns true if the fused distance NN implementation should be used. + * + * On Ampere (SM <= 8.x) always use fused. + * On Hopper (SM 9.x) use fused when m or n >= 4096. + * On Blackwell (SM >= 10.x) use unfused. + */ +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 >= 10) { + // On Blackwell onwards, use unfused + return false; + } + return false; +} + template struct SamplingOp { DataT* rnd; diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index e271a861e8..b15119599e 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -4,6 +4,7 @@ */ #include "../../distance/fused_distance_nn.cuh" +#include "../../distance/unfused_distance_nn.cuh" #include "kmeans_common.cuh" #include @@ -50,24 +51,50 @@ void minClusterAndDistanceCompute( raft::KeyValuePair initial_value(0, std::numeric_limits::max()); raft::matrix::fill(handle, minClusterAndDistance, initial_value); - workspace.resize((sizeof(int)) * n_samples, stream); - - cuvs::distance::fusedDistanceNNMinReduce, IndexT>( - minClusterAndDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); + bool should_use_fused = + use_fused(handle, n_samples, n_clusters, n_features); + + if (should_use_fused) { + workspace.resize((sizeof(int)) * n_samples, stream); + + cuvs::distance::fusedDistanceNNMinReduce, IndexT>( + minClusterAndDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + L2NormX.data_handle(), + centroidsNorm.data_handle(), + n_samples, + n_clusters, + n_features, + (void*)workspace.data(), + metric != cuvs::distance::DistanceType::L2Expanded, + false, + true, + metric, + 0.0f, + stream); + } else { + workspace.resize(sizeof(DataT) * n_samples * n_clusters, stream); + + cuvs::distance:: + unfusedDistanceNNMinReduce, IndexT>( + handle, + minClusterAndDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + L2NormX.data_handle(), + centroidsNorm.data_handle(), + n_samples, + n_clusters, + n_features, + (void*)workspace.data(), + metric != cuvs::distance::DistanceType::L2Expanded, + false, + true, + metric, + 0.0f, + stream); + } } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh new file mode 100644 index 0000000000..f85de31937 --- /dev/null +++ b/cpp/src/distance/unfused_distance_nn.cuh @@ -0,0 +1,321 @@ +/* + * 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 +#include + +namespace cuvs { +namespace distance { + +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 +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]; + // 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) { + // 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; + 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 > AccT(0) ? block_result.value : AccT(0)); + } + 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]; + } + } + } +} + +/** + * @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 that is 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, + 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); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +template +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) +{ + cudaDataType_t xyType, zType; + cublasComputeType_t computeType; + + 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; + // 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; + } + } 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(cublasSetStream(cublas_h, stream)); + + RAFT_CUBLAS_TRY(cublasGemmEx(cublas_h, + CUBLAS_OP_T, + CUBLAS_OP_N, + N, // Dimensions (swapped due to row/col-major difference) + M, + K, + reinterpret_cast(&alpha), + y, + xyType, + K, + x, + xyType, + K, + reinterpret_cast(&beta), + z, + zType, + N, + computeType, // Computation type + CUBLAS_GEMM_DEFAULT // Algorithm selection + )); +} + +/** + * \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. + * + * @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 + * + * @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(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. + * @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"); + + 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); + + ASSERT((metric == DistanceType::CosineExpanded) || (metric == DistanceType::L2Expanded) || + (metric == DistanceType::L2SqrtExpanded), + "Only cosine and L2 distance types are supported"); + + 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); + } +} + +/** @} */ + +} // namespace distance +} // namespace cuvs + +#endif diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index b30f108789..5987498b89 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -105,7 +105,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/distance_nn.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu new file mode 100644 index 0000000000..f31f3ebacf --- /dev/null +++ b/cpp/tests/neighbors/distance_nn.cu @@ -0,0 +1,268 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../test_utils.cuh" +#include "distance_nn_helper.cuh" + +#include "../../src/distance/fused_distance_nn.cuh" +#include "../../src/distance/unfused_distance_nn.cuh" + +#include +#include +#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; +}; + +__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) { + int hash = (i + seed_offset) * 1103515245 + 12345; + buff[i] = static_cast((hash >> 16) & 0xFF); + } +} + +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}; + 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)); + } + + // 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); + + // 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) { + workspace_size = m * n * sizeof(AccT); + } + + // 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); + } + + 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}, + {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 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}, +}; + +// 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, 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}, +}; + +// 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, 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}, +}; + +// DataT = int8_t, AccT = int32_t +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)); + +// DataT = int8_t, AccT = float +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 diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh new file mode 100644 index 0000000000..fda7b76573 --- /dev/null +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -0,0 +1,210 @@ +/* + * 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])); + } + + v1_norm = raft::sqrt(v1_norm); + v2_norm = raft::sqrt(v2_norm); + 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); + + RAFT_CUDA_TRY(cudaGetLastError()); + 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)); + raft::resource::sync_stream(handle, 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