diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 957ffd88d9..72a4d65a5a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -667,6 +667,7 @@ if(NOT BUILD_CPU_ONLY) src/distance/detail/kernels/kernel_matrices.cu ${pairwise_matrix_dispatch_inst_files} src/distance/distance.cu + src/distance/kde.cu src/distance/pairwise_distance.cu src/distance/sparse_distance.cu src/neighbors/all_neighbors/all_neighbors.cu diff --git a/cpp/include/cuvs/distance/distance.hpp b/cpp/include/cuvs/distance/distance.hpp index 13c8c7bd7e..6eaafb60da 100644 --- a/cpp/include/cuvs/distance/distance.hpp +++ b/cpp/include/cuvs/distance/distance.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -80,6 +80,21 @@ inline bool is_min_close(DistanceType metric) return select_min; } +/** + * @brief Density kernel type for Kernel Density Estimation. + * + * These are the smoothing kernels used in KDE — distinct from the dot-product + * kernels (RBF, Polynomial, etc.) in cuvs::distance::kernels used by SVMs. + */ +enum class DensityKernelType : int { + Gaussian = 0, + Tophat = 1, + Epanechnikov = 2, + Exponential = 3, + Linear = 4, + Cosine = 5 +}; + namespace kernels { enum KernelType { LINEAR, POLYNOMIAL, RBF, TANH }; diff --git a/cpp/include/cuvs/distance/kde.hpp b/cpp/include/cuvs/distance/kde.hpp new file mode 100644 index 0000000000..ea3ef8da78 --- /dev/null +++ b/cpp/include/cuvs/distance/kde.hpp @@ -0,0 +1,81 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +namespace cuvs::distance { + +/** + * @brief Compute log-density estimates for query points using kernel density estimation. + * + * Fuses pairwise distance computation, kernel evaluation, logsumexp reduction, + * and normalization into a single CUDA kernel pass. O(N+M) memory usage — + * the full N×M pairwise distance matrix is never materialised. + * + * Supports 13 distance metrics (all expressible as per-feature accumulation), + * 6 density kernel functions, float32 and float64, and both uniform and + * weighted training sets. + * + * When the query count is small relative to the number of GPU SMs, the + * training set is automatically split across a 2D grid (multi-pass mode) to + * keep the GPU fully utilised. Partial logsumexp results are merged by a + * reduction kernel. + * + * @tparam T float or double + * + * @param[in] handle RAFT resources handle for stream management + * @param[in] query Query points, row-major (n_query × n_features) + * @param[in] train Training points, row-major (n_train × n_features) + * @param[in] weights Per-training-point weights (n_train,), or nullopt for uniform + * @param[out] output Log-density estimates (n_query,) + * @param[in] bandwidth Kernel bandwidth (must be > 0) + * @param[in] sum_weights Sum of sample weights (or n_train if uniform) + * @param[in] kernel Density kernel function + * @param[in] metric Distance metric + * @param[in] metric_arg Metric parameter (e.g. p for Minkowski; ignored otherwise) + */ +template +void kde(raft::resources const& handle, + raft::device_matrix_view query, + raft::device_matrix_view train, + std::optional> weights, + raft::device_vector_view output, + T bandwidth, + T sum_weights, + DensityKernelType kernel, + cuvs::distance::DistanceType metric, + T metric_arg); + +extern template void kde( + raft::resources const&, + raft::device_matrix_view, + raft::device_matrix_view, + std::optional>, + raft::device_vector_view, + float, + float, + DensityKernelType, + cuvs::distance::DistanceType, + float); + +extern template void kde( + raft::resources const&, + raft::device_matrix_view, + raft::device_matrix_view, + std::optional>, + raft::device_vector_view, + double, + double, + DensityKernelType, + cuvs::distance::DistanceType, + double); + +} // namespace cuvs::distance diff --git a/cpp/src/distance/kde.cu b/cpp/src/distance/kde.cu new file mode 100644 index 0000000000..b9989bffec --- /dev/null +++ b/cpp/src/distance/kde.cu @@ -0,0 +1,684 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include + +namespace cuvs::distance { + +// ============================================================================ +// Distance accumulator ops — decomposed into init / accumulate / finalize +// so the tiled kernel can tile over features while accumulating partial +// distances in registers. +// +// Each specialisation defines: +// N_ACC — number of accumulator values per distance computation +// init(acc) — zero the accumulators +// accumulate(acc, a, b, p) — per-feature accumulation +// finalize(acc, d, p) — convert accumulators to final scalar distance +// ============================================================================ + +template +struct DistOp; + +// euclidean: sqrt(sum((a-b)^2)) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + T d = a - b; + acc[0] += d * d; + } + inline __device__ static T finalize(T* acc, int, T) { return sqrt(acc[0]); } +}; + +// sqeuclidean: sum((a-b)^2) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + T d = a - b; + acc[0] += d * d; + } + inline __device__ static T finalize(T* acc, int, T) { return acc[0]; } +}; + +// manhattan: sum(|a-b|) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) { acc[0] += abs(a - b); } + inline __device__ static T finalize(T* acc, int, T) { return acc[0]; } +}; + +// chebyshev: max(|a-b|) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + acc[0] = max(acc[0], abs(a - b)); + } + inline __device__ static T finalize(T* acc, int, T) { return acc[0]; } +}; + +// minkowski: (sum(|a-b|^p))^(1/p) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T p) { acc[0] += pow(abs(a - b), p); } + inline __device__ static T finalize(T* acc, int, T p) { return pow(acc[0], T(1) / p); } +}; + +// cosine: 1 - dot(a,b)/(||a||*||b||) +// acc[0]=dot, acc[1]=||a||^2, acc[2]=||b||^2 +template +struct DistOp { + static constexpr int N_ACC = 3; + inline __device__ static void init(T* acc) { acc[0] = acc[1] = acc[2] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + acc[0] += a * b; + acc[1] += a * a; + acc[2] += b * b; + } + inline __device__ static T finalize(T* acc, int, T) + { + T denom = sqrt(acc[1]) * sqrt(acc[2]); + return (denom > T(0)) ? (T(1) - acc[0] / denom) : T(0); + } +}; + +// correlation: cosine on mean-centred vectors (single-pass via sum identities) +// acc[0]=sum_a, acc[1]=sum_b, acc[2]=sum_a2, acc[3]=sum_b2, acc[4]=sum_ab +template +struct DistOp { + static constexpr int N_ACC = 5; + inline __device__ static void init(T* acc) { acc[0] = acc[1] = acc[2] = acc[3] = acc[4] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + acc[0] += a; + acc[1] += b; + acc[2] += a * a; + acc[3] += b * b; + acc[4] += a * b; + } + inline __device__ static T finalize(T* acc, int d, T) + { + T ma = acc[0] / T(d); + T mb = acc[1] / T(d); + T dot = acc[4] - T(d) * ma * mb; + T na = acc[2] - T(d) * ma * ma; + T nb = acc[3] - T(d) * mb * mb; + T den = sqrt(na) * sqrt(nb); + return (den > T(0)) ? (T(1) - dot / den) : T(0); + } +}; + +// canberra: sum(|a-b|/(|a|+|b|)) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + const T diff = abs(a - b); + const T add = abs(a) + abs(b); + acc[0] += ((add != T(0)) * diff / (add + (add == T(0)))); + } + inline __device__ static T finalize(T* acc, int, T) { return acc[0]; } +}; + +// hellinger: sqrt(1 - sum(sqrt(a)*sqrt(b))) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) { acc[0] += sqrt(a) * sqrt(b); } + inline __device__ static T finalize(T* acc, int, T) + { + const T val = T(1) - acc[0]; + return sqrt((!signbit(val)) * val); + } +}; + +// jensen-shannon +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + const T m = T(0.5) * (a + b); + const bool mz = (m == T(0)); + const T logM = (!mz) * log(m + mz); + const bool xz = (a == T(0)); + const bool yz = (b == T(0)); + acc[0] += (-a * (logM - log(a + xz))) + (-b * (logM - log(b + yz))); + } + inline __device__ static T finalize(T* acc, int, T) { return sqrt(T(0.5) * acc[0]); } +}; + +// hamming: count(a!=b)/d +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) { acc[0] += (a != b); } + inline __device__ static T finalize(T* acc, int d, T) { return acc[0] / T(d); } +}; + +// KL divergence: sum(a*log(a/b)) +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) + { + if (a > T(0) && b > T(0)) { acc[0] += a * log(a / b); } + } + inline __device__ static T finalize(T* acc, int, T) { return acc[0]; } +}; + +// Russell-Rao: (d - sum(a*b)) / d +template +struct DistOp { + static constexpr int N_ACC = 1; + inline __device__ static void init(T* acc) { acc[0] = T(0); } + inline __device__ static void accumulate(T* acc, T a, T b, T) { acc[0] += a * b; } + inline __device__ static T finalize(T* acc, int d, T) { return (T(d) - acc[0]) / T(d); } +}; + +// ============================================================================ +// Log-kernel traits — one specialisation per DensityKernelType +// ============================================================================ + +template +struct LogKernel; + +template +struct LogKernel { + inline __device__ static T eval(T x, T h) { return -(x * x) / (T(2) * h * h); } +}; + +template +struct LogKernel { + inline __device__ static T eval(T x, T h) + { + return (x < h) ? T(0) : cuda::std::numeric_limits::lowest(); + } +}; + +template +struct LogKernel { + inline __device__ static T eval(T x, T h) + { + T z = max(T(1) - (x * x) / (h * h), T(1e-30)); + return (x < h) ? log(z) : cuda::std::numeric_limits::lowest(); + } +}; + +template +struct LogKernel { + inline __device__ static T eval(T x, T h) { return -x / h; } +}; + +template +struct LogKernel { + inline __device__ static T eval(T x, T h) + { + T z = max(T(1) - x / h, T(1e-30)); + return (x < h) ? log(z) : cuda::std::numeric_limits::lowest(); + } +}; + +template +struct LogKernel { + inline __device__ static T eval(T x, T h) + { + T z = max(cos(T(0.5) * T(M_PI) * x / h), T(1e-30)); + return (x < h) ? log(z) : cuda::std::numeric_limits::lowest(); + } +}; + +// ============================================================================ +// Host-side normalization functions (mirror the Python implementations) +// ============================================================================ + +template +T logVn(int n) +{ + return T(0.5) * n * std::log(T(M_PI)) - std::lgamma(T(0.5) * n + T(1)); +} + +template +T logSn(int n) +{ + return std::log(T(2) * T(M_PI)) + logVn(n - 1); +} + +template +T norm_factor(DensityKernelType kernel, T h, int d) +{ + T factor; + switch (kernel) { + case DensityKernelType::Gaussian: factor = T(0.5) * d * std::log(T(2) * T(M_PI)); break; + case DensityKernelType::Tophat: factor = logVn(d); break; + case DensityKernelType::Epanechnikov: factor = logVn(d) + std::log(T(2) / T(d + 2)); break; + case DensityKernelType::Exponential: factor = logSn(d - 1) + std::lgamma(T(d)); break; + case DensityKernelType::Linear: factor = logVn(d) - std::log(T(d + 1)); break; + case DensityKernelType::Cosine: { + // Compute integral_0^1 cos(pi/2 * t) * t^{d-1} dt using the recurrence: + // I_n = (2/pi) - n*(n-1)*(2/pi)^2 * I_{n-2} + // I_0 = 2/pi, I_1 = 2/pi - (2/pi)^2 + // This is derived from repeated integration by parts; both sin and cos + // boundary terms must be included (the old loop-based formula missed the + // cos terms at t=0 for even d). + const T two_over_pi = T(2) / T(M_PI); + const T two_over_pi_sq = two_over_pi * two_over_pi; + T I_prev = two_over_pi; // I_0 + T I_curr = two_over_pi - two_over_pi_sq; // I_1 + const int n = d - 1; // need I_n + if (n == 0) { + factor = std::log(I_prev) + logSn(d - 1); + } else { + for (int j = 2; j <= n; ++j) { + T I_next = two_over_pi - T(j) * T(j - 1) * two_over_pi_sq * I_prev; + I_prev = I_curr; + I_curr = I_next; + } + factor = std::log(I_curr) + logSn(d - 1); + } + } break; + default: throw std::invalid_argument("Unsupported kernel type"); + } + return factor + d * std::log(h); +} + +// ============================================================================ +// Tiled CUDA kernel — edistance-style CELL_TILE + FEAT_TILE optimisation +// +// One thread per query point. Train vectors are cooperatively loaded into +// shared memory in tiles of [FEAT_TILE][CELL_TILE]. Each thread accumulates +// distances from its query point to CELL_TILE train points simultaneously, +// amortising query feature reads and reducing global memory traffic. +// +// Supports both single-pass (full train set → final output) and multi-pass +// (train subset → partial logsumexp) modes. Multi-pass mode is used when +// the query count is too small to fill the GPU, parallelising over the +// train dimension via a 2D grid. +// ============================================================================ + +template +__global__ void kde_tiled_kernel(const T* __restrict__ query, + const T* __restrict__ train, + const T* __restrict__ weights, + T* __restrict__ out_a, + T* __restrict__ out_b, + int n_query, + int n_train, + int d, + T bandwidth, + T metric_arg, + T log_norm, + int train_chunk, + int feat_tile) +{ + using DOp = DistOp; + + extern __shared__ char smem_raw[]; + T* smem_train = reinterpret_cast(smem_raw); // [feat_tile][CELL_TILE] + + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const bool valid = (i < n_query); + + constexpr int N_ACC = DOp::N_ACC; + + // Determine train range for this block + const int j_begin = blockIdx.y * train_chunk; + const int j_end = min(j_begin + train_chunk, n_train); + + // Initialize to lowest() (not -inf) so that out-of-support points returning + // lowest() don't produce 0*exp(+inf)=NaN via exp(-inf - lowest()) = exp(+inf). + T running_max = cuda::std::numeric_limits::lowest(); + T running_sum = T(0); + + // Tile over train points in groups of CELL_TILE + for (int j_base = j_begin; j_base < j_end; j_base += CELL_TILE) { + const int cells_in_tile = min(CELL_TILE, j_end - j_base); + + // Per-train-point accumulators in registers + T acc[CELL_TILE * N_ACC]; +#pragma unroll + for (int c = 0; c < CELL_TILE; ++c) + DOp::init(&acc[c * N_ACC]); + + // Tile over features + for (int feat_base = 0; feat_base < d; feat_base += feat_tile) { + const int feats_in_tile = min(feat_tile, d - feat_base); + + // Cooperatively load train tile into shared memory: smem[feat][cell] + const int total_elems = feat_tile * CELL_TILE; + for (int idx = threadIdx.x; idx < total_elems; idx += blockDim.x) { + const int cell = idx / feat_tile; + const int feat = idx % feat_tile; + T val = T(0); + if (cell < cells_in_tile && feat < feats_in_tile) { + val = train[static_cast(j_base + cell) * d + feat_base + feat]; + } + smem_train[feat * CELL_TILE + cell] = val; + } + + __syncthreads(); + + if (valid) { + for (int f = 0; f < feats_in_tile; ++f) { + const T val_q = query[static_cast(i) * d + feat_base + f]; +#pragma unroll + for (int c = 0; c < CELL_TILE; ++c) { + const T val_t = smem_train[f * CELL_TILE + c]; + DOp::accumulate(&acc[c * N_ACC], val_q, val_t, metric_arg); + } + } + } + + __syncthreads(); + } + + // Finalize distances and fold into streaming logsumexp + if (valid) { +#pragma unroll + for (int c = 0; c < CELL_TILE; ++c) { + if (c >= cells_in_tile) break; + T dist = DOp::finalize(&acc[c * N_ACC], d, metric_arg); + T log_k = LogKernel::eval(dist, bandwidth); + if (weights) log_k += log(weights[j_base + c]); + + if (log_k > running_max) { + running_sum = running_sum * exp(running_max - log_k) + T(1); + running_max = log_k; + } else { + running_sum += exp(log_k - running_max); + } + } + } + } + + if (valid) { + if (out_b == nullptr) { + // Single-pass: write final log-probability + out_a[i] = log(running_sum) + running_max - log_norm; + } else { + // Multi-pass: write partial (max, sum) for later reduction + const size_t idx = static_cast(i) * gridDim.y + blockIdx.y; + out_a[idx] = running_max; + out_b[idx] = running_sum; + } + } +} + +// ============================================================================ +// Reduction kernel — merges partial logsumexp results from multi-pass +// ============================================================================ + +template +__global__ void kde_reduce_kernel(const T* __restrict__ partial_max, + const T* __restrict__ partial_sum, + T* __restrict__ output, + int n_query, + int n_blocks, + T log_norm) +{ + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_query) return; + + T rmax = cuda::std::numeric_limits::lowest(); + T rsum = T(0); + + for (int b = 0; b < n_blocks; ++b) { + const size_t idx = static_cast(i) * n_blocks + b; + const T pm = partial_max[idx]; + const T ps = partial_sum[idx]; + if (pm > rmax) { + rsum = rsum * exp(rmax - pm) + ps; + rmax = pm; + } else { + rsum += ps * exp(pm - rmax); + } + } + output[i] = log(rsum) + rmax - log_norm; +} + +// ============================================================================ +// Double dispatch: runtime enum → compile-time template +// ============================================================================ + +template +void dispatch_metric(cuvs::distance::DistanceType metric, Fn&& fn) +{ + using DT = cuvs::distance::DistanceType; + switch (metric) { + case DT::L2SqrtUnexpanded: fn(std::integral_constant{}); break; + case DT::L2Expanded: fn(std::integral_constant{}); break; + case DT::L1: fn(std::integral_constant{}); break; + case DT::Linf: fn(std::integral_constant{}); break; + case DT::LpUnexpanded: fn(std::integral_constant{}); break; + case DT::CosineExpanded: fn(std::integral_constant{}); break; + case DT::CorrelationExpanded: fn(std::integral_constant{}); break; + case DT::Canberra: fn(std::integral_constant{}); break; + case DT::HellingerExpanded: fn(std::integral_constant{}); break; + case DT::JensenShannon: fn(std::integral_constant{}); break; + case DT::HammingUnexpanded: fn(std::integral_constant{}); break; + case DT::KLDivergence: fn(std::integral_constant{}); break; + case DT::RusselRaoExpanded: fn(std::integral_constant{}); break; + default: throw std::invalid_argument("Unsupported distance metric for KDE"); + } +} + +template +void dispatch_kernel(DensityKernelType kernel, Fn&& fn) +{ + switch (kernel) { + case DensityKernelType::Gaussian: + fn(std::integral_constant{}); + break; + case DensityKernelType::Tophat: + fn(std::integral_constant{}); + break; + case DensityKernelType::Epanechnikov: + fn(std::integral_constant{}); + break; + case DensityKernelType::Exponential: + fn(std::integral_constant{}); + break; + case DensityKernelType::Linear: + fn(std::integral_constant{}); + break; + case DensityKernelType::Cosine: + fn(std::integral_constant{}); + break; + default: throw std::invalid_argument("Unsupported kernel type for KDE"); + } +} + +// ============================================================================ +// Implementation: launches the tiled kernel (1-pass or 2-pass) +// ============================================================================ + +template +void kde(raft::resources const& handle, + raft::device_matrix_view query, + raft::device_matrix_view train, + std::optional> weights, + raft::device_vector_view output, + T bandwidth, + T sum_weights, + DensityKernelType kernel, + cuvs::distance::DistanceType metric, + T metric_arg) +{ + RAFT_EXPECTS(query.extent(0) <= std::numeric_limits::max() && + train.extent(0) <= std::numeric_limits::max() && + query.extent(1) <= std::numeric_limits::max(), + "n_query, n_train, and n_features must fit in int32"); + int n_query = static_cast(query.extent(0)); + int n_train = static_cast(train.extent(0)); + int d = static_cast(query.extent(1)); + + RAFT_EXPECTS(n_query > 0, "n_query must be > 0"); + RAFT_EXPECTS(n_train > 0, "n_train must be > 0"); + RAFT_EXPECTS(d > 0, "n_features must be > 0"); + RAFT_EXPECTS(bandwidth > T(0), "bandwidth must be > 0"); + + const T* query_ptr = query.data_handle(); + const T* train_ptr = train.data_handle(); + const T* weights_ptr = weights.has_value() ? weights->data_handle() : nullptr; + T* output_ptr = output.data_handle(); + + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + T log_norm = std::log(sum_weights) + norm_factor(kernel, bandwidth, d); + + // Cap feature tile to the actual dimension to avoid wasted shared memory + // and cooperative load cycles for low-dimensional data (e.g. 2D embeddings). + const int feat_tile = min(64, d); + // 512 threads for float32 (more cooperative load throughput, better GPU fill). + // 256 for float64 to avoid exceeding per-block register limits with + // CELL_TILE=64 double-precision accumulators (64×2 regs × 512 threads > 65536). + const int threads = (sizeof(T) == 4) ? 512 : 256; + int n_query_blocks = (n_query + threads - 1) / threads; + + dispatch_metric(metric, [&](auto metric_tag) { + dispatch_kernel(kernel, [&](auto kernel_tag) { + constexpr auto M = decltype(metric_tag)::value; + constexpr auto K = decltype(kernel_tag)::value; + + // Adapt CELL_TILE to keep accumulator register pressure under ~128 regs. + // LpUnexpanded uses pow() which decomposes to exp2+log2 with no HW + // intrinsic — cap its tile at 32 to avoid cudaErrorLaunchOutOfResources + // across all architectures (V100–Blackwell). + constexpr int N_ACC = DistOp::N_ACC; + constexpr int ACC_REGS = sizeof(T) / 4; + constexpr int RAW_TILE = 128 / (N_ACC * ACC_REGS); + constexpr int MAX_TILE = (M == cuvs::distance::DistanceType::LpUnexpanded) ? 32 : 64; + constexpr int CELL_TILE = RAW_TILE >= MAX_TILE ? MAX_TILE + : RAW_TILE >= 32 ? 32 + : RAW_TILE >= 16 ? 16 + : RAW_TILE >= 8 ? 8 + : 4; + + size_t smem_bytes = feat_tile * CELL_TILE * sizeof(T); + + // Determine whether to split the train dimension across blocks. + // When n_query is small the GPU is underutilised; splitting the train + // set across a 2D grid exposes more parallelism. + int dev, sm_count; + RAFT_CUDA_TRY(cudaGetDevice(&dev)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev)); + int target_blocks = sm_count * 4; + int n_train_blocks = max(1, target_blocks / n_query_blocks); + int min_train_chunk = CELL_TILE * 4; + n_train_blocks = min(n_train_blocks, max(1, n_train / min_train_chunk)); + + if (n_train_blocks <= 1) { + // Single-pass: process all train points, write directly to output + dim3 grid(n_query_blocks); + kde_tiled_kernel + <<>>(query_ptr, + train_ptr, + weights_ptr, + output_ptr, + static_cast(nullptr), + n_query, + n_train, + d, + bandwidth, + metric_arg, + log_norm, + n_train, + feat_tile); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } else { + // Multi-pass: split train dimension, write partial (max, sum), then reduce + int train_chunk = (n_train + n_train_blocks - 1) / n_train_blocks; + // Round up to CELL_TILE for clean tiling + train_chunk = ((train_chunk + CELL_TILE - 1) / CELL_TILE) * CELL_TILE; + // Recompute actual number of blocks after rounding + n_train_blocks = (n_train + train_chunk - 1) / train_chunk; + + size_t buf_elems = static_cast(n_query) * n_train_blocks; + rmm::device_uvector partial_max(buf_elems, stream); + rmm::device_uvector partial_sum(buf_elems, stream); + + dim3 grid(n_query_blocks, n_train_blocks); + kde_tiled_kernel + <<>>(query_ptr, + train_ptr, + weights_ptr, + partial_max.data(), + partial_sum.data(), + n_query, + n_train, + d, + bandwidth, + metric_arg, + log_norm, + train_chunk, + feat_tile); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + kde_reduce_kernel<<>>( + partial_max.data(), partial_sum.data(), output_ptr, n_query, n_train_blocks, log_norm); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + }); + }); +} + +// Explicit instantiations +template void kde( + raft::resources const&, + raft::device_matrix_view, + raft::device_matrix_view, + std::optional>, + raft::device_vector_view, + float, + float, + DensityKernelType, + cuvs::distance::DistanceType, + float); + +template void kde( + raft::resources const&, + raft::device_matrix_view, + raft::device_matrix_view, + std::optional>, + raft::device_vector_view, + double, + double, + DensityKernelType, + cuvs::distance::DistanceType, + double); + +} // namespace cuvs::distance diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index f5757da423..1d600e6768 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -330,6 +330,7 @@ ConfigureTest( distance/dist_lp_unexp.cu distance/dist_russell_rao.cu distance/gram.cu + distance/kde.cu distance/masked_nn.cu distance/sparse_distance.cu sparse/gram.cu diff --git a/cpp/tests/distance/generate_kde_golden.py b/cpp/tests/distance/generate_kde_golden.py new file mode 100644 index 0000000000..90844d9535 --- /dev/null +++ b/cpp/tests/distance/generate_kde_golden.py @@ -0,0 +1,506 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generate golden values for KDE gtest from sklearn / scipy.""" + +import numpy as np +from scipy.spatial.distance import cdist +from math import lgamma +from scipy.special import logsumexp + +try: + from sklearn.neighbors import KernelDensity + + HAS_SKLEARN = True +except ImportError: + HAS_SKLEARN = False + print("WARNING: sklearn not available, using manual computation only") + + +# ============================================================================ +# Manual KDE reference (for metrics sklearn doesn't support) +# ============================================================================ + + +def log_vn(d): + """Log volume of d-dimensional unit ball.""" + return 0.5 * d * np.log(np.pi) - lgamma(0.5 * d + 1) + + +def log_sn(d): + """Log surface area of d-dimensional unit sphere.""" + return np.log(2 * np.pi) + log_vn(d - 1) + + +def norm_factor(kernel, h, d): + if kernel == "gaussian": + factor = 0.5 * d * np.log(2 * np.pi) + elif kernel == "tophat": + factor = log_vn(d) + elif kernel == "epanechnikov": + factor = log_vn(d) + np.log(2.0 / (d + 2)) + elif kernel == "exponential": + factor = log_sn(d - 1) + lgamma(d) + elif kernel == "linear": + factor = log_vn(d) - np.log(d + 1) + elif kernel == "cosine": + two_over_pi = 2.0 / np.pi + two_over_pi_sq = two_over_pi**2 + I_prev = two_over_pi + I_curr = two_over_pi - two_over_pi_sq + n = d - 1 + if n == 0: + factor = np.log(I_prev) + log_sn(d - 1) + else: + for j in range(2, n + 1): + I_next = two_over_pi - j * (j - 1) * two_over_pi_sq * I_prev + I_prev = I_curr + I_curr = I_next + factor = np.log(I_curr) + log_sn(d - 1) + else: + raise ValueError(f"Unknown kernel: {kernel}") + return factor + d * np.log(h) + + +def log_kernel_eval(dist, h, kernel): + """Evaluate log-kernel (element-wise on arrays).""" + LOWEST = np.finfo(np.float64).min + if kernel == "gaussian": + return -(dist**2) / (2 * h**2) + elif kernel == "tophat": + return np.where(dist < h, 0.0, LOWEST) + elif kernel == "epanechnikov": + z = np.maximum(1 - (dist**2) / (h**2), 1e-30) + return np.where(dist < h, np.log(z), LOWEST) + elif kernel == "exponential": + return -dist / h + elif kernel == "linear": + z = np.maximum(1 - dist / h, 1e-30) + return np.where(dist < h, np.log(z), LOWEST) + elif kernel == "cosine": + z = np.maximum(np.cos(0.5 * np.pi * dist / h), 1e-30) + return np.where(dist < h, np.log(z), LOWEST) + else: + raise ValueError(f"Unknown kernel: {kernel}") + + +def manual_kde(query, train, bandwidth, kernel, dists, weights=None): + """Compute log-density using precomputed distances.""" + n_train = train.shape[0] + d = train.shape[1] + + log_k = log_kernel_eval(dists, bandwidth, kernel) + + if weights is not None: + log_k = log_k + np.log(weights)[np.newaxis, :] + sw = np.sum(weights) + else: + sw = float(n_train) + + log_sum = logsumexp(log_k, axis=1) + log_norm = np.log(sw) + norm_factor(kernel, bandwidth, d) + return log_sum - log_norm + + +def compute_dists_scipy(query, train, metric, metric_arg=None): + """Compute pairwise distances using scipy.""" + if metric == "euclidean": + return cdist(query, train, metric="euclidean") + elif metric == "sqeuclidean": + return cdist(query, train, metric="sqeuclidean") + elif metric == "manhattan": + return cdist(query, train, metric="cityblock") + elif metric == "chebyshev": + return cdist(query, train, metric="chebyshev") + elif metric == "minkowski": + return cdist(query, train, metric="minkowski", p=metric_arg) + elif metric == "cosine": + return cdist(query, train, metric="cosine") + elif metric == "correlation": + return cdist(query, train, metric="correlation") + elif metric == "canberra": + return cdist(query, train, metric="canberra") + elif metric == "hellinger": + # sqrt(max(0, 1 - sum(sqrt(a)*sqrt(b)))) + n_q, n_t = query.shape[0], train.shape[0] + d = np.zeros((n_q, n_t)) + for i in range(n_q): + for j in range(n_t): + val = 1.0 - np.sum(np.sqrt(query[i]) * np.sqrt(train[j])) + d[i, j] = np.sqrt(max(0.0, val)) + return d + elif metric == "jensenshannon": + n_q, n_t = query.shape[0], train.shape[0] + d = np.zeros((n_q, n_t)) + for i in range(n_q): + for j in range(n_t): + a, b = query[i], train[j] + m = 0.5 * (a + b) + acc = 0.0 + for f in range(len(a)): + logM = np.log(m[f]) if m[f] > 0 else 0.0 + logA = np.log(a[f]) if a[f] > 0 else 0.0 + logB = np.log(b[f]) if b[f] > 0 else 0.0 + acc += (-a[f] * (logM - logA)) + (-b[f] * (logM - logB)) + d[i, j] = np.sqrt(0.5 * acc) + return d + elif metric == "hamming": + return cdist(query, train, metric="hamming") + elif metric == "kldivergence": + n_q, n_t = query.shape[0], train.shape[0] + d = np.zeros((n_q, n_t)) + for i in range(n_q): + for j in range(n_t): + a, b = query[i], train[j] + acc = 0.0 + for f in range(len(a)): + if a[f] > 0 and b[f] > 0: + acc += a[f] * np.log(a[f] / b[f]) + d[i, j] = acc + return d + elif metric == "russellrao": + n_q, n_t = query.shape[0], train.shape[0] + dim = query.shape[1] + d = np.zeros((n_q, n_t)) + for i in range(n_q): + for j in range(n_t): + d[i, j] = (dim - np.sum(query[i] * train[j])) / dim + return d + else: + raise ValueError(f"Unknown metric: {metric}") + + +# ============================================================================ +# Data generation +# ============================================================================ + +np.random.seed(42) + +N_QUERY, N_TRAIN, D = 4, 8, 3 + +# General data in [0.1, 2.0] - works for most metrics +query_gen = np.random.uniform(0.1, 2.0, (N_QUERY, D)) +train_gen = np.random.uniform(0.1, 2.0, (N_TRAIN, D)) + +# Round to 4 decimal places for clean hardcoded values +query_gen = np.round(query_gen, 4) +train_gen = np.round(train_gen, 4) + +# Probability data (rows sum to 1) for Hellinger, JS, KL +query_prob = np.abs(np.random.uniform(0.1, 1.0, (N_QUERY, D))) +query_prob = query_prob / query_prob.sum(axis=1, keepdims=True) +train_prob = np.abs(np.random.uniform(0.1, 1.0, (N_TRAIN, D))) +train_prob = train_prob / train_prob.sum(axis=1, keepdims=True) +query_prob = np.round(query_prob, 6) +train_prob = np.round(train_prob, 6) +# Renormalize after rounding +query_prob = query_prob / query_prob.sum(axis=1, keepdims=True) +train_prob = train_prob / train_prob.sum(axis=1, keepdims=True) + +# Weights +weights = np.round(np.random.uniform(0.5, 3.0, N_TRAIN), 4) + +# High-dimensional dataset for feature/train tiling test (deterministic formula) +N_QUERY_HD, N_TRAIN_HD, D_HD = 4, 100, 128 +query_hd = np.zeros((N_QUERY_HD, D_HD)) +train_hd = np.zeros((N_TRAIN_HD, D_HD)) +for i in range(N_QUERY_HD): + for j in range(D_HD): + query_hd[i, j] = 0.1 + ((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9 +for i in range(N_TRAIN_HD): + for j in range(D_HD): + train_hd[i, j] = 0.1 + ((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9 + +# Large dataset for multi-pass test (deterministic formula) +N_TRAIN_LARGE = 2000 +query_large = np.zeros((2, D)) +train_large = np.zeros((N_TRAIN_LARGE, D)) +for i in range(2): + for j in range(D): + query_large[i, j] = ( + 0.1 + ((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9 + ) +for i in range(N_TRAIN_LARGE): + for j in range(D): + train_large[i, j] = ( + 0.1 + ((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9 + ) +weights_large = np.zeros(N_TRAIN_LARGE) +for i in range(N_TRAIN_LARGE): + weights_large[i] = 0.5 + ((i * 31 + 17) % 1000) / 1000.0 * 2.5 + + +# ============================================================================ +# Generate golden values +# ============================================================================ + + +def fmt_array(name, arr, type_str="double"): + """Format array as C++ initializer.""" + vals = ", ".join(f"{v:.15e}" for v in arr) + return f"const {type_str} {name}[] = {{{vals}}};" + + +def fmt_2d_array(name, arr, type_str="double"): + """Format 2D array as flattened C++ initializer (row-major).""" + flat = arr.flatten() + vals = ",\n ".join(f"{v:.15e}" for v in flat) + return f"const {type_str} {name}[] = {{\n {vals}}};" + + +results = {} + + +# --- 1. Each kernel with Euclidean metric --- +BW_KERNEL = 4.0 # Large enough for compact-support kernels +kernel_names = [ + "gaussian", + "tophat", + "epanechnikov", + "exponential", + "linear", + "cosine", +] + +print("// === Kernel tests (Euclidean metric, bandwidth=4.0) ===") +dists_euc = compute_dists_scipy(query_gen, train_gen, "euclidean") +for kname in kernel_names: + if HAS_SKLEARN: + kde = KernelDensity( + bandwidth=BW_KERNEL, kernel=kname, metric="euclidean" + ) + kde.fit(train_gen) + expected = kde.score_samples(query_gen) + else: + expected = manual_kde( + query_gen, train_gen, BW_KERNEL, kname, dists_euc + ) + + # Cross-validate: manual should match sklearn + manual_expected = manual_kde( + query_gen, train_gen, BW_KERNEL, kname, dists_euc + ) + if not np.allclose(expected, manual_expected, atol=1e-10): + print(f" WARNING: sklearn vs manual mismatch for kernel={kname}") + print(f" sklearn: {expected}") + print(f" manual: {manual_expected}") + + results[f"kernel_{kname}"] = expected + print(f"// kernel={kname}: {expected}") + + +# --- 2. Each metric with Gaussian kernel --- +BW_METRIC = 1.0 + +# Metrics that sklearn supports via BallTree +sklearn_metric_map = { + "L2SqrtUnexpanded": ("euclidean", "euclidean", None), + "L1": ("manhattan", "manhattan", None), + "Linf": ("chebyshev", "chebyshev", None), + "LpUnexpanded": ("minkowski", "minkowski", 3.0), +} + +# All metrics with scipy distance names +all_metrics = { + "L2SqrtUnexpanded": ("euclidean", None), + "L2Expanded": ("sqeuclidean", None), + "L1": ("manhattan", None), + "Linf": ("chebyshev", None), + "LpUnexpanded": ("minkowski", 3.0), + "CosineExpanded": ("cosine", None), + "CorrelationExpanded": ("correlation", None), + "Canberra": ("canberra", None), +} + +# Metrics needing probability data +prob_metrics = { + "HellingerExpanded": ("hellinger", None), + "JensenShannon": ("jensenshannon", None), + "KLDivergence": ("kldivergence", None), +} + +# Metrics needing special handling +special_metrics = { + "HammingUnexpanded": ("hamming", None), + "RusselRaoExpanded": ("russellrao", None), +} + +print("\n// === Metric tests (Gaussian kernel, bandwidth=1.0) ===") + +# Standard metrics with general data +for metric_name, (scipy_name, metric_arg) in all_metrics.items(): + dists = compute_dists_scipy(query_gen, train_gen, scipy_name, metric_arg) + expected = manual_kde(query_gen, train_gen, BW_METRIC, "gaussian", dists) + + # Cross-validate with sklearn where possible + if HAS_SKLEARN and metric_name in sklearn_metric_map: + sk_metric, _, sk_arg = sklearn_metric_map[metric_name] + kwargs = {} + if sk_arg is not None: + kwargs["metric_params"] = {"p": sk_arg} + kde = KernelDensity( + bandwidth=BW_METRIC, kernel="gaussian", metric=sk_metric, **kwargs + ) + kde.fit(train_gen) + sk_expected = kde.score_samples(query_gen) + if not np.allclose(expected, sk_expected, atol=1e-10): + print( + f" WARNING: sklearn vs manual mismatch for metric={metric_name}" + ) + print(f" sklearn: {sk_expected}") + print(f" manual: {expected}") + expected = sk_expected # Prefer sklearn values + + results[f"metric_{metric_name}"] = expected + print(f"// metric={metric_name}: {expected}") + +# Probability metrics +for metric_name, (scipy_name, metric_arg) in prob_metrics.items(): + dists = compute_dists_scipy(query_prob, train_prob, scipy_name, metric_arg) + expected = manual_kde(query_prob, train_prob, BW_METRIC, "gaussian", dists) + results[f"metric_{metric_name}"] = expected + print(f"// metric={metric_name} (prob data): {expected}") + +# Hamming and RussellRao with general data +for metric_name, (scipy_name, metric_arg) in special_metrics.items(): + dists = compute_dists_scipy(query_gen, train_gen, scipy_name, metric_arg) + expected = manual_kde(query_gen, train_gen, BW_METRIC, "gaussian", dists) + results[f"metric_{metric_name}"] = expected + print(f"// metric={metric_name}: {expected}") + + +# --- 3. Weighted test --- +print("\n// === Weighted tests ===") +if HAS_SKLEARN: + kde = KernelDensity( + bandwidth=BW_METRIC, kernel="gaussian", metric="euclidean" + ) + kde.fit(train_gen, sample_weight=weights) + expected_weighted = kde.score_samples(query_gen) +else: + dists = compute_dists_scipy(query_gen, train_gen, "euclidean") + expected_weighted = manual_kde( + query_gen, train_gen, BW_METRIC, "gaussian", dists, weights + ) + +# Cross-validate +dists_euc_m = compute_dists_scipy(query_gen, train_gen, "euclidean") +manual_weighted = manual_kde( + query_gen, train_gen, BW_METRIC, "gaussian", dists_euc_m, weights +) +if HAS_SKLEARN and not np.allclose( + expected_weighted, manual_weighted, atol=1e-6 +): + print(" WARNING: sklearn vs manual mismatch for weighted") + print(f" sklearn: {expected_weighted}") + print(f" manual: {manual_weighted}") + +results["weighted_gaussian_euclidean"] = expected_weighted +print(f"// weighted Gaussian+Euclidean: {expected_weighted}") + + +# --- 4. High-dimensional tiling test --- +print( + "\n// === High-dimensional tiling test (n_query=4, n_train=100, d=128) ===" +) +dists_hd = compute_dists_scipy(query_hd, train_hd, "euclidean") +expected_hd = manual_kde(query_hd, train_hd, BW_METRIC, "gaussian", dists_hd) +results["highd_gaussian"] = expected_hd +print(f"// high-d Gaussian+Euclidean: {expected_hd}") + + +# --- 5. Multi-pass test --- +print("\n// === Multi-pass tests (n_query=2, n_train=2000) ===") +BW_LARGE = 1.0 +if HAS_SKLEARN: + kde = KernelDensity( + bandwidth=BW_LARGE, kernel="gaussian", metric="euclidean" + ) + kde.fit(train_large) + expected_mp = kde.score_samples(query_large) +else: + dists_large = compute_dists_scipy(query_large, train_large, "euclidean") + expected_mp = manual_kde( + query_large, train_large, BW_LARGE, "gaussian", dists_large + ) +results["multipass_gaussian_euclidean"] = expected_mp +print(f"// multi-pass Gaussian+Euclidean: {expected_mp}") + +# Multi-pass weighted +if HAS_SKLEARN: + kde = KernelDensity( + bandwidth=BW_LARGE, kernel="gaussian", metric="euclidean" + ) + kde.fit(train_large, sample_weight=weights_large) + expected_mp_w = kde.score_samples(query_large) +else: + dists_large = compute_dists_scipy(query_large, train_large, "euclidean") + expected_mp_w = manual_kde( + query_large, + train_large, + BW_LARGE, + "gaussian", + dists_large, + weights_large, + ) +results["multipass_weighted"] = expected_mp_w +print(f"// multi-pass weighted: {expected_mp_w}") + + +# ============================================================================ +# Output C++ code +# ============================================================================ + +print( + "\n\n// ============================================================================" +) +print("// C++ golden value arrays (copy into kde.cu test)") +print( + "// ============================================================================\n" +) + +# Input data +print(fmt_2d_array("golden_query", query_gen)) +print(fmt_2d_array("golden_train", train_gen)) +print(fmt_2d_array("golden_query_prob", query_prob)) +print(fmt_2d_array("golden_train_prob", train_prob)) +print(fmt_array("golden_weights", weights)) + +# Kernel test expected values +for kname in kernel_names: + print(fmt_array(f"expected_kernel_{kname}", results[f"kernel_{kname}"])) + +# Metric test expected values +for metric_name in ( + list(all_metrics.keys()) + + list(prob_metrics.keys()) + + list(special_metrics.keys()) +): + safe_name = metric_name.replace("Expanded", "").replace("Unexpanded", "") + print( + fmt_array( + f"expected_metric_{safe_name}", results[f"metric_{metric_name}"] + ) + ) + +# Weighted +print(fmt_array("expected_weighted", results["weighted_gaussian_euclidean"])) + +# High-d +print(fmt_array("expected_highd_gaussian", results["highd_gaussian"])) + +# Multi-pass large data generation formula (for C++) +print("\n// Multi-pass: generate train_large in C++ with:") +print("// for i in [0, 2000): for j in [0, 3):") +print( + "// data[i*3+j] = 0.1 + ((i*1337 + j*7 + 42) % 1000) / 1000.0 * 1.9;" +) +print("// query_large same formula with i in [0, 2)") +print("// weights_large: 0.5 + ((i*31 + 17) % 1000) / 1000.0 * 2.5") +print(fmt_array("expected_multipass", results["multipass_gaussian_euclidean"])) +print(fmt_array("expected_multipass_weighted", results["multipass_weighted"])) + +# Summary +print(f"\n// Total test cases: {len(results)}") +print(f"// sklearn available: {HAS_SKLEARN}") diff --git a/cpp/tests/distance/kde.cu b/cpp/tests/distance/kde.cu new file mode 100644 index 0000000000..11b7f462ac --- /dev/null +++ b/cpp/tests/distance/kde.cu @@ -0,0 +1,814 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +// Golden values generated by generate_kde_golden.py and validated against +// scikit-learn's KernelDensity (cross-checked with scipy.spatial.distance). + +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::distance { +namespace { + +// ============================================================================ +// Golden input data (n_query=4, n_train=8, d=3) +// Generated with numpy seed=42, uniform [0.1, 2.0], rounded to 4 decimals. +// ============================================================================ + +// clang-format off +const double golden_query[] = { + 8.116e-01, 1.9064e+00, 1.4908e+00, + 1.2375e+00, 3.964e-01, 3.964e-01, + 2.104e-01, 1.7457e+00, 1.2421e+00, + 1.4453e+00, 1.391e-01, 1.9428e+00}; + +const double golden_train[] = { + 1.6816e+00, 5.034e-01, 4.455e-01, + 4.485e-01, 6.781e-01, 1.097e+00, + 9.207e-01, 6.533e-01, 1.2625e+00, + 3.65e-01, 6.551e-01, 7.961e-01, + 9.665e-01, 1.5918e+00, 4.794e-01, + 1.077e+00, 1.2256e+00, 1.883e-01, + 1.2543e+00, 4.24e-01, 2.236e-01, + 1.9029e+00, 1.9347e+00, 1.636e+00}; + +// Probability data for Hellinger, Jensen-Shannon, KL (rows sum to ~1) +const double golden_query_prob[] = { + 2.92794e-01, 1.47046e-01, 5.6016e-01, + 3.963926036073964e-01, 1.676488323511676e-01, 4.359585640414360e-01, + 9.473709473709474e-02, 6.644206644206645e-01, 2.408422408422408e-01, + 4.23298e-01, 2.31349e-01, 3.45353e-01}; + +const double golden_train_prob[] = { + 3.23336e-01, 1.45475e-01, 5.31189e-01, + 3.011573011573012e-01, 3.570113570113570e-01, 3.418313418313418e-01, + 3.65168e-01, 5.32028e-01, 1.02804e-01, + 3.41263e-01, 1.73734e-01, 4.85003e-01, + 2.74293e-01, 2.09901e-01, 5.15806e-01, + 3.09083e-01, 2.58995e-01, 4.31922e-01, + 1.86554e-01, 6.76021e-01, 1.37425e-01, + 4.79228e-01, 3.85546e-01, 1.35226e-01}; + +const double golden_weights[] = { + 5.138e-01, 2.5387e+00, 2.2671e+00, 2.3225e+00, + 2.4282e+00, 6.851e-01, 1.3962e+00, 7.897e-01}; +// clang-format on + +constexpr int N_QUERY = 4; +constexpr int N_TRAIN = 8; +constexpr int D = 3; +constexpr double BW_K = 4.0; // bandwidth for kernel tests +constexpr double BW_M = 1.0; // bandwidth for metric tests +constexpr double TOL_F = 0.01; // float tolerance +constexpr double TOL_D = 1e-8; // double tolerance + +// ============================================================================ +// Expected outputs — validated against sklearn.neighbors.KernelDensity +// ============================================================================ + +// clang-format off + +// --- Kernel tests: each kernel with Euclidean metric, bandwidth=4.0 ---------- +const double expected_kernel_gaussian[] = {-6.985831191382569e+00, -6.953814282065952e+00, -6.988906802303401e+00, -7.007510953007653e+00}; +const double expected_kernel_tophat[] = {-5.591295041660853e+00, -5.591295041660853e+00, -5.591295041660853e+00, -5.591295041660853e+00}; +const double expected_kernel_epanechnikov[] = {-4.827319358476036e+00, -4.755933236670615e+00, -4.834850378257059e+00, -4.879352298149549e+00}; +const double expected_kernel_exponential[] = {-7.746494343178423e+00, -7.620259536022148e+00, -7.751990400473733e+00, -7.800404376143302e+00}; +const double expected_kernel_linear[] = {-4.661923679475233e+00, -4.486935502588164e+00, -4.672012358013138e+00, -4.751694562578681e+00}; +const double expected_kernel_cosine[] = {-4.758868545881162e+00, -4.672056888256515e+00, -4.767563308290395e+00, -4.820902624764557e+00}; + +// --- Metric tests: each metric with Gaussian kernel, bandwidth=1.0 ----------- +const double expected_metric_L2Sqrt[] = {-3.773347385399441e+00, -3.237487197978773e+00, -3.792416676154021e+00, -4.081323070175513e+00}; +const double expected_metric_L2[] = {-4.351045181188024e+00, -3.265414863290537e+00, -4.308482695431518e+00, -5.026718429596215e+00}; +const double expected_metric_L1[] = {-4.534517691762289e+00, -3.590244898954789e+00, -4.690939219312805e+00, -5.420760479120334e+00}; +const double expected_metric_Linf[] = {-3.536183484952867e+00, -3.131636582078475e+00, -3.453812113010937e+00, -3.663484100823001e+00}; +const double expected_metric_Lp[] = {-3.634306192386314e+00, -3.177198157081673e+00, -3.617941711158735e+00, -3.865636492533764e+00}; +const double expected_metric_Cosine[] = {-2.775990213060389e+00, -2.776616645880969e+00, -2.803150466324275e+00, -2.797431562726978e+00}; +const double expected_metric_Correlation[] = {-3.299113322056466e+00, -3.196609763091426e+00, -3.310166132472229e+00, -3.350015447606014e+00}; +const double expected_metric_Canberra[] = {-3.264532146862635e+00, -3.174616384273250e+00, -3.545718987712004e+00, -3.734349603505341e+00}; +const double expected_metric_Hellinger[] = {-2.787131063986393e+00, -2.779555079618617e+00, -2.802456618620946e+00, -2.773723039913471e+00}; +const double expected_metric_JensenShannon[] = {-2.786096965201834e+00, -2.778968568160119e+00, -2.801208991538172e+00, -2.773469792696907e+00}; +const double expected_metric_KLDivergence[] = {-2.829684660629893e+00, -2.795168262992252e+00, -2.844215468422274e+00, -2.774843583859613e+00}; +const double expected_metric_Hamming[] = {-3.256815599614018e+00, -3.256815599614018e+00, -3.256815599614018e+00, -3.256815599614018e+00}; +const double expected_metric_RusselRao[] = {-2.879997942754293e+00, -2.842532898863741e+00, -2.842100803659594e+00, -2.835252865846702e+00}; + +// --- Weighted: Gaussian + Euclidean, bandwidth=1.0 --------------------------- +const double expected_weighted[] = {-3.691891610762810e+00, -3.264752149776808e+00, -3.610241442309405e+00, -3.977545737565916e+00}; + +// --- High-d tiling: Gaussian + Euclidean, bw=1.0 ---------------------------- +// d=128, nt=100: 2 full feat tiles, ~3 train tiles +const double expected_highd_128d[] = {-1.197941609669934e+02, -1.213662986398092e+02, -1.218918586452003e+02, -1.197682527483547e+02}; +// d=100, nt=100: partial last feat tile (64+36) +const double expected_highd_100d[] = {-9.394095288665109e+01, -9.563481208833603e+01, -9.616064201998816e+01, -9.388707692021224e+01}; +// d=200, nt=150: 3+ feat tiles, partial train tile +const double expected_highd_200d[] = {-1.879364765631665e+02, -1.879383487102569e+02, -1.887424577099825e+02, -1.883374949960477e+02}; + +// --- Multi-pass: n_query=2, n_train=2000, Gaussian + Euclidean, bw=1.0 ------ +const double expected_multipass[] = {-3.614608321851654e+00, -3.161743189037974e+00}; +const double expected_multipass_weighted[] = {-3.615504865329237e+00, -3.161944705945213e+00}; + +// clang-format on + +// ============================================================================ +// Helper: upload data, run GPU KDE, compare against golden expected values +// ============================================================================ + +template +void run_kde_golden(const double* query_data, + const double* train_data, + int n_query, + int n_train, + int d, + const double* weights_data, + T bandwidth, + DensityKernelType kernel, + DistanceType metric, + T metric_arg, + const double* expected, + T tolerance) +{ + // Convert input data to test precision + std::vector h_query(n_query * d); + std::vector h_train(n_train * d); + for (int i = 0; i < n_query * d; ++i) + h_query[i] = static_cast(query_data[i]); + for (int i = 0; i < n_train * d; ++i) + h_train[i] = static_cast(train_data[i]); + + std::vector h_weights; + T sum_weights = static_cast(n_train); + if (weights_data) { + h_weights.resize(n_train); + sum_weights = T(0); + for (int i = 0; i < n_train; ++i) { + h_weights[i] = static_cast(weights_data[i]); + sum_weights += h_weights[i]; + } + } + + std::vector h_expected(n_query); + for (int i = 0; i < n_query; ++i) + h_expected[i] = static_cast(expected[i]); + + // Device allocations + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + rmm::device_uvector d_query(n_query * d, stream); + rmm::device_uvector d_train(n_train * d, stream); + rmm::device_uvector d_output(n_query, stream); + rmm::device_uvector d_weights(weights_data ? n_train : 0, stream); + + raft::update_device(d_query.data(), h_query.data(), n_query * d, stream); + raft::update_device(d_train.data(), h_train.data(), n_train * d, stream); + if (weights_data) { raft::update_device(d_weights.data(), h_weights.data(), n_train, stream); } + + auto query_view = raft::make_device_matrix_view( + d_query.data(), static_cast(n_query), static_cast(d)); + auto train_view = raft::make_device_matrix_view( + d_train.data(), static_cast(n_train), static_cast(d)); + auto output_view = raft::make_device_vector_view( + d_output.data(), static_cast(n_query)); + + std::optional> weights_opt; + if (weights_data) { + weights_opt = raft::make_device_vector_view( + d_weights.data(), static_cast(n_train)); + } + + cuvs::distance::kde(handle, + query_view, + train_view, + weights_opt, + output_view, + bandwidth, + sum_weights, + kernel, + metric, + metric_arg); + + ASSERT_TRUE(cuvs::devArrMatchHost( + h_expected.data(), d_output.data(), n_query, cuvs::CompareApprox(tolerance), stream)); +} + +// ============================================================================ +// Multi-pass helper: deterministic data generation matching the Python script +// ============================================================================ + +template +void run_kde_multipass(const double* expected, bool weighted, T tolerance) +{ + constexpr int nq = 2, nt = 2000, dim = 3; + std::vector h_query(nq * dim); + std::vector h_train(nt * dim); + for (int i = 0; i < nq; ++i) + for (int j = 0; j < dim; ++j) + h_query[i * dim + j] = T(0.1) + T(((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9); + for (int i = 0; i < nt; ++i) + for (int j = 0; j < dim; ++j) + h_train[i * dim + j] = T(0.1) + T(((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9); + + std::vector h_weights(nt); + T sum_weights = static_cast(nt); + if (weighted) { + sum_weights = T(0); + for (int i = 0; i < nt; ++i) { + h_weights[i] = T(0.5) + T(((i * 31 + 17) % 1000) / 1000.0 * 2.5); + sum_weights += h_weights[i]; + } + } + + std::vector h_expected(nq); + for (int i = 0; i < nq; ++i) + h_expected[i] = static_cast(expected[i]); + + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + rmm::device_uvector d_query(nq * dim, stream); + rmm::device_uvector d_train(nt * dim, stream); + rmm::device_uvector d_output(nq, stream); + rmm::device_uvector d_weights(weighted ? nt : 0, stream); + + raft::update_device(d_query.data(), h_query.data(), nq * dim, stream); + raft::update_device(d_train.data(), h_train.data(), nt * dim, stream); + if (weighted) { raft::update_device(d_weights.data(), h_weights.data(), nt, stream); } + + auto query_view = raft::make_device_matrix_view( + d_query.data(), std::int64_t(nq), std::int64_t(dim)); + auto train_view = raft::make_device_matrix_view( + d_train.data(), std::int64_t(nt), std::int64_t(dim)); + auto output_view = + raft::make_device_vector_view(d_output.data(), std::int64_t(nq)); + + std::optional> weights_opt; + if (weighted) { + weights_opt = + raft::make_device_vector_view(d_weights.data(), std::int64_t(nt)); + } + + cuvs::distance::kde(handle, + query_view, + train_view, + weights_opt, + output_view, + T(1.0), + sum_weights, + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + T(2.0)); + + ASSERT_TRUE(cuvs::devArrMatchHost( + h_expected.data(), d_output.data(), nq, cuvs::CompareApprox(tolerance), stream)); +} + +// ============================================================================ +// High-dimensional helper: exercises feature tiling (d > feat_tile=64) +// and train tiling (n_train > CELL_TILE). +// ============================================================================ + +template +void run_kde_highd(int nt, int dim, const double* expected, T tolerance) +{ + constexpr int nq = 4; + std::vector h_query(nq * dim); + std::vector h_train(nt * dim); + for (int i = 0; i < nq; ++i) + for (int j = 0; j < dim; ++j) + h_query[i * dim + j] = T(0.1) + T(((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9); + for (int i = 0; i < nt; ++i) + for (int j = 0; j < dim; ++j) + h_train[i * dim + j] = T(0.1) + T(((i * 1337 + j * 7 + 42) % 1000) / 1000.0 * 1.9); + + std::vector h_expected(nq); + for (int i = 0; i < nq; ++i) + h_expected[i] = static_cast(expected[i]); + + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + rmm::device_uvector d_query(nq * dim, stream); + rmm::device_uvector d_train(nt * dim, stream); + rmm::device_uvector d_output(nq, stream); + + raft::update_device(d_query.data(), h_query.data(), nq * dim, stream); + raft::update_device(d_train.data(), h_train.data(), nt * dim, stream); + + auto query_view = raft::make_device_matrix_view( + d_query.data(), std::int64_t(nq), std::int64_t(dim)); + auto train_view = raft::make_device_matrix_view( + d_train.data(), std::int64_t(nt), std::int64_t(dim)); + auto output_view = + raft::make_device_vector_view(d_output.data(), std::int64_t(nq)); + + std::optional> weights_opt; + + cuvs::distance::kde(handle, + query_view, + train_view, + weights_opt, + output_view, + T(1.0), + T(nt), + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + T(2.0)); + + ASSERT_TRUE(cuvs::devArrMatchHost( + h_expected.data(), d_output.data(), nq, cuvs::CompareApprox(tolerance), stream)); +} + +} // namespace + +// ============================================================================ +// Each kernel with Euclidean metric (isolates kernel evaluation + normalization) +// ============================================================================ + +TEST(KdeKernelF, Gaussian) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_K), + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_kernel_gaussian, + float(TOL_F)); +} + +TEST(KdeKernelF, Tophat) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_K), + DensityKernelType::Tophat, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_kernel_tophat, + float(TOL_F)); +} + +TEST(KdeKernelF, Epanechnikov) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_K), + DensityKernelType::Epanechnikov, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_kernel_epanechnikov, + float(TOL_F)); +} + +TEST(KdeKernelF, Exponential) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_K), + DensityKernelType::Exponential, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_kernel_exponential, + float(TOL_F)); +} + +TEST(KdeKernelF, Linear) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_K), + DensityKernelType::Linear, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_kernel_linear, + float(TOL_F)); +} + +TEST(KdeKernelF, Cosine) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_K), + DensityKernelType::Cosine, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_kernel_cosine, + float(TOL_F)); +} + +// ============================================================================ +// Each metric with Gaussian kernel (isolates distance correctness) +// ============================================================================ + +TEST(KdeMetricF, L2Sqrt) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_metric_L2Sqrt, + float(TOL_F)); +} + +TEST(KdeMetricF, L2Expanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::L2Expanded, + 2.0f, + expected_metric_L2, + float(TOL_F)); +} + +TEST(KdeMetricF, L1) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::L1, + 2.0f, + expected_metric_L1, + float(TOL_F)); +} + +TEST(KdeMetricF, Linf) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::Linf, + 2.0f, + expected_metric_Linf, + float(TOL_F)); +} + +TEST(KdeMetricF, LpUnexpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::LpUnexpanded, + 3.0f, + expected_metric_Lp, + float(TOL_F)); +} + +TEST(KdeMetricF, CosineExpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::CosineExpanded, + 2.0f, + expected_metric_Cosine, + float(TOL_F)); +} + +TEST(KdeMetricF, CorrelationExpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::CorrelationExpanded, + 2.0f, + expected_metric_Correlation, + float(TOL_F)); +} + +TEST(KdeMetricF, Canberra) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::Canberra, + 2.0f, + expected_metric_Canberra, + float(TOL_F)); +} + +TEST(KdeMetricF, HellingerExpanded) +{ + run_kde_golden(golden_query_prob, + golden_train_prob, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::HellingerExpanded, + 2.0f, + expected_metric_Hellinger, + float(TOL_F)); +} + +TEST(KdeMetricF, JensenShannon) +{ + run_kde_golden(golden_query_prob, + golden_train_prob, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::JensenShannon, + 2.0f, + expected_metric_JensenShannon, + float(TOL_F)); +} + +TEST(KdeMetricF, KLDivergence) +{ + run_kde_golden(golden_query_prob, + golden_train_prob, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::KLDivergence, + 2.0f, + expected_metric_KLDivergence, + float(TOL_F)); +} + +TEST(KdeMetricF, HammingUnexpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::HammingUnexpanded, + 2.0f, + expected_metric_Hamming, + float(TOL_F)); +} + +TEST(KdeMetricF, RusselRaoExpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::RusselRaoExpanded, + 2.0f, + expected_metric_RusselRao, + float(TOL_F)); +} + +// ============================================================================ +// Weighted inputs +// ============================================================================ + +TEST(KdeWeightedF, GaussianEuclidean) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + golden_weights, + float(BW_M), + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + 2.0f, + expected_weighted, + float(TOL_F)); +} + +// ============================================================================ +// High-dimensional tiling (exercises feature and train tiling loops) +// ============================================================================ + +// d=128, nt=100: 2 full feature tiles, ~3 train tiles +TEST(KdeHighDimF, D128) { run_kde_highd(100, 128, expected_highd_128d, float(TOL_F)); } + +TEST(KdeHighDimD, D128) { run_kde_highd(100, 128, expected_highd_128d, TOL_D); } + +// d=100, nt=100: partial last feature tile (64+36) +TEST(KdeHighDimF, D100) { run_kde_highd(100, 100, expected_highd_100d, float(TOL_F)); } + +TEST(KdeHighDimD, D100) { run_kde_highd(100, 100, expected_highd_100d, TOL_D); } + +// d=200, nt=150: 3+ feature tiles, partial train tile +TEST(KdeHighDimF, D200) { run_kde_highd(150, 200, expected_highd_200d, float(TOL_F)); } + +TEST(KdeHighDimD, D200) { run_kde_highd(150, 200, expected_highd_200d, TOL_D); } + +// ============================================================================ +// Multi-pass (n_query=2, n_train=2000 forces 2D grid on most GPUs) +// ============================================================================ + +TEST(KdeMultiPassF, GaussianEuclidean) +{ + run_kde_multipass(expected_multipass, false, float(TOL_F)); +} + +TEST(KdeMultiPassF, GaussianEuclideanWeighted) +{ + run_kde_multipass(expected_multipass_weighted, true, float(TOL_F)); +} + +// ============================================================================ +// Double-precision tests (subset for tighter tolerance validation) +// ============================================================================ + +TEST(KdeKernelD, Gaussian) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + BW_K, + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + 2.0, + expected_kernel_gaussian, + TOL_D); +} + +TEST(KdeKernelD, Epanechnikov) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + BW_K, + DensityKernelType::Epanechnikov, + DistanceType::L2SqrtUnexpanded, + 2.0, + expected_kernel_epanechnikov, + TOL_D); +} + +TEST(KdeKernelD, Cosine) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + BW_K, + DensityKernelType::Cosine, + DistanceType::L2SqrtUnexpanded, + 2.0, + expected_kernel_cosine, + TOL_D); +} + +TEST(KdeMetricD, L2Sqrt) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + BW_M, + DensityKernelType::Gaussian, + DistanceType::L2SqrtUnexpanded, + 2.0, + expected_metric_L2Sqrt, + TOL_D); +} + +TEST(KdeMetricD, CosineExpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + BW_M, + DensityKernelType::Gaussian, + DistanceType::CosineExpanded, + 2.0, + expected_metric_Cosine, + TOL_D); +} + +TEST(KdeMetricD, CorrelationExpanded) +{ + run_kde_golden(golden_query, + golden_train, + N_QUERY, + N_TRAIN, + D, + nullptr, + BW_M, + DensityKernelType::Gaussian, + DistanceType::CorrelationExpanded, + 2.0, + expected_metric_Correlation, + 1e-7); +} + +TEST(KdeMultiPassD, GaussianEuclidean) +{ + run_kde_multipass(expected_multipass, false, TOL_D); +} + +TEST(KdeMultiPassD, GaussianEuclideanWeighted) +{ + run_kde_multipass(expected_multipass_weighted, true, TOL_D); +} + +} // namespace cuvs::distance